#lang racket
(require redex)
(require plot)

; Task-Parallel Assembly Language (TPAL)
; ===============================
;
; A RISC machine that features native support for task parallelism. We define the grammar for TPAL
; by incremental extensions of the core language, namely TPAL-0.
 
; TPAL-0 syntax
; -------------
; There are two kinds of memory: registers and a heap. Registers always start with the character
; 'r', and the number of registers is unlimited. This generalization of registers is sometimes
; called "pseudo registers". Our model uses pseudo registers because it makes writing example
; programs easier and comes at no loss of generality. The heap memory in TPAL-0 stores only
; code segments, and is read only.

(define-language TPAL-0

  ; Instructions and operands
  (v ::= n l r j)                               ; operands
  (r ::=                                        ; pseudo registers (must start with the letter "r")
     (side-condition (name n variable-not-otherwise-mentioned)
                     (eq? #\r (string-ref (symbol->string (term n)) 0))))
  (l ::=                                        ; labels (must start with the letter "l")
     (side-condition (name n variable-not-otherwise-mentioned)
                     (eq? #\l (string-ref (symbol->string (term n)) 0))))
  (n ::= number)                                ; integer literals
  (j ::=                                        ; join-record identifiers (must start with the letter "j")
     (side-condition (name n variable-not-otherwise-mentioned)
                     (eq? #\j (string-ref (symbol->string (term n)) 0))))
  (⊙ ::= + - * / % < > <= >= == <>)             ; binary operators
  (i ::=                                        ; instructions:
     (:= r v)                                   ;   assignment
     (:= r (⊙ r v))                             ;   binary operation
     (:= r (jralloc l))                         ;   join-record creation
     (fork r v)                                 ;   task creation
     (if-jump r v))                             ;   conditional branch
  (I ::=                                        ; instruction sequences:
     (jump v)                                   ;   jump
     (join v)                                   ;   join with other tasks
     (halt)                                     ;   machine halt
     (i I))                                     ;   sequence

  ; Code blocks
  (jp ::= assoc assoc-commute)                  ; join-resolution policies
  (ΔR ::= ((r r) ...))                          ; register relocation mappings
  (★ ::=                                        ; code-block annotations:
     ·                                          ;   empty
     (prppt l)                                  ;   promotion-ready program point
     (jtppt                                     ;   join-target program point:
      jp                                        ;     policy for combining task results
      ΔR                                        ;     register relocation map
      l))                                       ;     label of the combining block
  (B ::= (code ★ I))                            ; code blocks
  (h ::= B))                                    ; heap values

; TPAL-1 syntax
; -------------
; This syntax extends TPAL-0 to support stack-based storage.

(define-extended-language TPAL-1
  TPAL-0
  
  (v ::= ....
     (uptr h)             ; unique pointers
     (prmark))            ; promotion-ready stack marks

  (i ::= ....
     (:= (mem r v) r)     ; memory store
     (:= r (mem r v))     ; memory load
     (salloc r n)         ; stack allocation
     (sfree r n)          ; stack deallocation
     (:= r (prmempty r))  ; check for promotion-ready marks in the stack
     (prmsplit r r r))    ; split the stack by the least recent promotion-ready mark
  
  (h ::= ....
     (tup v ...)))        ; tuples

; TPAL
; ----
; This syntax specifies register files, heaps, and task records.

(define-extended-language TPAL
  TPAL-1

  (R ::= ((r v) ...))                           ; register files
  (H ::= ((l h) ...))                           ; heaps
  (l+ ::= · l)                                  ; label-or-nil option
  (◇ ::= natural)                               ; cycle counters
  (T ::=                                        ; tasks:
     (l+                                        ;   label of enclosing block (non nil only if at start of the block)
      ◇                                         ;   cycle counter
      H                                         ;   private heap
      R                                         ;   register file
      I)))                                      ;   continuation

; TPAL-⇓
; ------
; This syntax specifies activation records for the join points of parallel tasks and
; the cost-graph data structure that specifies the cost model of TPAL.

(define-extended-language TPAL-⇓
  TPAL

  (js ::=                                       ; join statuses:
      open-join                                 ;   open, if ∃ at least two running tasks on join, and
      closed-join)                              ;   closed, otherwise
  (jr ::=                                       ; join-activation records:
      (j                                        ;   unique id of the join
       l                                        ;   label of the continuation block of the join
       js))                                     ;   status of the join
  (J ::= (jr ...))                              ; join-activation-record maps

  (g ::=                                        ; cost graphs:
     $0                                         ;   identity cost
     $1                                         ;   unit cost
     (∘ g g)                                    ;   sequential composition
     (∥ g g)))                                  ;   parallel composition

; Metafunctions for dynamic semantics
; -----------------------------------

(define-metafunction TPAL
  Lookup-in-R-by-r : R r -> v
  [(Lookup-in-R-by-r ((r_b v_b) ... (r_1 v_1) (r_a v_a) ...) r_1) v_1])

(define-metafunction TPAL
  Lookup-in-R-by-v : R v -> v
  [(Lookup-in-R-by-v R_1 r_1) (Lookup-in-R-by-r R_1 r_1)]
  [(Lookup-in-R-by-v R_1 v) v])

(define-metafunction TPAL
  Sort-R : R -> R
  [(Sort-R R_1) ,(sort (term R_1) #:key car symbol<?)])

(define-metafunction TPAL
  Update-R-by-r-helper : R r v -> R
  [(Update-R-by-r-helper ((r_b v_b) ... (r_t v_t) (r_a v_a) ...) r_t v_t2)
   ((r_b v_b) ... (r_t v_t2) (r_a v_a) ...)]
  [(Update-R-by-r-helper ((r_b v_b) ...) r_t v_t)
   ((r_t v_t) (r_b v_b) ...)])

(define-metafunction TPAL
  Update-R-by-r : R r v -> R
  [(Update-R-by-r R_1 r_1 v_1) (Sort-R (Update-R-by-r-helper R_1 r_1 v_1))])

(define-metafunction TPAL
  Lookup-in-H : H l -> h
  [(Lookup-in-H ((l_b h_b) ... (l_1 h_1) (l_a h_a) ...) l_1) h_1])

(define-metafunction TPAL
  Load-jump-target : H R v -> (l B)
  [(Load-jump-target H_1 R_1 v_t)
   (l_t B_t)
   (where l_t (Lookup-in-R-by-v R_1 v_t))
   (where B_t (Lookup-in-H H_1 l_t))
   (where (code ★ I_t) B_t)])

(define-metafunction TPAL
  Lookup-★-in-H : H l -> ★ or #f
  [(Lookup-★-in-H H l) ★
   (where (code ★ _) (Lookup-in-H H l))]
  [(Lookup-★-in-H H l) #f])

(define-metafunction TPAL
  Apply-ΔR-to-T2-helper : R ΔR -> R
  [(Apply-ΔR-to-T2-helper ((r_s v_s) (r_a v_a) ...) ((r_s r_t) (r_Δas r_Δat) ...))
   ((r_t v_s) (r_Δas2 r_Δat2) ...)
   (where ((r_Δas2 r_Δat2) ...) (Apply-ΔR-to-T2-helper ((r_a v_a) ...) ((r_Δas r_Δat) ...)))]
  [(Apply-ΔR-to-T2-helper (_ (r_a v_a) ...) ΔR)
   (Apply-ΔR-to-T2-helper ((r_a v_a) ...) ΔR)]
  [(Apply-ΔR-to-T2-helper () ())
   ()])

(define-metafunction TPAL
  Apply-ΔR-to-T2 : R ΔR -> R
  [(Apply-ΔR-to-T2 R ΔR) (Apply-ΔR-to-T2-helper (Sort-R R) (Sort-R ΔR))])

(define-metafunction TPAL
  Apply-ΔR-to-T1 : R ΔR -> R
  [(Apply-ΔR-to-T1 R ΔR)
   ,(filter-not (λ (rp) (member rp (term ΔR) (λ (p1 p2) (equal? (car p1) (cadr p2))))) (term R))])

(define-metafunction TPAL
  Merge-Rs-by-ΔR : R R ΔR -> R
  [(Merge-Rs-by-ΔR R_1 R_2 ΔR)
   (Sort-R ((r_1 v_1) ... (r_2 v_2) ...))
   (where ((r_1 v_1) ...) (Apply-ΔR-to-T1 R_1 ΔR))
   (where ((r_2 v_2) ...) (Apply-ΔR-to-T2 R_2 ΔR))])

(define-metafunction TPAL
  Machine-bool : boolean -> number
  [(Machine-bool #t) 0]
  [(Machine-bool #f) 1])

(define-metafunction TPAL
  Apply-⊙ : ⊙ n n -> n
  [(Apply-⊙ + n_1 n_2) ,(+ (term n_1) (term n_2))]
  [(Apply-⊙ - n_1 n_2) ,(- (term n_1) (term n_2))]
  [(Apply-⊙ * n_1 n_2) ,(* (term n_1) (term n_2))]
  [(Apply-⊙ / n_1 n_2) ,(floor (/ (term n_1) (term n_2)))]
  [(Apply-⊙ % n_1 n_2) ,(remainder (term n_1) (term n_2))]
  [(Apply-⊙ < n_1 n_2) (Machine-bool ,(< (term n_1) (term n_2)))]
  [(Apply-⊙ > n_1 n_2) (Machine-bool ,(> (term n_1) (term n_2)))]
  [(Apply-⊙ <= n_1 n_2) (Machine-bool ,(<= (term n_1) (term n_2)))]
  [(Apply-⊙ >= n_1 n_2) (Machine-bool ,(>= (term n_1) (term n_2)))]
  [(Apply-⊙ == n_1 n_2) (Machine-bool ,(= (term n_1) (term n_2)))]
  [(Apply-⊙ <> n_1 n_2) (Machine-bool ,(not (= (term n_1) (term n_2))))])

(define-metafunction TPAL
  Tick : ◇ -> ◇
  [(Tick ◇) ,(add1 (term ◇))])

(define ◇-identity 0)

(define-metafunction TPAL
  Promotion-ready? : H l+ ◇ ◇ -> boolean
  [(Promotion-ready? H l_1 ◇_♥ ◇)
   ,(>= (term ◇) (term ◇_♥))
   (where (prppt _) (Lookup-★-in-H H l_1))]
  [(Promotion-ready? _ _ _ _)
   #f])

(define-metafunction TPAL
  Not-promotion-ready? : H l+ ◇ ◇ -> boolean
  [(Not-promotion-ready? H l+ ◇_♥ ◇)
   ,(not (term (Promotion-ready? H l+ ◇_♥ ◇)))])

(define-metafunction TPAL-⇓
  Sort-J : J -> J
  [(Sort-J J_1) ,(sort (term J_1) #:key car symbol<?)])

(define-metafunction TPAL-⇓
  Lookup-in-J : J j -> jr or #f
  [(Lookup-in-J (jr_b ... (j l js) jr_a ...) j) (j l js)]
  [(Lookup-in-J _ _) #f])

(define-metafunction TPAL-⇓
  Update-J-helper : J jr -> J
  [(Update-J-helper (jr_b ... (j _ _) jr_a ...) (j l js))
   (jr_b ... (j l js) jr_a ...)]
  [(Update-J-helper (jr ...) (j l js))
   ((j l js) jr ...)])

(define-metafunction TPAL-⇓
  Update-J : J jr -> J
  [(Update-J J jr) (Sort-J (Update-J-helper J jr))])

(define-metafunction TPAL-⇓
  Merge-Js : J J -> J
  [(Merge-Js (jr_1 ...) J_2)
   ,(foldl (λ (jr J) (term (Update-J ,J ,jr))) (term J_2) (term (jr_1 ...)))])

(define-metafunction TPAL-⇓
  Remove-jr-in-J : J j -> J or #f
  [(Remove-jr-in-J (jr_b ... (j l js) jr_a ...) j) (jr_b ... jr_a ...)]
  [(Remove-jr-in-J _ _) #f])

(define-metafunction TPAL-⇓
  Split-by-promotion-ready-mark : (v ...) -> ((v ...) (v ...)) or #f
  [(Split-by-promotion-ready-mark (v ...))
   ((v_b ...) (v_a ...))
   (where (n_1 n_i ...) ,(indexes-where (term (v ...)) (λ (x) (equal? x (term (prmark))))))
   (where n ,(last (term (n_1 n_i ...))))
   (where (v_b ...) ,(take (term (v ...)) (term n)))
   (where (v_a ...) ,(cdr (drop (term (v ...)) (term n))))]
  [(Split-by-promotion-ready-mark (v ...))
   #f])

(define-metafunction TPAL-⇓
  Work : g -> natural
  [(Work $0) 0]
  [(Work $1) 1]
  [(Work (∘ g_1 g_2)) ,(+ (term (Work g_1)) (term (Work g_2)))]
  [(Work (∥ g_1 g_2)) ,(+ (term (Work g_1)) (term (Work g_2)))])

(define-metafunction TPAL-⇓
  Span : g -> natural
  [(Span $0) 0]
  [(Span $1) 1]
  [(Span (∘ g_1 g_2)) ,(+ (term (Span g_1)) (term (Span g_2)))]
  [(Span (∥ g_1 g_2)) ,(max (term (Span g_1)) (term (Span g_2)))])

; Big-step dynamic semantics
; --------------------------

(define-judgment-form TPAL-⇓
  #:mode (⇓-0 I I I O O O)
  #:contract (⇓-0 ◇ J T J T g)

  [------------------------------------------------ HALT
   (⇓-0 ◇_♥ J (_ ◇ H R (halt)) J (· ◇ H R (halt)) $0)]

  [(side-condition (Not-promotion-ready? H l+ ◇_♥ ◇))
   (where R_2 (Update-R-by-r R_1 r_d (Lookup-in-R-by-v R_1 v_s)))
   (⇓-0 ◇_♥ J (· (Tick ◇) H R_2 I) J_2 T g)
   -------------------------------------------------------------- MOV
   (⇓-0 ◇_♥ J (l+ ◇ H R_1 ((:= r_d v_s) I)) J_2 T (∘ $1 g))]

  [(side-condition (Not-promotion-ready? H l+ ◇_♥ ◇))
   (where n_1 (Lookup-in-R-by-r R_1 r_s1))
   (where n_2 (Lookup-in-R-by-v R_1 v_s2))
   (where n_3 (Apply-⊙ ⊙ n_1 n_2))
   (where R_2 (Update-R-by-r R_1 r_d n_3))
   (⇓-0 ◇_♥ J (· (Tick ◇) H R_2 I) J_2 T g)
   ------------------------------------------------------------- BINOP
   (⇓-0 ◇_♥ J (l+ ◇ H R_1 ((:= r_d (⊙ r_s1 v_s2)) I)) J_2 T (∘ $1 g))]

  [(side-condition (Not-promotion-ready? H l+ ◇_♥ ◇))
   (where (l_t (code _ I_t)) (Load-jump-target H R v))
   (⇓-0 ◇_♥ J (l_t (Tick ◇) H R I_t) J_2 T g)
   ----------------------------------------------------- JUMP
   (⇓-0 ◇_♥ J (l+ ◇ H R (jump v)) J_2 T (∘ $1 g))]

  [(side-condition (Not-promotion-ready? H l+ ◇_♥ ◇))
   (where (l_t (code _ I_t)) (Load-jump-target H R v_t))
   (side-condition ,(equal? (term (Lookup-in-R-by-r R r_c)) 0))
   (⇓-0 ◇_♥ J (l_t (Tick ◇) H R I_t) J_2 T g)
   ------------------------------------------------------------ IF-EQ
   (⇓-0 ◇_♥ J (l+ ◇ H R ((if-jump r_c v_t) I_1)) J_2 T (∘ $1 g))]

  [(side-condition (Not-promotion-ready? H l+ ◇_♥ ◇))
   (side-condition ,(not (equal? (term (Lookup-in-R-by-r R r_c)) 0)))
   (⇓-0 ◇_♥ J (· (Tick ◇) H R I) J_2 T g)
   ------------------------------------------------------------------ IF-NEQ
   (⇓-0 ◇_♥ J (l+ ◇ H R ((if-jump r_c _) I)) J_2 T (∘ $1 g))]

  [(side-condition (Not-promotion-ready? H l+ ◇_♥ ◇))
   (where j ,(variable-not-in (term J) (term j)))
   (where J_1 (Update-J J (j l closed-join)))
   (where R_2 (Update-R-by-r R_1 r j))
   (⇓-0 ◇_♥ J_1 (· ◇ H R_2 I) J_2 T g)
   --------------------------------------------------------- JOIN-RECORD-ALLOC
   (⇓-0 ◇_♥ J (l+ ◇ H R_1 ((:= r (jralloc l)) I)) J_2 T (∘ $1 g))]

  [(side-condition (Promotion-ready? H l_1 ◇_♥ ◇))
   (where (prppt l_t) (Lookup-★-in-H H l_1))
   (where (l_t (code _ I_t)) (Load-jump-target H R l_t))
   (⇓-0 ◇_♥ J (l_t ,◇-identity H R I_t) J_2 T g)
   ----------------------------------------------------- TRY-PROMOTE
   (⇓-0 ◇_♥ J (l_1 ◇ H R I) J_2 T (∘ $1 g))]

  [(side-condition (Not-promotion-ready? H l+ ◇_♥ ◇))
   (where j (Lookup-in-R-by-r R r_j))
   (where (j l_cont js_0) (Lookup-in-J J j))
   (where J_0 (Update-J J (j l_cont open-join)))
   (where (l_t (code _ I_t)) (Load-jump-target H R v_t))
   (⇓-0 ◇_♥ J_0 (· ,◇-identity H R I) J_1 (_ ◇_1 H_1 R_1 (join r_j1)) g_1)
   (⇓-0 ◇_♥ J_0 (l_t ,◇-identity H R I_t) J_2 (_ ◇_2 H_2 R_2 (join r_j2)) g_2)
   (where j (Lookup-in-R-by-r R_1 r_j1))
   (where j (Lookup-in-R-by-r R_2 r_j2))
   (where (jtppt _ ΔR l_comb) (Lookup-★-in-H H l_cont))
   (where (l_comb (code _ I_comb)) (Load-jump-target H R l_comb))
   (where R_cont (Merge-Rs-by-ΔR R_1 R_2 ΔR))
   (where J_cont (Update-J (Merge-Js J_1 J_2) (j l_cont js_0)))
   (⇓-0 ◇_♥ J_cont (· ,◇-identity H R_cont I_comb) J_3 T_3 g_3)
   --------------------------------------------------------------------- FORK
   (⇓-0 ◇_♥ J (l+ ◇ H R ((fork r_j v_t) I)) J_3 T_3 (∘ $1 (∘ (∥ g_1 g_2) g_3)))]

  [(side-condition (Not-promotion-ready? H l+ ◇_♥ ◇))
   (where j (Lookup-in-R-by-r R r_j))
   (where (jr_b ... (j l open-join) jr_a ...) J)
   ------------------------------------------------------------ JOIN-AND-BLOCK
   (⇓-0 ◇_♥ J (l+ ◇ H R (join r_j)) J (· ◇ H R (join r_j)) $1)]

  [(side-condition (Not-promotion-ready? H l+ ◇_♥ ◇))
   (where j (Lookup-in-R-by-r R r_j))
   (where (j l closed-join) (Lookup-in-J J j))
   (where J_1 (Remove-jr-in-J J j))
   (where (l (code _ I)) (Load-jump-target H R l))
   (⇓-0 ◇_♥ J_1 (· ◇ H R I) J_2 T_2 g_2)
   ------------------------------------------------------------ JOIN-AND-CONTINUE
   (⇓-0 ◇_♥ J (l+ ◇ H R (join r_j)) J_2 T_2 (∘ $1 g_2))])

(define-extended-judgment-form TPAL-⇓ ⇓-0
  #:mode (⇓ I I I O O O)
  #:contract (⇓ ◇ J T J T g)
  
  [(side-condition (Not-promotion-ready? H l+ ◇_♥ ◇))
   (where (uptr (tup v_1 ...)) (Lookup-in-R-by-r R_1 r))
   (where (v_a ...) ,(append (make-list (term n) 0) (term (v_1 ...))))
   (where v_2 (uptr (tup v_a ...)))
   (where R_2 (Update-R-by-r R_1 r v_2))
   (⇓ ◇_♥ J (· (Tick ◇) H R_2 I) J_2 T g_2)
   ------------------------------------------------------------------- STACK-ALLOC
   (⇓ ◇_♥ J (l+ ◇ H R_1 ((salloc r n) I)) J_2 T (∘ $1 g_2))]

  [(side-condition (Not-promotion-ready? H l+ ◇_♥ ◇))
   (where (uptr (tup v_stk ...)) (Lookup-in-R-by-r R_1 r))
   (side-condition ,(<= (term n) (length (term (v_stk ...)))))
   (where (v_stk2 ...) ,(drop (term (v_stk ...)) (term n)))
   (where R_2 (Update-R-by-r R_1 r (uptr (tup v_stk2 ...))))
   (⇓ ◇_♥ J (· (Tick ◇) H R_2 I) J_2 T g_2)
   ------------------------------------------------------------ STACK-FREE
   (⇓ ◇_♥ J (l+ ◇ H R_1 ((sfree r n) I)) J_2 T (∘ $1 g_2))]
  
  [(side-condition (Not-promotion-ready? H l+ ◇_♥ ◇))
   (where (uptr (tup v_stk ...)) (Lookup-in-R-by-r R_1 r_d))
   (where n_d (Lookup-in-R-by-v R_1 v_d))
   (side-condition ,(< (term n_d) (length (term (v_stk ...)))))
   (where v_s (Lookup-in-R-by-r R_1 r_s))
   (where (v_stk2 ...) ,(list-set (term (v_stk ...)) (term n_d) (term v_s)))
   (where R_2 (Update-R-by-r R_1 r_d (uptr (tup v_stk2 ...))))
   (⇓ ◇_♥ J (· (Tick ◇) H R_2 I) J_2 T g_2)
   ---------------------------------------------------------------------- STACK-STORE
   (⇓ ◇_♥ J (l+ ◇ H R_1 ((:= (mem r_d v_d) r_s) I)) J_2 T (∘ $1 g_2))]

  [(side-condition (Not-promotion-ready? H l+ ◇_♥ ◇))
   (where (uptr (tup v_stk ...)) (Lookup-in-R-by-r R_1 r_s))
   (where n_s (Lookup-in-R-by-v R_1 v_s))
   (side-condition ,(< (term n_s) (length (term (v_stk ...)))))
   (where v_t ,(list-ref (term (v_stk ...)) (term n_s)))
   (where R_2 (Update-R-by-r R_1 r_d v_t))
   (⇓ ◇_♥ J (· (Tick ◇) H R_2 I) J_2 T g_2)
   ------------------------------------------------------------------ STACK-LOAD
   (⇓ ◇_♥ J (l+ ◇ H R_1 ((:= r_d (mem r_s v_s)) I)) J_2 T (∘ $1 g_2))]

  [(side-condition (Not-promotion-ready? H l+ ◇_♥ ◇))
   (where (uptr (tup v_1 ...)) (Lookup-in-R-by-r R_1 r))
   (where boolean_1 ,(if (term (Split-by-promotion-ready-mark (v_1 ...))) #f #t))
   (where R_2 (Update-R-by-r R_1 r_d (Machine-bool boolean_1)))
   (⇓ ◇_♥ J (· (Tick ◇) H R_2 I) J_2 T g_2)
   ------------------------------------------------------------------------------- PR-MARK-EMPTY
   (⇓ ◇_♥ J (l+ ◇ H R_1 ((:= r_d (prmempty r)) I)) J_2 T (∘ $1 g_2))]

  [(side-condition (Not-promotion-ready? H l+ ◇_♥ ◇))
   (where (uptr (tup v_1 ...)) (Lookup-in-R-by-r R_1 r_s))
   (where ((v_b ...) (v_a ...)) (Split-by-promotion-ready-mark (v_1 ...)))
   (where R_2 (Update-R-by-r R_1 r_s (uptr (tup v_b ...))))
   (where R_3 (Update-R-by-r R_2 r_d (uptr (tup v_a ...))))
   (where R_4 (Update-R-by-r R_3 r_p ,(sub1 (length (term (v_b ...))))))
   (⇓ ◇_♥ J (· (Tick ◇) H R_4 I) J_2 T g_2)
   ----------------------------------------------------------------------- PR-MARK-SPLIT
   (⇓ ◇_♥ J (l+ ◇ H R_1 ((prmsplit r_s r_d r_p) I)) J_2 T (∘ $1 g_2))])

; TPAL-0 programs
; ---------------

(define-metafunction TPAL
  ▹ : i ... I -> I
  [(▹ I) I]
  [(▹ i_h i_t ... I) (i_h (▹ i_t ... I))])

(define simple-H
  (term
   ([l-entry
     (code ·
           (▹
            (:= r1 0)
            (:= r1 (+ r1 1))
            (jump l-body)))]
    [l-body
     (code (prppt l-try-promote)
           (▹
            (:= r1 (+ r1 1))
            (jump l-exit)))]
    [l-exit
     (code (jtppt assoc-commute ((r1 r2)) l-comb)
           (halt))]
    [l-body-par
     (code ·
           (▹
            (:= r1 (+ r1 1))
            (join r-j)))]
    [l-try-promote
     (code ·
           (▹
            (:= r-j (jralloc l-exit))
            (fork r-j l-body-par)
            (jump l-body-par)))]
    [l-comb
     (code ·
           (▹
            (join r-j)))])))

(define simple-T
  (term (· 0 ,simple-H () (jump l-entry))))

; Parallel product
; ~~~~~~~~~~~~~~~~
;
; Compute c = a × b for integers a and b in registers r-a and r-b, respectively, leaving
; the result c in register r-c.

(define (Mk-prod-H-seq r-ret ★-loop ★-exit)
  (term
   ([l-entry
     (code ·
           (▹
            (:= r-acc 0)
            (jump l-loop)))]
    [l-loop
     (code ,★-loop
           (▹
            (if-jump r-a l-exit)
            (:= r-acc (+ r-b r-acc))
            (:= r-a (- r-a 1))
            (jump l-loop)))]   
    [l-exit
     (code ,★-exit
           (▹
            (:= r-c r-acc)
            (jump ,r-ret)))])))

(define (Mk-prod-H-par r-a ★-loop-par)
  (term
   ([l-loop-try-promote
     (code ·
           (▹
            (:= r-t (< ,r-a 2))
            (if-jump r-t l-loop)
            (:= r-j (jralloc l-exit))
            (jump l-loop-promote)))]
    [l-comb
     (code ·
           (▹
            (:= r-acc (+ r-acc r-acc2))
            (join r-j)))]
    [l-loop-par
     (code ,★-loop-par
           (▹
            (if-jump ,r-a l-join)
            (:= r-acc (+ r-b r-acc))
            (:= ,r-a (- ,r-a 1))
            (jump l-loop-par)))]
    [l-loop-par-try-promote
     (code ·
           (▹
            (:= r-t (< ,r-a 2))
            (if-jump r-t l-loop-par)
            (jump l-loop-promote)))]
    [l-loop-promote
     (code ·
           (▹
            (:= r-m (/ ,r-a 2))
            (:= r-n (% ,r-a 2))
            (:= ,r-a r-m)
            (:= r-t-acc r-acc)
            (:= r-acc 0)
            (fork r-j l-loop-par)
            (:= ,r-a (+ r-m r-n))
            (:= r-acc r-t-acc)
            (jump l-loop-par)))]
    [l-join
     (code ·
           (▹
            (join r-j)))])))

(define (Mk-prod-H ★-loop ★-loop-par)
  (append
   (Mk-prod-H-seq
    (term r-ret)
    ★-loop
    (term (jtppt assoc-commute ((r-acc r-acc2)) l-comb)))
   (Mk-prod-H-par
    (term r-a)
    ★-loop-par)))

(define prod-H
  (Mk-prod-H (term (prppt l-loop-try-promote)) (term (prppt l-loop-par-try-promote))))

(define (prod-R-init a b)
  (term ((r-a ,a) (r-b ,b))))

(define (prod-T-init a b)
  (let* ([halt-B
          (term ((l-halt
                  (code ·
                        (halt)))))]
         [H (append prod-H halt-B)]
         [R (prod-R-init a b)]
         [I (term ((:= r-ret l-halt) (jump l-entry)))])
    (term (· 0 ,H ,R ,I))))

; Parallel raise-to-the-power
; ~~~~~~~~~~~~~~~~~~~~~~~~~~~
;
; Compute f = d ^ e for integers d and e in registers r-d and r-e, respectively, leaving
; the result f in register r-f.

(define (Mk-pow-H-seq ★-pow-loop ★-pow-exit)
  (term
   ([l-pow-entry
     (code ·
           (▹
            (:= r-pow-acc 1)
            (:= r-pow-j 0)
            (jump l-pow-loop)))]
    [l-pow-loop
     (code ,★-pow-loop
           (▹
            (if-jump r-e l-pow-exit)
            ; jump to prod loop, which will return to l-pow-loop-cont
            (:= r-a r-d)
            (:= r-b r-pow-acc)
            (:= r-ret l-pow-loop-cont)
            (jump l-entry)))]
    [l-pow-loop-cont
     (code ·
           (▹
            (:= r-pow-acc r-c)
            (:= r-e (- r-e 1))
            (jump l-pow-loop)))]
    [l-pow-exit
     (code ,★-pow-exit
           (▹
            (:= r-f r-pow-acc)
            (halt)))])))

(define (Mk-pow-H-par ★-pow-loop-par)
  (term
   ([l-pow-loop-try-promote-0
     (code ·
           (▹
            (:= r-promote-abort l-pow-loop)
            (:= r-pow-loop-promote-cont l-pow-loop-par)
            (if-jump r-pow-j l-pow-loop-try-promote)
            (:= r-promote-abort l-pow-loop-par)
            (jump l-pow-loop-par-try-promote)))]
    [l-pow-loop-try-promote
     (code ·
           (▹
            (:= r-t (< r-e 2))
            (if-jump r-t r-promote-abort)
            (:= r-pow-j (jralloc l-pow-exit))
            (jump l-pow-loop-promote)))]
    [l-pow-comb
     (code ·
           (▹
            (:= r-pow-acc (* r-pow-acc r-pow-acc2))
            (join r-pow-j)))]
    [l-pow-loop-par
     (code ,★-pow-loop-par
           (▹
            (if-jump r-e l-pow-join)
            ; jump to prod loop, which will return to l-pow-loop-cont
            (:= r-a r-d)
            (:= r-b r-pow-acc)
            (:= r-ret l-pow-loop-par-cont)
            (jump l-entry)))]
    [l-pow-loop-par-cont
     (code ·
           (▹
            (:= r-pow-acc r-c)
            (:= r-e (- r-e 1))
            (jump l-pow-loop-par)))]
    [l-pow-loop-par-try-promote
     (code ·
           (▹
            (:= r-t (< r-e 2))
            (if-jump r-t r-promote-abort)
            (jump l-pow-loop-promote)))]
    [l-pow-loop-promote
     (code ·
           (▹
            (:= r-ret l-pow-loop-par-cont) ; hack? this instr applies only to inner (prod) promotion
            (:= r-m (/ r-e 2))
            (:= r-n (% r-e 2))
            (:= r-e r-m)
            (:= r-t-acc r-pow-acc)
            (:= r-pow-acc 1)
            (fork r-pow-j l-pow-loop-par)
            (:= r-e (+ r-m r-n))
            (:= r-pow-acc r-t-acc)
            (jump r-pow-loop-promote-cont)))]
    [l-pow-join
     (code ·
           (▹
            (join r-pow-j)))]
    [l-loop-try-promote-via-pow ; to override the prppt of prod
     (code ·
           (▹
            (:= r-promote-abort l-loop-try-promote)
            (:= r-pow-loop-promote-cont l-loop)
            (if-jump r-pow-j l-pow-loop-try-promote)
            (jump l-pow-loop-par-try-promote)))]
    [l-loop-par-try-promote-via-pow ; to override the prppt of prod
     (code ·
           (▹
            (:= r-promote-abort l-loop-par-try-promote)
            (:= r-pow-loop-promote-cont l-loop-par)
            (if-jump r-pow-j l-pow-loop-try-promote)
            (jump l-pow-loop-par-try-promote)))])))

(define (pow-R-init d e)
  (term ((r-d ,d) (r-e ,e))))

(define (pow-T-init d e)
  (let* ([pow-H-seq (Mk-pow-H-seq
                     (term (prppt l-pow-loop-try-promote-0))
                     (term (jtppt assoc-commute ((r-pow-acc r-pow-acc2)) l-pow-comb)))]
         [pow-H-par (Mk-pow-H-par (term (prppt l-pow-loop-try-promote-0)))]
         [pow-H (append pow-H-seq pow-H-par)]
         [prod-H (Mk-prod-H (term (prppt l-loop-try-promote-via-pow))
                            (term (prppt l-loop-par-try-promote-via-pow)))]
         [H (append prod-H pow-H)]
         [R (pow-R-init d e)]
         [I (term (jump l-pow-entry))])
    (term (· 0 ,H ,R ,I))))

; TPAL-1 programs
; ---------------

(define stack-H
  (term
   ([l-entry
     (code ·
           (▹
            (:= r-sp (uptr (tup)))
            (salloc r-sp 2)
            (halt)))])))

(define stack-T-init
  (term
   (· 0 ,stack-H () (jump l-entry))))

; Fibonacci function
; ~~~~~~~~~~~~~~~~~~~~~~~~~~~
; Compute fib(n) s.t.
;   fib(0) = 0
;   fib(1) = 1
;   fib(n) = fib(n-1) + fib(n-2)
; for the integer n in r-n, leaving the result in register r-d.

(define t-fib-exit 990)
(define t-fib-branch1 991)
(define t-fib-branch2 992)

(define (Mk-fib-H-loop l-loop l-pop l-branch1 l-exit ★-loop ★-pop)
  (term
   ([,l-loop
     (code ,★-loop
           (▹
            (:= r-res r-n)
            (:= r-tmp (< r-n 2))
            (if-jump r-tmp ,l-pop)
            (:= r-res 0)
            (salloc r-sp 3)
            (:= r-tmp ,t-fib-branch1)
            (:= (mem r-sp 0) r-tmp)
            (:= r-tmp (- r-n 2))
            (:= (mem r-sp 1) r-tmp)
            (:= r-tmp (prmark))
            (:= (mem r-sp 2) r-tmp)
            (:= r-n (- r-n 1))
            (jump ,l-loop)))]
    [,l-pop
     (code ,★-pop
           (▹
            (:= r-tag (mem r-sp 0))
            (:= r-tmp (== r-tag ,t-fib-exit))
            (if-jump r-tmp ,l-exit)
            (:= r-tmp (== r-tag ,t-fib-branch1))
            (if-jump r-tmp ,l-branch1)
            ; branch2
            (:= r-tmp (mem r-sp 1))
            (:= r-res (+ r-res r-tmp))
            (sfree r-sp 3)
            (jump ,l-pop)))]
    [,l-branch1
     (code ·
           (▹
            (:= r-n (mem r-sp 1))
            (:= r-tmp ,t-fib-branch2)
            (:= (mem r-sp 0) r-tmp)
            (:= (mem r-sp 1) r-res)
            (:= r-tmp 0)
            (:= (mem r-sp 2) r-tmp) ; to overwrite the mark
            (jump ,l-loop)))])))

(define fib-H-par
  (term
   ([l-fib
     (code ·
           (▹
            (:= r-sp (uptr (tup)))
            (salloc r-sp 1)
            (:= r-tmp ,t-fib-exit)
            (:= (mem r-sp 0) r-tmp)
            (jump l-loop)))]
    [l-exit
     (code ·
           (▹
            (sfree r-sp 1)
            (:= r-ans r-res)
            (jump r-ret)))]
    [l-loop-try-promote
     (code ·
           (▹
            (:= r-tmp (prmempty r-sp))
            (if-jump r-tmp l-loop)
            (:= r-j (jralloc l-pop))
            (:= r-n1 r-n)
            (prmsplit r-sp r-sp-cont r-addr)
            (:= r-addr-tag (- r-addr 1))
            (:= r-tmp ,t-fib-exit)
            (:= (mem r-sp r-addr-tag) r-tmp)
            (:= r-addr-arg (- r-addr 0))
            (:= r-n (mem r-sp r-addr-arg))
            (:= r-sp1 r-sp)
            (:= r-sp (uptr (tup)))
            (salloc r-sp 1)
            (:= r-tmp ,t-fib-exit)
            (:= (mem r-sp 0) r-tmp)
            (fork r-j l-loop-par)
            (:= r-sp r-sp1)
            (:= r-n r-n1)
            (jump l-loop-par)))]
    [l-loop-par-try-promote
     (code ·
           (▹
            (:= r-tmp (prmempty r-sp))
            (if-jump r-tmp l-loop-par)
            (:= r-n1 r-n)
            (prmsplit r-sp r-sp-cont r-addr)
            (:= r-addr-tag (- r-addr 1))
            (:= r-tmp ,t-fib-exit)
            (:= (mem r-sp r-addr-tag) r-tmp)
            (:= r-addr-arg (- r-addr 0))
            (:= r-n (mem r-sp r-addr-arg))
            (:= r-sp1 r-sp)
            (:= r-sp (uptr (tup)))
            (salloc r-sp 1)
            (:= r-tmp ,t-fib-exit)
            (:= (mem r-sp 0) r-tmp)
            (fork r-j l-loop-par)
            (:= r-sp r-sp1)
            (:= r-n r-n1)
            (jump l-loop-par)))]
    [l-comb
     (code ·
           (▹
            (:= r-res (+ r-res r-res2))
            (join r-j)))]
    [l-join
     (code ·
           (▹
            (:= r-sp r-sp-cont)
            (join r-j)))])))

(define fib-H
  (let* ([fib-H-loop-seq
          (Mk-fib-H-loop (term l-loop)
                         (term l-pop)
                         (term l-branch1)
                         (term l-exit)
                         (term (prppt l-loop-try-promote))
                         (term (jtppt assoc-commute ((r-res r-res2)) l-comb)))]
         [fib-H-loop-par
          (Mk-fib-H-loop (term l-loop-par)
                         (term l-pop-par)
                         (term l-branch1-par)
                         (term l-join)
                         (term (prppt l-loop-par-try-promote))
                         (term (jtppt assoc-commute ((r-res r-res2)) l-comb)))])
    (append
     fib-H-loop-seq
     fib-H-loop-par
     fib-H-par)))

(define (fib-R-init n)
  (term ((r-n ,n))))

(define (fib-T-init n)
  (let* ([halt-B
          (term ((l-halt
                  (code ·
                        (halt)))))]
         [H (append fib-H halt-B)]
         [R (fib-R-init n)]
         [I (term ((:= r-ret l-halt) (jump l-fib)))])
    (term (· 0 ,H ,R ,I))))

(define (fib n)
  (if (< n 2)
      n
      (+ (fib (- n 1)) (fib (- n 2)))))

; Testing and debugging
; ---------------------

(define-metafunction TPAL
  Lookup-ans-in-R-of-T : T r -> v or #f
  [(Lookup-ans-in-R-of-T (_ _ _ R _) r) v
   (where v (Lookup-in-R-by-r R r))]
  [(Lookup-ans-in-R-of-T _ _) #f])

(define (test-TPAL ♥ T0 r a)
  (let ([p (λ (Ts c)
             (begin
               (if (not (= (length Ts) 1))
                   (printf "error: zero or multiple results~n")
                   '())
               (let* ([T (car Ts)]
                      [ans (term (Lookup-ans-in-R-of-T ,T ,r))])
                 (equal? ans c))))])
    (test-equal (judgment-holds (⇓ ,♥ () ,T0 J T g) T) a  #:equiv p)))

(define (test-prod ♥ a b)
  (test-TPAL ♥ (prod-T-init a b) (term r-c) (* a b)))

(define (test-pow ♥ d e)
  (test-TPAL ♥ (pow-T-init d e) (term r-f) (expt d e)))

(define (test-fib ♥ n)
  (test-TPAL ♥ (fib-T-init n) (term r-ans) (fib n)))

; Cost analysis (work and span)
; -----------------------------

; takes a TPAL term T0 and a heartbeat parameter ♥ and returns
; as a pair of numbers the work and span cost of executing
; the term to completion
(define (cost-of-exec ♥ T0)
    (let ([g (car (judgment-holds (⇓ ,♥ () ,T0 J T g) g))])
      (cons (term (Work ,g)) (term (Span ,g)))))

(define (work-of work-span)
  (car work-span))

(define (span-of work-span)
  (cdr work-span))

(define (average-parallelism-of work-span)
  (/ (work-of work-span) (span-of work-span)))

(define (plot-work-for-varying-♥ benchmark-name lines-list)
  (parameterize ([plot-title benchmark-name]
                 [plot-x-label "Heartbeat parameter ♥"]
                 [plot-y-label "Work"])
    (plot lines-list #:x-min 10 #:y-min 0 #:legend-anchor 'top-right
          #:out-file (string-append benchmark-name "_work.jpg") #:out-kind 'jpeg)))

(define (plot-span-for-varying-♥ benchmark-name lines-list)
  (parameterize ([plot-title benchmark-name]
                 [plot-x-label "Heartbeat parameter ♥"]
                 [plot-y-label "Span"])
    (plot lines-list #:x-min 10 #:y-min 0 #:legend-anchor 'bottom-left
          #:out-file (string-append benchmark-name "_span.jpg") #:out-kind 'jpeg)))

(define (plot-average-parallelism-for-varying-♥ benchmark-name lines-list)
  (parameterize ([plot-title benchmark-name]
                 [plot-x-label "Heartbeat parameter ♥"]
                 [plot-y-label "Average parallelism (Work/Span)"])
    (plot lines-list #:x-min 10 #:y-min 1 #:legend-anchor 'top-right
          #:out-file (string-append benchmark-name "_avgpar.jpg") #:out-kind 'jpeg)))

(define (cost-of-prod ♥ a b)
  (cost-of-exec ♥ (prod-T-init a b)))

(define (mk-f-curve-of-prod f ♥-lo ♥-hi)
  (let ([b 5]
        [♥-range (map (λ (n) (* 10 n)) (range ♥-lo ♥-hi))])
    (let ([mk-line (λ (a style)
               (lines (map (λ (♥) (vector ♥ (f (cost-of-prod ♥ a b)))) ♥-range)
                      #:label (string-append (number->string a) " * " (number->string b))
                      #:style style))])
        (list (mk-line 20 'solid) (mk-line 30 'dot) (mk-line 40 'long-dash)))))

(define (plot-for-prod ♥-upper)
  (let* ([♥-lo 1]
         [♥-hi (max (add1 ♥-lo) ♥-upper)]
         [n "prod"])
    (begin
      (plot-work-for-varying-♥ n (mk-f-curve-of-prod work-of ♥-lo ♥-hi))
      (plot-span-for-varying-♥ n (mk-f-curve-of-prod span-of ♥-lo ♥-hi))
      (plot-average-parallelism-for-varying-♥ n (mk-f-curve-of-prod average-parallelism-of ♥-lo ♥-hi)))))

(define (cost-of-pow ♥ a b)
  (cost-of-exec ♥ (pow-T-init a b)))

(define (mk-f-curve-of-pow f ♥-lo ♥-hi)
  (let ([a 5]
        [♥-range (map (λ (n) (* 10 n)) (range ♥-lo ♥-hi))])
    (let ([mk-line (λ (b style)
               (lines (map (λ (♥) (vector ♥ (f (cost-of-pow ♥ a b)))) ♥-range)
                      #:label (string-append (number->string a) "^" (number->string b))
                      #:style style))])
        (list (mk-line 3 'solid) (mk-line 5 'dot) (mk-line 7 'long-dash)))))

(define (plot-for-pow ♥-upper)
  (let* ([♥-lo 1]
         [♥-hi (max (add1 ♥-lo) ♥-upper)]
         [n "pow"])
    (begin
      (plot-work-for-varying-♥ n (mk-f-curve-of-pow work-of ♥-lo ♥-hi))
      (plot-span-for-varying-♥ n (mk-f-curve-of-pow span-of ♥-lo ♥-hi))
      (plot-average-parallelism-for-varying-♥ n (mk-f-curve-of-pow average-parallelism-of ♥-lo ♥-hi)))))

(define (cost-of-fib ♥ n)
  (cost-of-exec ♥ (fib-T-init n)))

(define (mk-f-curve-of-fib f ♥-lo ♥-hi)
  (let ([♥-range (map (λ (n) (* 10 n)) (range ♥-lo ♥-hi))])
    (let ([mk-line (λ (n style)
               (lines (map (λ (♥) (vector ♥ (f (cost-of-fib ♥ n)))) ♥-range)
                      #:label (string-append "n=" (number->string n))
                      #:style style))])
        (list (mk-line 4 'solid) (mk-line 5 'dot) (mk-line 6 'long-dash)))))

(define (plot-for-fib ♥-upper)
  (let* ([♥-lo 1]
         [♥-hi (max (add1 ♥-lo) ♥-upper)]
         [n "fib"])
    (begin
      (plot-work-for-varying-♥ n (mk-f-curve-of-fib work-of ♥-lo ♥-hi))
      (plot-span-for-varying-♥ n (mk-f-curve-of-fib span-of ♥-lo ♥-hi))
      (plot-average-parallelism-for-varying-♥ n (mk-f-curve-of-fib average-parallelism-of ♥-lo ♥-hi)))))

(define (plot-all)
  (begin
    (plot-for-prod 16)
    (plot-for-pow 16)
    (plot-for-fib 16)))