Main compiler file from CSCI-P 423 Compilers course. We were asked not to make the code public (to prevent cheating I guess), but I think using a "secret gist" should be fine. Aside from some course-provided scaffolding, all the code in this file is my own. Based on the textbook "Essentials of Compilation" by Jeremy Siek.

csci-p423-compiler.rkt

#lang racket
(require racket/set racket/stream)
(require racket/fixnum)
(require "interp-Lfun-prime.rkt")
(require "interp-Cfun.rkt")
(require "interp.rkt")
(require "type-check-Lthread.rkt")
(require "type-check-Cthread.rkt")
(require "utilities.rkt")
(require graph)
(require data/queue)
(require "multigraph.rkt")
(require "graph-printing.rkt")
(require "priority_queue.rkt")
(require "heap.rkt")
(provide (all-defined-out))

; TODO: debug why the one test case segfaults

(define mutex-byte-size 64)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Lint examples
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; The following compiler pass is just a silly one that doesn't change
;; anything important, but is nevertheless an example of a pass. It
;; flips the arguments of +. -Jeremy
(define (flip-exp e)
  (match e
    [(Var x) e]
    [(Prim 'read '()) (Prim 'read '())]
    [(Prim '- (list e1)) (Prim '- (list (flip-exp e1)))]
    [(Prim '+ (list e1 e2)) (Prim '+ (list (flip-exp e2) (flip-exp e1)))]))

(define (flip-Lint e)
  (match e
    [(Program info e) (Program info (flip-exp e))]))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Passes
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;;;;; SHRINK

; this needs to come first so that short-circuiting and/or is handled properly

(define (shrink-exp expr)
  (match expr
    [(Var x) (Var (racket-id->c-id x))]
    [(Int n) (Int n)]
    [(Bool b) (Bool b)]
    [(Void) (Void)]
    [(Let x e body) (Let (racket-id->c-id x) (shrink-exp e) (shrink-exp body))]
    [(SetBang var rhs) (SetBang var (shrink-exp rhs))]
    [(If cmp thn els)
     (If (shrink-exp cmp) (shrink-exp thn) (shrink-exp els))]
    [(WhileLoop cnd body)
     (WhileLoop (shrink-exp cnd) (shrink-exp body))]
    [(Begin es body)
     (Begin (map shrink-exp es) (shrink-exp body))]
    [(Apply rator rands)
     (Apply (shrink-exp rator) (map shrink-exp rands))]
    [(Prim 'and (list e1 e2))
     (If (shrink-exp e1) (shrink-exp e2) (Bool #f))]
    [(Prim 'or (list e1 e2))
     (If (shrink-exp e1) (Bool #t) (shrink-exp e2))]
    [(Prim p es)
     (Prim p (map shrink-exp es))]
    [(HasType p t)
     (shrink-exp p)]))

(define (shrink-def d)
  (match-let
    ([(Def name arg-ts ret-t dinfo de) d])
    (Def (racket-id->c-id name) arg-ts ret-t dinfo (shrink-exp de))))

(define (shrink p)
  (match p
    [(Program info e)
     (ProgramDefs info (list (Def 'main '() 'Integer '() (shrink-exp e))))]
    [(ProgramDefsExp info defs e)
     (let*
       ([main (Def 'main '() 'Integer '() e)]
        [edefs (cons main defs)]
        [shrunken-defs
          (map shrink-def edefs)])
       (ProgramDefs info shrunken-defs))]))

;;;;; PARTIAL EVALUATION

; textbook doesn't say to, so I won't, but I think it's also possible to substitute in vars with residual vals

(define (pe-neg r)
  (match r
    [(Var x) (Prim '- (list (Var x)))]
    [(Prim 'read '()) (Prim '- (list (Prim 'read '())))]
    [(Prim '- (list (Var x))) (Var x)]
    [(Prim '- (list (Prim 'read '()))) (Prim 'read '())]
    [(Prim '+ (list (Int x) y))
     (Prim '+ (list (Int (fx- 0 x)) (Prim '- (list y))))]
    [(Prim '+ (list x y)) (Prim '- (list (Prim '+ (list x y))))]
    [(Let x e body) (Prim '- (list (Let x e (pe-exp body))))]
    [(Int n) (Int (fx- 0 n))]
    [(If cnd thn els)
     (If (pe-exp cnd) (pe-neg thn) (pe-neg els))]
    [else (Prim '- (list r))]))

(define (pe-add e1 e2)
  (match* (e1 e2)
    [((Int m) (Int n)) (Int (fx+ m n))]
    [((Prim '+ (list (Int n1) i1)) (Prim '+ (list (Int n2) i2)))
     (Prim
       '+
       (list
         (Int (fx+ n1 n2))
         (Prim '+ (list i1 i2))))]
    [((Prim '+ (list (Int m) i)) (Int n))
     (Prim '+ (list (Int (fx+ m n)) i))]
    [((Int m) (Prim '+ (list (Int n) i)))
     (Prim '+ (list (Int (fx+ m n)) i))]
    [(a b) (Prim '+ (list a b))]))

(define (pe-exp e)
  (match e
    [(Var x) (Var x)]
    [(Int n) (Int n)]
    [(Bool b) (Bool b)]
    [(Void) (Void)]
    [(Let x e body)
     (let
       ([pe-e (pe-exp e)]
        [pe-body (pe-exp body)])
       (Let x pe-e pe-body))]
    [(SetBang var rhs) (SetBang var (pe-exp rhs))]
    [(If test t f) (If (pe-exp test) (pe-exp t) (pe-exp f))]
    [(WhileLoop cnd body) (WhileLoop (pe-exp cnd) (pe-exp body))]
    [(Begin es body)
     (Begin (map pe-exp es) (pe-exp body))]
    [(Apply rator rands)
     (Apply (pe-exp rator) (map pe-exp rands))]
    [(Prim '+ (list e1 e2))
     (pe-add
       (pe-exp e1)
       (pe-exp e2))]
    [(Prim '- (list e))
     (pe-neg (pe-exp e))]
    [(Prim '- (list e1 e2))
     (pe-exp (Prim '+ (list e1 (Prim '- (list e2)))))]
    [(Prim 'read '()) (Prim 'read '())]
    [(Prim op es) (Prim op (map pe-exp es))]))

(define (pe-def d)
  (match-let ([(Def name arg-ts ret-t dinfo de) d])
    (Def name arg-ts ret-t dinfo (pe-exp de))))

(define (pe-Lint p)
  (match p
    [(ProgramDefs info defs)
     (ProgramDefs
       info
       (map pe-def defs))]))

;;;;; UNIQUIFY

(define lookup-uniquify-env
  (lambda (x env)
    (cdr (assv x env))))

(define extend-uniquify-env
  (lambda (x v env)
    (cons (cons x v) env)))

(define (uniquify-exp env)
  (lambda (e)
    (match e
      [(Var x)
       (Var (lookup-uniquify-env x env))]
      [(Int n) (Int n)]
      [(Bool b) (Bool b)]
      [(Void) (Void)]
      [(Let x e body)
       (let*
         ([new-e ((uniquify-exp env) e)]
          [new-x (gensym x)]
          [new-env (extend-uniquify-env x new-x env)]
          [new-body ((uniquify-exp new-env) body)])
         (Let new-x new-e new-body))]
      [(SetBang var rhs)
       (SetBang (lookup-uniquify-env var env) ((uniquify-exp env) rhs))]
      [(If test t f)
       (If
         ((uniquify-exp env) test)
         ((uniquify-exp env) t)
         ((uniquify-exp env) f))]
      [(WhileLoop cnd body)
       (WhileLoop
         ((uniquify-exp env) cnd)
         ((uniquify-exp env) body))]
      [(Begin es body)
       (Begin
         (map (uniquify-exp env) es)
         ((uniquify-exp env) body))]
      [(Apply rator rands)
       (Apply ((uniquify-exp env) rator) (map (uniquify-exp env) rands))]
      [(Prim op es)
       (Prim op (for/list ([e es]) ((uniquify-exp env) e)))])))

(define (get-function-names-env defs)
  (map
    (lambda (d)
      (match-let
        ([(Def name _ _ _ _) d])
        (cons name name))) ; do function names need to be mangled?
    defs))

(define (uniquify-defs env)
  (lambda (d)
    (match-let*
      ([(Def name arg-ts ret-t dinfo de) d]
       [arg-names
         (map
           (lambda (x)
             (cons (car x) (car x))) ; uniquify arguments
           arg-ts)])
      (Def name arg-ts ret-t dinfo ((uniquify-exp (append arg-names env)) de)))))

;; uniquify : Lvar -> Lvar
(define (uniquify p)
  (match p
    [(ProgramDefs info defs)
     (ProgramDefs
       info
       (map (uniquify-defs (get-function-names-env defs)) defs))]))

;;;;; UNCOVER GET

(define (collect-set! e)
  (match e
    [(Var x) (set)]
    [(Int n) (set)]
    [(Bool b) (set)]
    [(Void) (set)]
    [(Let x e body)
     (set-union (collect-set! e) (collect-set! body))]
    [(SetBang var rhs)
     (set-union (set var) (collect-set! rhs))]
    [(If test thn els)
     (set-union (collect-set! test) (collect-set! thn) (collect-set! els))]
    [(WhileLoop cnd body)
     (set-union (collect-set! cnd) (collect-set! body))]
    [(Begin es body)
     (set-union (apply set-union (cons (set) (map collect-set! es))) (collect-set! body))]
    [(Apply rator rands)
     (set-union (collect-set! rator) (apply set-union (cons (set) (map collect-set! rands))))]
    [(Prim op es)
     (if (null? es)
       (set)
       (apply set-union (cons (set) (map collect-set! es))))]))

(define ((uncover-get!-exp set!-vars) e)
  (match e
    [(Var x)
     (if (set-member? set!-vars x)
       (GetBang x)
       (Var x))]
    [(Int n) (Int n)]
    [(Bool b) (Bool b)]
    [(Void) (Void)]
    [(Let x e body)
     (Let x ((uncover-get!-exp set!-vars) e) ((uncover-get!-exp set!-vars) body))]
    [(SetBang var rhs) (SetBang var ((uncover-get!-exp set!-vars) rhs))]
    [(If test thn els)
     (If
       ((uncover-get!-exp set!-vars) test)
       ((uncover-get!-exp set!-vars) thn)
       ((uncover-get!-exp set!-vars) els))]
    [(WhileLoop cnd body)
     (WhileLoop
       ((uncover-get!-exp set!-vars) cnd)
       ((uncover-get!-exp set!-vars) body))]
    [(Begin es body)
     (Begin
       (map (uncover-get!-exp set!-vars) es)
       ((uncover-get!-exp set!-vars) body))]
    [(Apply rator rands)
     (Apply
       ((uncover-get!-exp set!-vars) rator)
       (map (uncover-get!-exp set!-vars) rands))]
    [(Prim op es)
     (Prim op (map (uncover-get!-exp set!-vars) es))]))

(define (uncover-get!-def d)
  (match-let*
    ([(Def name arg-ts ret-t dinfo de) d]
     [set!-vars (collect-set! de)])
    (Def name arg-ts ret-t dinfo ((uncover-get!-exp set!-vars) de))))

(define (uncover-get! p)
  (match p
    [(ProgramDefs info ds)
     (ProgramDefs info (map uncover-get!-def ds))]))

;;;;; REVEAL FUNCTIONS

(define (reveal-functions-var x env)
  (let ([i (assv x env)])
    (if i
      (FunRef x (cdr i))
      (Var x))))

(define (reveal-functions-exp env)
  (lambda (e)
    (match e
      [(Var x) (reveal-functions-var x env)]
      [(GetBang x) (GetBang x)]
      [(Int n) (Int n)]
      [(Bool b) (Bool b)]
      [(Void) (Void)]
      [(Let x e body)
       (Let
         x
         ((reveal-functions-exp env) e)
         ((reveal-functions-exp env) body))]
      [(SetBang var rhs)
       (SetBang var ((reveal-functions-exp env) rhs))]
      [(If cnd thn els)
       (If
         ((reveal-functions-exp env) cnd)
         ((reveal-functions-exp env) thn)
         ((reveal-functions-exp env) els))]
      [(WhileLoop cnd body)
       (WhileLoop
         ((reveal-functions-exp env) cnd)
         ((reveal-functions-exp env) body))]
      [(Begin es body)
       (Begin
         (map (reveal-functions-exp env) es)
         ((reveal-functions-exp env) body))]
      [(Apply rator rands)
       (Apply
         ((reveal-functions-exp env) rator)
         (map (reveal-functions-exp env) rands))]
      [(Prim op es) (Prim op (map (reveal-functions-exp env) es))])))

(define (reveal-functions-def env)
  (lambda (d)
    (match-let
      ([(Def name arg-ts ret-t info de) d])
      (Def name arg-ts ret-t info ((reveal-functions-exp env) de)))))

(define (gen-functions-env defs)
  (map
    (lambda (d)
      (match-let
        ([(Def name arg-ts ret-t info de) d])
        (cons name (length arg-ts))))
    defs))

(define (reveal-functions p)
  (match p
    [(ProgramDefs info defs)
     (ProgramDefs
       info
       (map
         (reveal-functions-def (gen-functions-env defs))
         defs))]))

;;;;; LIMIT FUNCTIONS

;; generate-lets : (var . exp) list -> exp -> exp
;  used for expose allocation, rco, limit functions
(define (generate-lets env)
  (lambda (e)
    (match env
      ['() e]
      [`((,tvar . ,texp) . ,cdr-env)
        (let ([next-let (Let tvar texp e)])
          ((generate-lets cdr-env) next-let))])))

(define (limit-functions-exp e)
  (match e
    [(Var x) (Var x)]
    [(GetBang x) (GetBang x)]
    [(FunRef x ar)
     (FunRef x (min ar 6))]
    [(Int n) (Int n)]
    [(Bool b) (Bool b)]
    [(Void) (Void)]
    [(Let x e body)
     (Let
       x
       (limit-functions-exp e)
       (limit-functions-exp body))]
    [(SetBang var rhs)
     (SetBang var (limit-functions-exp rhs))]
    [(If cnd thn els)
     (If
       (limit-functions-exp cnd)
       (limit-functions-exp thn)
       (limit-functions-exp els))]
    [(WhileLoop cnd body)
     (WhileLoop
       (limit-functions-exp cnd)
       (limit-functions-exp body))]
    [(Begin es body)
     (Begin
       (map limit-functions-exp es)
       (limit-functions-exp  body))]
    [(Apply rator rands)
     (if (> (length rands) 6)
       (let
         ([plain-args (take rands 5)]
          [xtra-args (drop rands 5)])
         (Apply
           rator
           (append
             plain-args
             (list (Prim 'vector xtra-args)))))
       (Apply rator rands))]
    [(Prim op es) (Prim op (map limit-functions-exp es))]))

(define (limit-functions-def d)
  (match-let
    ([(Def name arg-ts ret-t dinfo de) d])
    (if (> (length arg-ts) 6)
      (let* 
        ([vec-arg (gensym 'temp)]
         [plain-args (take arg-ts 5)]
         [xtra-args (drop arg-ts 5)]
         [6th-arg
          `(,vec-arg : (Vector . ,(map caddr xtra-args)))]
         [new-arg-ts (append plain-args (list 6th-arg))]
         [xtra-args-env
           (map
             (lambda (arg i)
               (cons
                 (car arg)
                 (Prim 'vector-ref (list (Var vec-arg) (Int i)))))
             xtra-args
             (range (length xtra-args)))])
        (Def
          name
          new-arg-ts
          ret-t
          dinfo
          ((generate-lets xtra-args-env) (limit-functions-exp de))))
      (Def name arg-ts ret-t dinfo (limit-functions-exp de)))))

(define (limit-functions p)
  (match p
    [(ProgramDefs info defs)
     (ProgramDefs
       info
       (map limit-functions-def defs))]))

;;;;; EXPOSE ALLOCATION

(define (set-vec-mems vec-name env)
  (map
    (lambda (m idx)
      (Prim 'vector-set! (list (Var vec-name) (Int idx) (Var (car m)))))
    env
    (range (length env))))

(define (vectoroffy t)
  (match t
    [(cons 'Vector ts)
     (cons
       'Vector
       (cons
         '(Vectorof Integer)
         (map vectoroffy ts)))]
    [other other]))

(define (generate-vec-allocation v t) ; vector and type
  (match v
    [(Prim 'vector mems)
     (let*
       ([expose-alloc-mems (map expose-allocation-exp mems)]
        [vec-mem-env
          (map
            (lambda (m) (cons (gensym 'x) m))
            expose-alloc-mems)]
        [bytes-needed (+ 8 (+ (* 8 (add1 (length mems))) mutex-byte-size))]
        [vec-name (gensym 'v)]
        [mutex-name (gensym 'mutex)])
       ((generate-lets (reverse vec-mem-env))
        (Begin
          (list
            (Prim 'mutex-lock (list (GlobalAddressValue 'allocation_mutex)))
            (If
              (Prim
                '<
                (list
                  (Prim '+ (list (GlobalValue 'free_ptr) (Int bytes-needed)))
                  (GlobalValue 'fromspace_end)))
              (Void)
              (Collect bytes-needed)))
          (Let
            mutex-name
            (AllocateArray (Int mutex-byte-size) '(Vectorof Integer))
            (Begin
              (list (Prim 'mutex-init (list (Var mutex-name))))
              (Let
                vec-name
                (Allocate (add1 (length mems)) (vectoroffy t))
                (Begin
                  (set-vec-mems vec-name (append (list (cons mutex-name mutex-name)) vec-mem-env))
                  (Begin
                    (list
                      (Prim 'mutex-unlock (list (GlobalAddressValue 'allocation_mutex)))
                      )
                    (Var vec-name)))))))))]
     [else (error "generate-vec-allocation only works on vectors")]))

(define (expose-allocation-exp e)
  (match e
    [(Var x) (Var x)]
    [(GetBang x) (GetBang x)]
    [(FunRef f ar) (FunRef f ar)]
    [(Int n) (Int n)]
    [(Bool b) (Bool b)]
    [(Void) (Void)]
    [(Let x e body)
     (Let x (expose-allocation-exp e) (expose-allocation-exp body))]
    [(SetBang var rhs) (SetBang var (expose-allocation-exp rhs))]
    [(If cnd thn els)
     (If
       (expose-allocation-exp cnd)
       (expose-allocation-exp thn)
       (expose-allocation-exp els))]
    [(WhileLoop cnd body)
     (WhileLoop
       (expose-allocation-exp cnd)
       (expose-allocation-exp body))]
    [(Begin es body)
     (Begin
       (map expose-allocation-exp es)
       (expose-allocation-exp body))]
    [(Apply rator rands)
     (Apply
       (expose-allocation-exp rator)
       (map expose-allocation-exp rands))]
    [(Prim 'vector-ref (list v (Int n)))
     (let ([tempvec (gensym 'tempvec)]
           [temp (gensym 'temp)])
       (Let
         tempvec
         (expose-allocation-exp v)
         (Begin
           (list
             (Prim 'mutex-lock (list (Prim 'vector-ref (list (Var tempvec) (Int 0))))))
           (Let
             temp
             (Prim
               'vector-ref
               (list
                 (Var tempvec)
                 (Int (add1 n))))
             (Begin
               (list
                 (Prim 'mutex-unlock (list (Prim 'vector-ref (list (Var tempvec) (Int 0))))))
               (Var temp))))))]
    [(Prim 'vector-set! (list v (Int n) e))
     (let ([tempvec (gensym 'tempvec)]
           [temp (gensym 'temp)])
       (Let
         temp
         (expose-allocation-exp e)
         (Let
           tempvec
           (expose-allocation-exp v)
           (Begin
             (list
               (Prim 'mutex-lock (list (Prim 'vector-ref (list (Var tempvec) (Int 0)))))
               (Prim
                 'vector-set!
                 (list
                   (Var tempvec)
                   (Int (add1 n))
                   (Var temp))))
             (Prim 'mutex-unlock (list (Prim 'vector-ref (list (Var tempvec) (Int 0)))))))))]
    [(Prim 'vector-length (list v))
     (Prim
       '-
       (list
         (Prim '(vector-length (list (expose-allocation-exp v))))
         (Int 1)))]
    [(Prim op es) (Prim op (map expose-allocation-exp es))]
    [(HasType e t)
     (generate-vec-allocation e t)]))

(define (vectoroffy-args args)
  (match args
    [(cons (list x ': t) d)
     (cons
       (list x ': (vectoroffy t))
       (vectoroffy-args d)) ]
    ['() '()]))

(define (expose-allocation-def d)
  (match-let
    ([(Def name arg-ts ret-t info de) d])
    (Def name (vectoroffy-args arg-ts) (vectoroffy ret-t) info (expose-allocation-exp de))))

(define (expose-allocation p)
  (let ([typed-vec-p (type-check-Lthread-has-type p)])
    (match typed-vec-p
      [(ProgramDefs info ds)
       (ProgramDefs info (map expose-allocation-def ds))])))

;;;;; REMOVE COMPLEX OPERANDS

;; rco-exp : exp -> exp
(define (rco-exp p)
  (match p
    [(Var x) (Var x)]
    [(GetBang x) (Var x)]
    [(FunRef f ar) (FunRef f ar)]
    [(Int n) (Int n)]
    [(Bool b) (Bool b)]
    [(Void) (Void)]
    [(Let x e body)
     (Let x (rco-exp e) (rco-exp body))]
    [(SetBang var rhs) (SetBang var (rco-exp rhs))]
    [(If test e-true e-false)
     (If (rco-exp test) (rco-exp e-true) (rco-exp e-false))]
    [(WhileLoop cnd body)
     (WhileLoop (rco-exp cnd) (rco-exp body))]
    [(Begin es body)
     (Begin (map rco-exp es) (rco-exp body))]
    [(Apply rator rands)
     (match-let*
       ([`(,rands-atoms . ,rands-env) (rco-atoms rands)]
        [`(,rator-atom . ,rator-env) (rco-atom rator)])
       ((generate-lets (append rator-env rands-env))
        (Apply rator-atom rands-atoms)))]
    [(Collect n) (Collect n)]
    [(Allocate n t) (Allocate n t)]
    [(AllocateArray n t) (AllocateArray n t)] ; n guaranteed to be (Int ...) b/c not actually adding array support
    [(GlobalValue x) (GlobalValue x)]
    [(Prim op es)
     (match-let ([`(,new-atoms . ,new-env) (rco-atoms es)])
       ((generate-lets new-env) (Prim op new-atoms)))]))

;; rco-atom : Lvar -> atom, env
(define (rco-atom p)
  (match p
    [(Var x) (cons (Var x) '())]
    [(GetBang x) (gen-temp-var-env (Var x))]
    [(FunRef f ar) (gen-temp-var-env (FunRef f ar))]
    [(Int n) (cons (Int n) '())]
    [(Bool b) (cons (Bool b) '())]
    [(Void) (cons (Void) '())]
    [(Let x e body)
     (gen-temp-var-env
       (Let x (rco-exp e) (rco-exp body)))]
    [(SetBang var rhs)
     (gen-temp-var-env
       (SetBang var (rco-exp rhs)))]
    [(If test e-true e-false)
     (gen-temp-var-env
       (If (rco-exp test) (rco-exp e-true) (rco-exp e-false)))]
    [(WhileLoop cnd body)
     (gen-temp-var-env (WhileLoop (rco-exp cnd) (rco-exp body)))]
    [(Begin es body)
     (gen-temp-var-env
       (Begin (map rco-exp es) (rco-exp body)))]
    [(Apply rator rands)
     (gen-temp-var-env (rco-exp (Apply rator rands)))]
    [(Prim op es)
     (match-let ([`(,new-atoms . ,new-env) (rco-atoms es)])
       (let ([temp-var (gensym 'temp)])
         (cons (Var temp-var) (cons (cons temp-var (Prim op new-atoms)) new-env))))]
    [(Collect n) (gen-temp-var-env (Collect n))]
    [(Allocate n t) (gen-temp-var-env (Allocate n t))]
    [(GlobalValue x) (gen-temp-var-env (GlobalValue x))]
    [(GlobalAddressValue x) (gen-temp-var-env (GlobalAddressValue x))]
    ))

(define (gen-temp-var-env e)
  (let ([temp-var (gensym 'temp)])
    (cons
      (Var temp-var)
      (list (cons temp-var e)))))

;; prim-atoms-env : exp list -> (atom list, ((var . exp) list))
(define (rco-atoms ps)
  (let*
    ([new-atoms-envs (map rco-atom ps)]
     [new-atoms
       (map car new-atoms-envs)]
     [new-envs
       (map cdr new-atoms-envs)]
     ; reversing so that (reads) are in right order when lets generated
     [new-env (append* (reverse new-envs))])
    (cons new-atoms new-env)))

(define (rco-def d)
  (match-let
    ([(Def name arg-ts ret-t dinfo de) d])
    (Def name arg-ts ret-t dinfo (rco-exp de))))

;; remove-complex-opera* : Lvar -> Lvar^mon
(define (remove-complex-opera* p)
  (match p
    [(ProgramDefs info ds)
     (ProgramDefs
       info
       (map rco-def ds))]))

;;;;; EXPLICATE CONTROL

(define (create-block tail bs) ; -> tail, block dict
  (match tail
    [(Goto label) (cons (Goto label) bs)]
    [else
      (let ([label (gensym 'block)])
        (cons (Goto label) (cons (cons label tail) bs)))]))

(define (create-while-loop cnd body cont bs)
  (match-let*
    ([loop-label (gensym 'loop)]
     [(cons e-body bs1) (explicate-effect body (Goto loop-label) bs)]
     [(cons e-cnd bs2) (explicate-pred cnd e-body cont bs1)])
    (cons
      (Goto loop-label)
      (cons (cons loop-label e-cnd) bs2))))

(define (explicate-tail p bs) ; not confident about the returns of anything other than atoms
  (match p
    [(Var x) (cons (Return (Var x)) bs)]
    [(FunRef f a) (cons (Return (FunRef f a)))]
    [(Int n) (cons (Return (Int n)) bs)]
    [(Bool b) (cons (Return (Bool b)) bs)]
    [(Void) (cons (Return (Void)) bs)]
    [(Let x e body)
     (match-let ([(cons e-body bs1) (explicate-tail body bs)])
       (explicate-assign e x e-body bs1))]
    [(SetBang x rhs)
     (explicate-assign rhs x (Return (Void)) bs)]
    [(If cnd thn els)
     (match-let*
       ([(cons e-thn bs1) (explicate-tail thn bs)]
        [(cons e-els bs2) (explicate-tail els bs1)])
       (explicate-pred cnd e-thn e-els bs2))]
    [(WhileLoop cnd body) (create-while-loop cnd body (Void) bs)]
    [(Begin es body)
     (match-let*
       ([(cons e-body bs1) (explicate-tail body bs)])
       (explicate-effects es e-body bs1))]
    [(Prim op es) (cons (Return (Prim op es)) bs)] ; seems to work...
    [(Apply rator rands) (cons (TailCall rator rands) bs)]
    [(Collect n) (cons (Seq (Collect n) (Return (Void))) bs)]
    [(Allocate n t) (cons (Return (Allocate n t)) bs)]
    [(GlobalValue x) (cons (Return (GlobalValue x)) bs)]
    ))

(define (explicate-assign e x cont bs)
  (match e
    [(Var nx) (cons (Seq (Assign (Var x) (Var nx)) cont) bs)]
    [(FunRef f a) (cons (Seq (Assign (Var x) (FunRef f a)) cont) bs)]
    [(Int n) (cons (Seq (Assign (Var x) (Int n)) cont) bs)]
    [(Bool b) (cons (Seq (Assign (Var x) (Bool b)) cont) bs)]
    [(Void) (cons (Seq (Assign (Var x) (Void)) cont) bs)]
    [(Let y ye body)
     (match-let ([(cons ea-body bs1) (explicate-assign body x cont bs)])
       (explicate-assign ye y ea-body bs1))]
    [(SetBang x rhs)
     (match-let*
       ([(cons e-rhs bs1) (explicate-tail rhs bs)])
       (explicate-assign e-rhs x cont bs1))]
    [(If cnd thn els)
     (match-let*
       ([(cons e-thn bs1) (explicate-assign thn x cont bs)]
        [(cons e-els bs2) (explicate-assign els x cont bs1)])
     (explicate-pred cnd e-thn e-els bs2))]
    [(WhileLoop cnd body)
     (create-while-loop cnd body cont bs)]
    [(Begin es body)
     (match-let
       ([(cons e-body bs1) (explicate-assign body x cont bs)])
       (explicate-effects es e-body bs1))]
    [(Prim op es) (cons (Seq (Assign (Var x) (Prim op es)) cont) bs)]
    [(Apply rator rands)
     (cons (Seq (Assign (Var x) (Call rator rands)) cont) bs)]
    [(Collect n) (cons (Seq (Collect n) cont) bs)]
    [(Allocate n t)
     (cons
       (Seq
         (Assign (Var x) (Allocate n t))
         cont)
       bs)]
    [(AllocateArray n t)
     (cons
       (Seq
         (Assign (Var x) (AllocateArray n t))
         cont)
       bs)]
    [(GlobalValue y)
     (cons
       (Seq
         (Assign (Var x) (GlobalValue y))
         cont)
       bs)]
    [(GlobalAddressValue y)
     (cons
       (Seq
         (Assign (Var x) (GlobalValue y))
         cont)
       bs)]))

(define cmp-ops '(< <= eq? >= >))

(define (explicate-pred cnd thn els bs) ; assume thn and els are tails
  (match-let*
    ([(cons thn-jmp bs1) (create-block thn bs)]
     [(cons els-jmp bs2) (create-block els bs1)])
    (match cnd
      [(Var x)
       (cons
         (IfStmt
           (Prim 'eq? (list (Var x) (Bool #t)))
           thn-jmp
           els-jmp)
         bs2)]
      [(FunRef f a) (error "should not have funref in explicate-pred")]
      [(Bool b)
       (if b
         (cons thn bs)
         (cons els bs))]
      [(Void) (error "should not have void in explicate-pred")]
      [(Let x rhs body)
       (match-let
         ([(cons cont bs3) (explicate-pred body thn-jmp els-jmp bs2)])
         (explicate-assign rhs x cont bs3))]
      [(SetBang var rhs) (error "should not have setbang in explicate-pred")]
      [(If cnd^ thn^ els^)
       (match-let*
         ([(cons thn^-tail bs3) (explicate-pred thn^ thn-jmp els-jmp bs2)]
          [(cons els^-tail bs4) (explicate-pred els^ thn-jmp els-jmp bs3)]
          [(cons thn^-jmp bs5) (create-block thn^-tail bs4)]
          [(cons els^-jmp bs6) (create-block els^-tail bs5)])
         (explicate-pred cnd^ thn^-jmp els^-jmp bs6))]
      [(WhileLoop cnd body) (error "should not have while in explicate-pred")]
      [(Begin es body)
       (match-let
         ([(cons e-body bs1) (explicate-pred body thn-jmp els-jmp bs)])
         (explicate-effects es e-body bs1))]
      [(Prim 'not (list e))
       (explicate-pred e els thn bs)]
      [(Prim 'vector-ref (list v idx))
       (match-let*
         ([t (gensym 'temp)]
          [(cons ep bs2) (explicate-pred (Var t) thn els bs)])
         (explicate-assign 
           (Prim 'vector-ref (list v idx))
           t
           ep
           bs2))]
      [(Prim op es)
       #:when (memv op cmp-ops)
       (cons (IfStmt (Prim op es) thn-jmp els-jmp) bs2)]
      [(Apply rator rands)
       (match-let*
         ([temp (gensym 'temp)]
          [(cons ep b1) (explicate-pred (Var temp) thn els bs)])
         (explicate-assign
           (Apply rator rands)
           temp
           ep
           bs1))]
        
      [(Collect n) (error "should not have collect in explicate-pred")]
      [(Allocate n t) (error "should not have allocate in explicate-pred")]
      [(GlobalValue x)
       (cons
         (IfStmt
           (Prim 'eq? (list (GlobalValue x) (Bool #t)))
           thn-jmp
           els-jmp)
         bs2)])))

(define (explicate-effect e cont bs) ; -> tail, cont, blocks -> tail, blocks
  (match e
    [(Let x rhs body)
     (match-let*
       ([(cons e-body bs1) (explicate-effect body cont bs)])
       (explicate-assign rhs x e-body bs1))]
    [(SetBang var rhs) (explicate-assign rhs var cont bs)]
    [(If cnd thn els)
     (match-let*
       ([(cons thn-tail bs1) (explicate-effect thn cont bs)]
        [(cons els-tail bs2) (explicate-effect els cont bs1)])
       (explicate-pred cnd thn-tail els-tail bs2))]
    [(WhileLoop cnd body)
     (create-while-loop cnd body cont bs)]
    [(Begin es body)
     (match-let*
       ([(cons e-cont bs1) (explicate-effect body cont bs)])
       (explicate-effects es e-cont bs1))]
    [(Prim 'read (list)) (cons (Seq (Prim 'read (list)) cont) bs)]
    [(Prim 'vector-set! (list vec len val))
     (cons (Seq (Prim 'vector-set! (list vec len val)) cont) bs)]
    [(Prim 'print (list n))
     (cons (Seq (Prim 'print (list n)) cont) bs)]
    [(Prim 'sleep (list n))
     (cons (Seq (Prim 'sleep (list n)) cont) bs)]
    [(Prim 'call-parallel (list f1 v1 f2 v2))
     (cons (Seq (Prim 'call-parallel (list f1 v1 f2 v2)) cont) bs)]
    [(Prim 'mutex-lock (list m))
     (cons (Seq (Prim 'mutex-lock (list m)) cont) bs)]
    [(Prim 'mutex-unlock (list m))
     (cons (Seq (Prim 'mutex-unlock (list m)) cont) bs)]
    [(Prim 'mutex-init (list m))
     (cons (Seq (Prim 'mutex-init (list m)) cont) bs)]
    [(Apply rator rands)
     (cons (Seq (Call rator rands) cont) bs)]
    [(Collect n) (cons (Seq (Collect n) cont) bs)]
    [e (cons cont bs)])) ; catch-all for non-effectful stuff

(define (explicate-effects es cont bs)
  (cond
    [(null? es) (cons cont bs)]
    [else
      (match-let*
        ([(cons e-cdr bs2) (explicate-effects (cdr es) cont bs)])
        (explicate-effect (car es) e-cdr bs2))]))

(define (list-variables p) ; Lmonvar -> symbol list
  (match p
    [(Let x v body)
     (cons x (append (list-variables v) (list-variables body)))]
    [(If cmp thn els)
     (append
       (list-variables cmp)
       (list-variables thn)
       (list-variables els))]
    [(SetBang x rhs) (list-variables rhs)]
    [(Begin exps body)
     (append* (list-variables body) (map list-variables exps))]
    [(WhileLoop cnd body)
     (append (list-variables cnd) (list-variables body))]
    [(Prim p args)
     (append* (map list-variables args))]
    [(Apply rator rands)
     (append* (list-variables rator) (map list-variables rands))]
    [_ '()]))

#| (define (symbol-append . ls) |#
#|   (string->symbol |#
#|     (apply |#
#|       string-append |#
#|       (map symbol->string ls)))) |#

(define (explicate-control-def d)
  (match-let
    ([(Def name arg-ts ret-t dinfo de) d])
    (Def
      name
      arg-ts
      ret-t
      (cons (cons 'locals (list-variables de)) dinfo)
      (match-let
        ([(cons tail blocks) (explicate-tail de '())])
        (cons
          (cons
            (symbol-append name '_start)
            tail)
          blocks)))))

;; explicate-control : Lvar^mon -> Cvar
(define (explicate-control p)
  (match p
    [(ProgramDefs info ds)
     (ProgramDefs
       info
       (map explicate-control-def ds))]))

;;;;; SELECT INSTRUCTIONS

(define arg-pass-regs
  (list 'rdi 'rsi 'rdx 'rcx 'r8 'r9))

(define (arg-instr a)
  (match a
    [(Var x) (Var x)]
    [(Int n) (Imm n)]
    [(Bool b) (Imm (if b 1 0))]
    [(Void) (Imm 0)]
    [(GlobalValue x) (Global x)]
    [(GlobalAddressValue x) (Global x)]))

(define (cmp->cc cmp)
  (match cmp
    ['< 'l]
    ['<= 'le]
    ['eq? 'e]
    ['>= 'ge]
    ['> 'g]))

; list binary representation has left lsb, but handles conversion to int properly
(define (int->binary n)
  (cond
    [(zero? n) '(0)]
    [(zero? (sub1 n)) '(1)]
    [else
     (cons
       (if (zero? (modulo n 2)) 0 1)
       (int->binary (arithmetic-shift n -1)))]))

(define (binary->int b)
  (cond
    [(null? b) 0]
    [else (+ (car b) (* 2 (binary->int (cdr b))))]))

(define (pad-binary-to-length b n)
  (append
    b
    (build-list
      (- n (length b))
      (lambda (_) 0))))

(define (gen-pointer-mask ts)
  (match ts
    [`() '()]
    [`((Vector ,vts ...) . ,d)
      (cons 1 (gen-pointer-mask d))]
    [`(,a . ,d)
      (cons 0 (gen-pointer-mask d))]))

(define (generate-vec-tag t) ; type -> int (binary string)
  (let
    ([forwarding '(1)]
     [vec-len
       (pad-binary-to-length (int->binary (length t)) 6)]
     [pointer-mask
       (pad-binary-to-length (gen-pointer-mask (cdr t)) 50)])
    (pad-binary-to-length
      (append
        forwarding
        vec-len
        pointer-mask)
      64)))

(define (generate-mutex-array-tag t) ; type -> int (binary string)
  (let
    ([forwarding '(1)]
     [pointer-mask '(1)]
     [array-len (pad-binary-to-length (int->binary mutex-byte-size) 60)]
     [tag-type '(1)])
    (pad-binary-to-length
      (append
        forwarding
        pointer-mask
        array-len
        tag-type)
      64)))

(define (select-assign-instructions x p)
  (match p
    [(Int n)
     (list (Instr 'movq (list (Imm n) x)))]
    [(Var y)
     (list (Instr 'movq (list (Var y) x)))]
    [(Bool b)
     (list (Instr 'movq (list (arg-instr (Bool b)) x)))]
    [(Void)
     (list (Instr 'movq (list (Imm 0) x)))]
    [(FunRef f n)
     (list (Instr 'leaq (list (Global f) x)))]
    [(Call rator rands) ; also add call to select-stmt-instructions?
     (append
       (map
         (lambda (arg reg)
           (Instr 'movq (list (arg-instr arg) (Reg reg))))
         rands
         (take arg-pass-regs (length rands)))
       (list
         (IndirectCallq rator (length rands))
         (Instr 'movq (list (Reg 'rax) x))))]
    [(Prim 'read '())
     (list
       (Callq 'read_int 0)
       (Instr 'movq (list (Reg 'rax) x)))]
    [(Prim '+ (list e1 e2))
     (cond
       [(eqv? x e1)
        (list (Instr 'addq (list (arg-instr e2) x)))]
       [(eqv? x e2)
        (list (Instr 'addq (list (arg-instr e1) x)))]
       [else
         (list
           (Instr 'movq (list (arg-instr e1) x))
           (Instr 'addq (list (arg-instr e2) x)))])]
    [(Prim '- (list e))
     (cond
       [(eqv? x e)
        (list (Instr 'negq (list x)))]
       [else
         (list
           (Instr 'movq (list (arg-instr e) x))
           (Instr 'negq (list x)))])]
    [(Prim '- (list e1 e2))
     (cond
       [(eqv? x e1) ; x = x - e2
        (list (Instr 'subq (list (arg-instr e2) x)))]
       [(eqv? x e2) ; x = e1 - x
        (list
          (Instr 'negq (list x))
          (Instr 'addq (list (arg-instr e1) x)))]
       [else ; x = e1 - e2
         (list
           (Instr 'movq (list (arg-instr e1) x))
           (Instr 'subq (list (arg-instr e2) x)))])]
    [(Prim 'not (list e))
     (cond
       [(eqv? x e) ; x = not x
        (list (Instr 'xorq (list (Imm 1) x)))]
       [else
         (list
           (Instr 'movq (list (arg-instr e) x))
           (Instr 'xorq (list (Imm 1) x)))])]
    [(Prim 'vector-ref (list vec (Int idx)))
     (list
       (Instr 'movq (list vec (Reg 'r11)))
       (Instr 'movq (list (Deref 'r11 (* 8 (add1 idx))) x)))]
    [(Prim 'vector-set! (list vec (Int idx) val))
     (list
       (Instr 'movq (list vec (Reg 'r11)))
       (Instr 'movq (list (arg-instr val) (Deref 'r11 (* 8 (add1 idx)))))
       (Instr 'movq (list (Imm 0) x)))]
    [(Prim 'vector-length (list vec))
     (list
       (Instr 'movq (list vec (Reg 'r11)))
       (Instr 'andq (list (Imm #b111110) (Reg 'r11)))
       (Instr 'sarq (list (Imm 1) (Reg 'r11)))
       (Instr 'movq (list (Reg 'r11) x)))]
    [(Prim 'print (list n))
     (list
       (Instr 'movq (list (arg-instr n) (Reg 'rdi)))
       (Callq 'print_int 1)
       (Instr 'movq (list (Imm 0) x)))]
    [(Prim 'mutex-unlock (list m))
     (list
       (Instr 'movq (list (arg-instr m) (Reg 'rdi)))
       (Instr 'addq (list (Imm 1) (Reg 'rdi))) ; add 1 to get actual mutex
       (Callq 'unlock_mutex 1)
       (Instr 'movq (list (Imm 0) x)))]
    [(Prim 'call-parallel (list f1 v1 f2 v2))
     (list
       (Instr 'movq (list (arg-instr f1) (Reg 'rdi)))
       (Instr 'movq (list (arg-instr v1) (Reg 'rsi)))
       (Instr 'movq (list (arg-instr f2) (Reg 'rdx)))
       (Instr 'movq (list (arg-instr v2) (Reg 'rcx)))
       (Callq 'call_parallel 4)
       (Instr 'movq (list (Imm 0) (Reg 'rax))))]
    [(Prim cmp (list e1 e2))
     #:when (memv cmp cmp-ops)
     (list
       (Instr 'cmpq (list (arg-instr e2) (arg-instr e1)))
       (Instr 'set (list (cmp->cc cmp) (ByteReg 'al)))
       (Instr 'movzbq (list (ByteReg 'al) x)))]
    [(GlobalValue y)
     (list (Instr 'movq (list (Global y) x)))]
    [(Allocate n t)
     (let ([tag (binary->int (generate-vec-tag t))])
       (list
         (Instr 'movq (list (Global 'free_ptr) (Reg 'r11)))
         (Instr 'addq (list (Imm (* 8 (add1 n))) (Global 'free_ptr)))
         (Instr 'movq (list (Imm tag) (Deref 'r11 0)))
         (Instr 'movq (list (Reg 'r11) x))))]
    [(AllocateArray (Int n) t)
     (let ([tag (binary->int (generate-mutex-array-tag t))]) ; mutex array specific generation
       (list
         (Instr 'movq (list (Global 'free_ptr) (Reg 'r11)))
         (Instr 'addq (list (Imm (* 8 (add1 n))) (Global 'free_ptr)))
         (Instr 'movq (list (Imm tag) (Deref 'r11 0)))
         (Instr 'movq (list (Reg 'r11) x))))]
    ))

(define (select-stmt-instructions p)
  (match p
    [(Assign x e) (select-assign-instructions x e)]
    [(Call rator rands)
     (append
       (map
         (lambda (arg reg)
           (Instr 'movq (list (arg-instr arg) (Reg reg))))
         rands
         (take arg-pass-regs (length rands)))
       (list
         (IndirectCallq rator (length rands))))]
    [(Prim 'read (list))
     (list
       (Callq 'read_int 0))]
    [(Prim 'vector-set! (list vec (Int idx) val))
     (list
       (Instr 'movq (list vec (Reg 'r11)))
       (Instr 'movq (list (arg-instr val) (Deref 'r11 (* 8 (add1 idx))))))]
    [(Prim 'print (list n))
     (list
       (Instr 'movq (list (arg-instr n) (Reg 'rdi)))
       (Callq 'print_int 1))]
    [(Prim 'sleep (list n))
     (list
       (Instr 'movq (list (arg-instr n) (Reg 'rdi)))
       (Callq 'sleep_secs 1))]
    [(Prim 'mutex-lock (list m))
     (list
       (Instr 'movq (list (arg-instr m) (Reg 'rdi)))
       (Instr 'addq (list (Imm 1) (Reg 'rdi))) ; add 1 to get to actual mutex
       (Callq 'lock_mutex 1))]
    [(Prim 'mutex-unlock (list m))
     (list
       (Instr 'movq (list (arg-instr m) (Reg 'rdi)))
       (Instr 'addq (list (Imm 1) (Reg 'rdi))) ; add 1 to get actual mutex
       (Callq 'unlock_mutex 1))]
    [(Prim 'mutex-init (list m))
     (list
       (Instr 'movq (list (arg-instr m) (Reg 'rdi)))
       (Instr 'addq (list (Imm 1) (Reg 'rdi))) ; add 1 to get actual mutex
       (Callq 'init_mutex 1))]
    [(Prim 'call-parallel (list f1 v1 f2 v2))
     (list
       (Instr 'movq (list (arg-instr f1) (Reg 'rdi)))
       (Instr 'movq (list (arg-instr v1) (Reg 'rsi)))
       (Instr 'movq (list (arg-instr f2) (Reg 'rdx)))
       (Instr 'movq (list (arg-instr v2) (Reg 'rcx)))
       (Callq 'call_parallel 4))]
    [(Collect n)
     (list
       (Instr 'movq (list (Global 'rootstack_ptr) (Reg 'rdi)))
       (Instr 'movq (list (Imm n) (Reg 'rsi)))
       (Callq 'collect 2))]))

(define (select-tail-instructions fname p)
  (match p
    [(Seq stmt tail)
     (append
       (select-stmt-instructions stmt)
       (select-tail-instructions fname tail)) ]
    [(Return e)
     (append
       (select-assign-instructions (Reg 'rax) e)
       (list
         (Jmp
           (symbol-append fname '_conclusion))))]
    [(Goto l) (list (Jmp l))]
    [(TailCall rator rands)
     (append
       (map
         (lambda (arg reg)
           (Instr 'movq (list (arg-instr arg) (Reg reg))))
         rands
         (take arg-pass-regs (length rands)))
       (list (TailJmp rator (length rands))))]
    [(IfStmt (Prim cmp (list e1 e2)) (Goto l-thn) (Goto l-els))
     (list
       (Instr 'cmpq (list (arg-instr e2) (arg-instr e1)))
       (JmpIf (cmp->cc cmp) l-thn)
       (Jmp l-els))]))

(define (arg-regs-to-vars arg-ts)
  (map
    (lambda (reg var)
      (Instr 'movq (list (Reg reg) (Var (car var)))))
    (take arg-pass-regs (length arg-ts))
    arg-ts))

(define (select-block-instructions fname arg-ts)
  (lambda (b)
    (match-let ([(cons label tail) b])
      (cond
        [(string-suffix? (symbol->string label) "_start")
         (cons
           label
           (Block
             '()
             (append
               (arg-regs-to-vars arg-ts)
               (select-tail-instructions fname tail))))]
        [else (cons label (Block '() (select-tail-instructions fname tail)))]))))

(define (select-instructions-def d)
  (match-let ([(Def name arg-ts ret-t dinfo dbs) d])
    (Def
      name
      '()
      ret-t
      (cons (cons 'num-params (length arg-ts)) dinfo)
      (map (select-block-instructions name arg-ts) dbs))))

;; select-instructions : Cvar -> x86var
(define (select-instructions p)
  (match p
    [(ProgramDefs info ds)
     (X86ProgramDefs info (map select-instructions-def ds))]))

;;;;; UNCOVER LIVE

(define caller-saved-regs
  (set
    'rax 'rcx 'rdx 'rsi 'rdi
    'r8 'r9 'r10 'r11))

(define callee-saved-regs
  (set
    'rsp 'rbp 'rbx
    'r12 'r13 'r14 'r15))

(define reg-list
  (list
    ; full regs
    'rax 'rcx 'rdx 'rsi 'rdi
    'r8 'r9 'r10 'r11
    'rsp 'rbp 'rbx
    'r12 'r13 'r14 'r15))

(define return-val-reg (set 'rax))

(define (block-deps-instr! g bk)
  (lambda (instr)
    (match instr
      [(Jmp label)
       (add-directed-edge! g bk label)]
      [(JmpIf cc label)
       (add-directed-edge! g bk label)]
      [_ (void)])))

(define (block-deps! g)
  (lambda (bk) ; bk should be (cons label tail)
    (match-let ([(Block info is) (cdr bk)])
      (map
        (block-deps-instr! g (car bk))
        is))))

(define (find-block-deps bks)
  (let ([g (make-multigraph '())])
    (map
      (lambda (bk) (add-vertex! g (car bk)))
      bks)
    (add-vertex! g 'conclusion)
    (map (block-deps! g) bks)
    g))

(define (used-args-from-arity a)
  (let ([used-a (min a (length arg-pass-regs))])
    (list->set (take arg-pass-regs used-a))))

(define (symbfy-arg x)
  (match x
    [(Imm n) #f]
    [(Deref r n) #f]
    [(Var x) x]
    [(Reg r) r]
    [(ByteReg r) r] ; more for handling bytereg?]
    [(Global x) #f]))

(define (symbfy-args locs)
  (let* ([locs-l (set->list locs)]
         [flocs (filter-map symbfy-arg locs-l)])
    (list->set flocs)))

; l2l stand for label to live

(define (uncover-read-locs instr l2l)
  (match instr
    [(Instr 'movq (list src tgt)) (symbfy-args (set src))]
    [(Instr 'movzbq (list src tgt)) (symbfy-args (set src))]
    [(Instr 'set (list cc arg)) (set)]
    [(Instr 'cmpq (list arg1 arg2)) (symbfy-args (set arg1 arg2))]
    [(Instr i (list src tgt)) (symbfy-args (set src tgt))]
    [(Instr i (list tgt)) (symbfy-args (set tgt))] ; assuming no pushq or popq
    [(Callq label arity) (used-args-from-arity arity)]
    [(IndirectCallq f n)
     (set-union (used-args-from-arity n) (symbfy-args (set f)))]
    [(TailJmp f n)
     (set-union (used-args-from-arity n) (symbfy-args (set f)))]
    [(Retq) (set)] ; don't need to handle retq but included here for completeness
    [(Jmp lbl) (cdr (assv lbl l2l))]
    [(JmpIf cc lbl) (cdr (assv lbl l2l))]))

(define (uncover-write-locs instr l2l)
  (match instr
    [(Instr 'movq (list src tgt)) (symbfy-args (set tgt))]
    [(Instr 'movzbq (list src tgt)) (symbfy-args (set tgt))]
    [(Instr 'set (list cc arg)) (symbfy-args (set arg))]
    [(Instr 'cmpq (list arg1 arg2)) (set)]
    [(Instr i (list src tgt)) (symbfy-args (set tgt))]
    [(Instr i (list tgt)) (symbfy-args (set tgt))] ; assuming no pushq or popq
    [(Callq label arity) caller-saved-regs]
    [(IndirectCallq f n) caller-saved-regs]
    [(TailJmp f n) caller-saved-regs] ; not sure of this is correct
    [(Retq) (set)] ; don't need to handle retq but included here for completeness
    [(Jmp lbl) (cdr (assv lbl l2l))]
    [(JmpIf cc lbl) (cdr (assv lbl l2l))]))

(define (uncover-instrs-live instrs l2l)
  (cond
    [(null? instrs) (list (set))] ; (list block-live-after) ?
    [else
      (let*
        ([cdr-instrs-live (uncover-instrs-live (cdr instrs) l2l)]
         [live-after (car cdr-instrs-live)]
         [read-locs (uncover-read-locs (car instrs) l2l)]
         [written-locs (uncover-write-locs (car instrs) l2l)])
        (cons
          (set-union (set-subtract live-after written-locs) read-locs)
          cdr-instrs-live))]))

(define (uncover-block-live b l2l) ; -> block, l2l
  (match-let*
    ([(Block info is) b]
     [lives (uncover-instrs-live is l2l)]
     [info-no-la (remv (assv 'live-afters info) info)])
    (cons
      (Block
        (cons (cons 'live-afters (cdr lives)) info-no-la)
        is)
      (car lives))))

(define (uncover-instr-used-blocks instrs)
  (cond
    [(null? instrs) '()]
    [else
      (match (car instrs)
        [(Jmp lbl) #:when (not (eqv? lbl 'conclusion))
         (cons lbl (uncover-instr-used-blocks (cdr instrs)))]
        [(JmpIf cc lbl) #:when (not (eqv? lbl 'conclusion))
         (cons lbl (uncover-instr-used-blocks (cdr instrs)))]
        [else (uncover-instr-used-blocks (cdr instrs))])]))

(define (uncover-block-used-blocks b)
  (match-let
    ([(Block info is) b])
    (uncover-instr-used-blocks is)))

(define (iterate-analyze-dataflow! G q l2l bs) ; -> label-block list
  (if (not (queue-empty? q))
    (match-let*
      ([cur-label (dequeue! q)]
       [dependents (get-neighbors G cur-label)]
       [prev-live-before (cdr (assv cur-label l2l))]
       #| [input (apply set-union (set) (map (lambda (l) (cdr (assv l l2l))) deps)) ] |#
       [cur-label-block (assv cur-label bs)]
       [used-blocks (uncover-block-used-blocks (cdr cur-label-block))]
       [(cons ub ub-live-before) (uncover-block-live (cdr cur-label-block) l2l)]
       [del-l2l (remv (assv cur-label l2l) l2l)]
       [next-l2l (cons (cons cur-label ub-live-before) del-l2l)]
       [del-bs (remv (assv cur-label bs) bs)]
       [next-bs (cons (cons cur-label ub) del-bs)])
      (begin
        (if (set=? prev-live-before ub-live-before)
          (void)
          (map
            (lambda (lbl) (enqueue! q lbl))
            dependents))
        (iterate-analyze-dataflow! G q next-l2l next-bs)))
    bs))

(define (analyze-dataflow fname G bs) ; -> label-block list
  (let*
    ([q (make-queue)]
     [gen-bottom
       (map
         (lambda (label-block)
           (cons
             (car label-block)
             (set)))
         bs)]
     [bottom
       (cons
         (cons
           (symbol-append fname '_conclusion)
           (set 'rax 'rsp))
         gen-bottom)]
     )
    (begin
      (map
        (lambda (label-block)
          (enqueue! q (car label-block)))
        bs)
      (iterate-analyze-dataflow! G q bottom bs))))

(define (uncover-live-def d)
  (match-let ([(Def name arg-ts ret-t dinfo dbs) d])
    (let ([block-deps (find-block-deps dbs)])
      (Def
        name
        arg-ts
        ret-t
        (cons (cons 'block-deps block-deps) dinfo)
        (analyze-dataflow name (transpose block-deps) dbs)))))

(define (uncover-live p)
  (match p
    [(X86ProgramDefs info ds)
     (X86ProgramDefs info (map uncover-live-def ds))]))

;;;;; BUILD INTERFERENCE

(define (connect-vertices! g v1 v2)
  (set-map
    v1
    (lambda (e1)
      (set-map v2 (lambda (e2) (add-edge! g e1 e2))))))

(define (build-instr-interfere! g tuple-vars)
  (lambda (instr live-after)
    (match instr
      [(Instr 'movq (list src dest))
       (connect-vertices!
         g
         (set (symbfy-arg dest))
         (set-remove (set-remove live-after (symbfy-arg src)) (symbfy-arg dest)))]
      [(Instr 'movzbq (list src dest))
       (connect-vertices!
         g
         (set (symbfy-arg dest))
         (set-remove (set-remove live-after (symbfy-arg src)) (symbfy-arg dest)))]
      [(Instr 'set (list cc arg))
       (connect-vertices! g (set (symbfy-arg arg)) live-after)]
      [(Instr 'cmpq (list arg1 arg2)) (void)]
      [(Instr i (list src dest))
       (connect-vertices! g (set (symbfy-arg dest)) live-after)]
      [(Instr i (list dest)) ; assuming no pushq or popq
       (connect-vertices! g (set (symbfy-arg dest)) live-after)]
      [(Callq label arity)
       (let ([live-tuples (set-union live-after tuple-vars)])
         (connect-vertices! g caller-saved-regs live-after)
         (connect-vertices! g callee-saved-regs live-tuples))]
      [(IndirectCallq f arity)
       (let ([live-tuples (set-union live-after tuple-vars)])
         (connect-vertices! g caller-saved-regs live-after)
         (connect-vertices! g callee-saved-regs live-tuples))]
      [(TailJmp f arity)
       (let ([live-tuples (set-union live-after tuple-vars)])
         (connect-vertices! g caller-saved-regs live-after)
         (connect-vertices! g callee-saved-regs live-tuples))]
      [(Retq) (void)] ; don't need to handle retq but included here for completeness
      [(Jmp lbl) (void)]
      [(JmpIf cc lbl) (void)])))

(define (build-instr-mov-related! mg)
  (lambda (instr)
    (match instr
      [(Instr 'movq (list src dest))
       (connect-vertices! mg (set (symbfy-arg src)) (set (symbfy-arg dest)))]
      [else (void)])))

(define (build-instrs-interfere! g tuple-vars)
  (lambda (instrs live-afters)
    (map (build-instr-interfere! g tuple-vars) instrs live-afters)))

(define (build-instrs-mov-related! mg)
  (lambda (instrs)
    (map (build-instr-mov-related! mg) instrs)))

(define (build-block-interfere! g tuple-vars)
  (lambda (b)
    (match b
      [(Block info is)
       ((build-instrs-interfere! g tuple-vars)
        is
        (cdr (assv 'live-afters info)))])))

(define (build-block-mov-related! mg)
  (lambda (b)
    (match b
      [(Block info is)
       ((build-instrs-mov-related! mg) is)])))

(define (get-tuple-vars var-types)
  (map
    car
    (filter
      (lambda (x) (list? (cdr x))) ; might need better criteria eventually
      var-types)))

(define (build-interference-def d)
  (match d
    [(Def name arg-ts ret-t dinfo bs)
     (let ([g (unweighted-graph/undirected '())] ; interference graph
           [mg (unweighted-graph/undirected '())] ; mov-related graph
           [tuple-vars
             (list->set
               (get-tuple-vars (cdr (assv 'locals-types dinfo))))])
       (map
         (lambda (r) (add-vertex! g r))
         reg-list)
       (map
         (lambda (r) (add-vertex! mg r))
         (cons 'al reg-list))
       (map
         (lambda (v) (add-vertex! mg v))
         (cdr (assv 'locals dinfo)))
       (map
         (lambda (label-tail)
           ((build-block-interfere! g tuple-vars) (cdr label-tail)))
         bs)
       (map
         (lambda (label-tail)
           ((build-block-mov-related! mg) (cdr label-tail)))
         bs)
       (Def
         name
         arg-ts
         ret-t
         (cons
           (cons 'move-related mg)
           (cons
             (cons 'conflicts g)
             dinfo))
         bs))]))

(define (build-interference p)
  (match p
    [(X86ProgramDefs info ds)
     (X86ProgramDefs info (map build-interference-def ds))]))

;;;;; ALLOCATE REGISTERS

; colorinfo is `(,node ,(box (cons saturation color)))
; color is int (#f for no color)
; saturation is int list
; because node is first arg, can use assv

(define allocable-regs
  (list
    'rcx 'rdx 'rsi 'rdi
    'r8 'r9 'r10 'rbx
    'r12 'r13 'r14))

(define non-allocable-regs
  (list
    (gensym 'rfake) 'rax
    'rsp 'rbp 'r11 'r15))

(define (empty-colorinfos vars)
  (map
    (lambda (v)
      (cons v (box (list '() #f))))
    vars))

(define (saturate-colorinfo! ci c)
  (match-let*
    ([`(,node . ,b) ci]
     [`(,saturation ,color . ,tail) (unbox b)])
    (set-box!
      b
      `(,(cons c saturation) ,color . ,tail))))

(define (color-colorinfo! ci c)
  (match-let*
    ([`(,node . ,b) ci]
     [`(,saturation ,color . ,tail) (unbox b)])
    (set-box!
      b
      `(,saturation ,c . ,tail))))

(define (get-color-colorinfo ci)
  (match-let*
    ([`(,node . ,b) ci]
     [`(,saturation ,color . ,tail) (unbox b)])
    color))

(define (pqhandle-colorinfo! ci pqh)
  (match-let*
    ([`(,node . ,b) ci]
     [`(,saturation ,color . ,tail) (unbox b)])
    (set-box!
      b
      `(,saturation ,color ,pqh))))

(define (find-lowest-except l)
  (letrec
    ([f (lambda (x)
          (if (memv x l)
            (f (add1 x))
            x))])
    (f 0)))

(define (choose-color g cis mg l v)
  (let*
    ([lowest-valid-color (find-lowest-except l)]
     [invalid-colors (list->set l)]
     [move-related (get-neighbors mg v)]
     [move-related-colors
       (filter-map 
         (lambda (x)
           (let ([c (assv x cis)])
             (if c
               (get-color-colorinfo (assv x cis))
               #f)))
         move-related)]
     [move-related-regs
       (list->set
         (filter
           (lambda (x)
             (and (<= x 10) (>= x 0)))
           move-related-colors))]
     [move-related-stacklocs
       (list->set
         (filter
           (lambda (x) (> x 10))
           move-related-colors))]
     [valid-moverel-regs (set-subtract move-related-regs invalid-colors)]
     [valid-moverel-stacklocs (set-subtract move-related-stacklocs invalid-colors)])
    (cond
      [(not (set-empty? valid-moverel-regs))
       (car (set->list valid-moverel-regs))]
      [(<= lowest-valid-color 10)
       lowest-valid-color]
      [(not (set-empty? valid-moverel-stacklocs))
       (car (set->list valid-moverel-stacklocs))]
      [else
        lowest-valid-color])))
  

(define (reg->color r)
  (if (memv r allocable-regs)
    (index-of allocable-regs r)
    (- (index-of non-allocable-regs r))))

(define (precolor-regs! g regs cis)
  (map
    (lambda (v)
      (let ([color (reg->color v)])
        (if (memv v reg-list)
          (begin
            (color-colorinfo!  (assv v cis) color)
            (map
              (lambda (nv)
                (saturate-colorinfo! (assv nv cis) color))
              (get-neighbors g v)))
          (void))))
      regs))

(define (initialize-pqueue! vars pq cis)
  (map
    (lambda (v)
      (let
        ([pqh (pqueue-push! pq v)])
        (pqhandle-colorinfo! (assv v cis) pqh)))
    vars))

(define (dsatur! g pq cis mg) ; graph, priorityqueue (of nodes), (var  -> colorinfos
  (if (not (zero? (pqueue-count pq)))
    (match-let*
      ([v (pqueue-pop! pq)]
       [ci (assv v cis)]
       [`(,node . ,b) ci]
       [`(,saturation ,pcolor . ,tail) (unbox b)]
       [color (choose-color g cis mg saturation v)]) ; (define (choose-color g cis mg l v)
      (begin
        (color-colorinfo! ci color)
        (map
          (lambda (nv)
            (saturate-colorinfo! (assv nv cis) color))
          (get-neighbors g v))
        (dsatur! g pq cis mg)))
    (void)))

(define (color-graph! g mg)
  (let*
    ([vertices (get-vertices g)]
     [cis (empty-colorinfos vertices)]
     [pqueue-cmp
       (lambda (var1 var2)
         (match-let*
           ([`(,saturation1 ,color1 . ,tail1) (unbox (cdr (assv var1 cis)))]
            [`(,saturation2 ,color2 . ,tail2) (unbox (cdr (assv var2 cis)))])
           (cond
             [(= (length saturation1) (length saturation2))
              (let*
                ([moverels1 (list->set (get-neighbors mg var1))]
                 [moverels2 (list->set (get-neighbors mg var2))]
                 [validmoverels1 (set-subtract moverels1 (list->set saturation1))]
                 [validmoverels2 (set-subtract moverels2 (list->set saturation2))])
                (cond
                  [(set-empty? validmoverels1) #f]
                  [(set-empty? validmoverels2) #t]
                  [else (<= (length (set->list validmoverels1)) (length (set->list validmoverels2)))]))]
             [else (<= (length saturation1) (length saturation2))])))]
     [pq (make-pqueue pqueue-cmp)]
     [_ (precolor-regs! g reg-list cis)]
     [_ (initialize-pqueue! (remv* reg-list vertices) pq cis)])
    (dsatur! g pq cis mg)
    (map
      (lambda (ci)
        (match-let*
          ([`(,var . ,b) ci]
           [`(,saturation ,color . ,tail) (unbox b)])
          (cons var color)))
      cis)))

(define (assign-arg-homes env)
  (lambda (p)
    (match p
      [(Imm n) (Imm n)]
      [(Reg r) (Reg r)]
      [(ByteReg r) (ByteReg r)]
      [(Var x)
       (cdr (assv x env))]
      [(Global x) (Global x)]
      [(Deref reg n) (Deref reg n)])))

(define (assign-instr-homes env)
  (lambda (p)
    (match p
      [(IndirectCallq f arity)
       (IndirectCallq ((assign-arg-homes env) f) arity)]
      [(TailJmp f arity)
       (TailJmp ((assign-arg-homes env) f) arity)]
      [(Instr 'set (list cc dst))
       (Instr 'set (list cc ((assign-arg-homes env) dst)))]
      [(Instr op es)
       (Instr op (map (assign-arg-homes env) es))]
      [i i])))

(define (assign-block-homes env)
  (lambda (p)
    (match p
      [(Block info is)
       (Block info (map (assign-instr-homes env) is))])))

(define (homes->used-callee-saved-regs h)
  (let*
    ([var-homes
       (filter
         (lambda (x) (not (memv (car x) reg-list)))
         h)]
     [used-regs (map cdr var-homes)])
    (filter
      (lambda (x)
        (set-member? callee-saved-regs (symbfy-arg x)))
      used-regs)))

(define (color-to-home c)
  (cond
    [(< c 0) (Reg (list-ref non-allocable-regs (- c)))]
    [(<= c 10) (Reg (list-ref allocable-regs c))]
    [(>= c 11) (Deref 'rbp (* -8 (- c 11)))]))

(define (count-rootstack-homes env)
  (match env
    ['() 0]
    [`((,vname . ,(Deref 'r15 n)) . ,d)
      (add1 (count-stack-homes d))]
    [`(,a . ,d) (count-stack-homes d)]))

(define (count-stack-homes env)
  (match env
    ['() 0]
    [`((,vname . ,(Deref 'rbp n)) . ,d)
      (add1 (count-stack-homes d))]
    [`(,a . ,d) (count-stack-homes d)]))

(define (assign-colors-to-homes colors types) ; -> (var . home) list
  (foldl
    (lambda (col res)
      (cond
        [(< (cdr col) 0)
         (cons
           (cons
             (car col)
             (Reg (list-ref non-allocable-regs (- (cdr col)))))
           res)]
        [(<= (cdr col) 10)
         (cons
           (cons
             (car col)
             (Reg (list-ref allocable-regs (cdr col))))
           res)]
        [(>= (cdr col) 11)
         (match (cdr (assv (car col) types))
           [`(Vector ,ts ...)
             (cons
               (cons
                 (car col)
                 (Deref 'r15 (* 8 (add1 (count-rootstack-homes res)))))
               res)]
           [else
             (cons
               (cons
                 (car col)
                 (Deref 'rbp (* -8 (add1 (count-stack-homes res)))))
               res)])]))
    '()
    colors))

(define (allocate-registers-def d)
  (match d
    [(Def name arg-ts ret-t info es)
     (let*
       ([g (cdr (assv 'conflicts info))]
        [mg (cdr (assv 'move-related info))]
        [var-types (cdr (assv 'locals-types info))]
        [assigned-colors (color-graph! g mg)]
        [assigned-homes
          (assign-colors-to-homes assigned-colors var-types)]
        [num-spills (count-stack-homes assigned-homes)]
        [num-root-spills (count-rootstack-homes assigned-homes)])
       (Def
         name
         arg-ts
         ret-t
         (append
           (list
             (cons 'num-spills num-spills)
             (cons 'num-root-spills num-root-spills)
             (cons 'used-callee (homes->used-callee-saved-regs assigned-homes)))
           info)
         (map
           (lambda (label-tail)
             (cons
               (car label-tail)
               ((assign-block-homes assigned-homes) (cdr label-tail))))
           es)))]))

(define (allocate-registers p)
  (match p
    [(X86ProgramDefs info ds)
     (X86ProgramDefs info (map allocate-registers-def ds))]))

;;;;; PATCH INSTRUCTIONS

(define (patch-instr-instructions p) ; Instr -> Instr list
  (match p
    [(Instr 'addq (list (Deref reg1 pos1) (Deref reg2 pos2)))
     (list
       (Instr 'movq (list (Deref reg1 pos1) (Reg 'rax)))
       (Instr 'addq (list (Reg 'rax) (Deref reg2 pos2))))]
    [(Instr 'addq (list (Imm n) (Deref reg int))) #:when (> n 65536)
     (list
       (Instr 'movq (list (Imm n) (Reg 'rax)))
       (Instr 'addq (list (Reg 'rax) (Deref reg int))))]
    [(Instr 'subq (list (Deref reg1 pos1) (Deref reg2 pos2)))
     (list
       (Instr 'movq (list (Deref reg1 pos1) (Reg 'rax)))
       (Instr 'subq (list (Reg 'rax) (Deref reg2 pos2))))]
    [(Instr 'subq (list (Imm n) (Deref reg int))) #:when (> n 65536)
     (list
       (Instr 'movq (list (Imm n) (Reg 'rax)))
       (Instr 'subq (list (Reg 'rax) (Deref reg int))))]
    [(Instr 'movq (list p p))
     (list)]
    [(Instr 'movq (list (Deref reg1 pos1) (Deref reg2 pos2)))
     (list
       (Instr 'movq (list (Deref reg1 pos1) (Reg 'rax)))
       (Instr 'movq (list (Reg 'rax) (Deref reg2 pos2))))]
    [(Instr 'movq (list (Imm n) (Deref reg int))) #:when (> n 65536)
     (list
       (Instr 'movq (list (Imm n) (Reg 'rax)))
       (Instr 'movq (list (Reg 'rax) (Deref reg int))))]
    [(Instr 'cmpq (list a (Imm n)))
     (list
       (Instr 'movq (list (Imm n) (Reg 'rax)))
       (Instr 'cmpq (list a (Reg 'rax))))]
    [(Instr 'cmpq (list (Deref reg1 pos1) (Deref reg2 pos2)))
     (list
       (Instr 'movq (list (Deref reg1 pos1) (Reg 'rax)))
       (Instr 'cmpq (list (Reg 'rax) (Deref reg2 pos2))))]
    [(Instr 'movzbq (list (ByteReg br) (Deref reg2 pos2))) ; will never movzbq into an imm, so don't worry
     (list
       (Instr 'movsbq (list (ByteReg br) (Reg 'rax)))
       (Instr 'mov (list (Reg 'rax) (Deref reg2 pos2))))]
    [(Instr 'leaq (list f (Deref reg pos)))
     (list
       (Instr 'leaq (list f (Reg 'rax)))
       (Instr 'movq (list (Reg 'rax) (Deref reg pos))))]
    [(TailJmp (Reg 'rax) arity)
     (list (TailJmp (Reg 'rax) arity))]
    [(TailJmp f arity)
     (list
       (Instr 'movq (list f (Reg 'rax)))
       (TailJmp (Reg 'rax) arity))]
    [else (list p)]))
  
(define (patch-block-instructions p)
  (match p
    [(Block info is)
     (Block info (append-map patch-instr-instructions is))]))

(define (patch-instructions-def d)
  (match-let ([(Def name arg-ts ret-t dinfo bs) d])
    (Def
      name
      arg-ts
      ret-t
      dinfo
      (map
        (lambda (label-tail)
          (cons
            (car label-tail)
            (patch-block-instructions (cdr label-tail))))
        bs))))

;; patch-instructions : x86var -> x86int
(define (patch-instructions p)
  (match p
    [(X86ProgramDefs info ds)
     (X86ProgramDefs
       info
       (map patch-instructions-def ds))]))

;;;;; PRELUDE AND CONCLUSION

(define (prelude name nvars nrootvars used-callees is-main)
  (append
    ; handle stack stuff
    (map
      (lambda (x)
        (Instr 'pushq (list x)))
      used-callees)
    (list
      (Instr 'pushq (list (Reg 'rbp)))
      (Instr 'movq (list (Reg 'rsp) (Reg 'rbp)))
      (Instr 'subq (list (Imm (align (* 8 nvars) 16)) (Reg 'rsp))))
    ; handle rootstack stuff
    (if is-main
      (list
        (Instr 'movq (list (Imm 65536) (Reg 'rdi)))
        (Instr 'movq (list (Imm 65536) (Reg 'rsi)))
        (Callq 'initialize 0)
        (Instr 'movq (list (Global 'rootstack_begin) (Reg 'r15)))
        (Instr 'movq (list (Reg 'r15) (Global 'rootstack_ptr))))
      '())
    (append-map
      (lambda (n)
        (list
          (Instr 'movq (list (Global 'rootstack_ptr) (Reg 'r15)))
          (Instr 'movq (list (Imm 0) (Deref 'r15 (* n 8))))))
      (range nrootvars))
    (list
      (Instr 'addq (list (Imm (* 8 nrootvars)) (Global 'rootstack_ptr))))
    ; start
    (list (Jmp (symbol-append name '_start)))))

(define (cleanup nvars nrootvars used-callees)
  (append
    ; handle rootstack stuff
    (list
      (Instr 'subq (list (Imm (* 8 nrootvars)) (Global 'rootstack_ptr))))
    ; handle stack stuff
    (list
      (Instr 'addq (list (Imm (align (* 8 nvars) 16)) (Reg 'rsp)))
      (Instr 'popq (list (Reg 'rbp))))
    (map
      (lambda (x)
        (Instr 'popq (list x)))
      (reverse used-callees))))

(define (conclusion nvars nrootvars used-callees)
  (append
    (cleanup nvars nrootvars used-callees)
    (list (Retq))))

(define (translate-tailjmp-instrs nvars nrootvars used-callees)
  (lambda (instrs)
    (match instrs
      ['() '()]
      [(list (TailJmp f arity))
       (append
         (cleanup nvars nrootvars used-callees)
         (list (IndirectJmp f)))]
      [(cons a d)
       (cons a ((translate-tailjmp-instrs nvars nrootvars used-callees) d))])))

(define (translate-tailjmp-block nvars nrootvars used-callees)
  (lambda (label-tail)
    (match-let ([(cons label (Block info is)) label-tail])
      (cons
        label
        (Block
          info
          ((translate-tailjmp-instrs nvars nrootvars used-callees) is))))))

(define (prelude-and-conclusion-def d)
  (match d
    [(Def name arg-ts ret-t info es) ; handle main
     (let*
       ([nvars (cdr (assv 'num-spills info))]
        [nrootvars (cdr (assv 'num-root-spills info))]
        [used-callees (cdr (assv 'used-callee info))]
        [aligned-used-callees
          (if (odd? (length used-callees))
            (cons (Reg 'rbp) used-callees)
            used-callees)]
        [is-main (eqv? name 'main)])
       (Def
         name
         arg-ts
         ret-t
         info
         (append
           (list
             (cons
               name
               (Block
                 '()
                 (prelude name nvars nrootvars aligned-used-callees is-main)))
             (cons
               (symbol-append name '_conclusion)
               (Block '() (conclusion nvars nrootvars aligned-used-callees))))
           (map (translate-tailjmp-block nvars nrootvars used-callees) es))))]))

;; prelude-and-conclusion : x86int -> x86int
(define (prelude-and-conclusion p)
  (match p
    [(X86ProgramDefs info ds)
     (let*
       ([updated-defs (map prelude-and-conclusion-def ds)]
        [def-infos
          (map
            (lambda (d)
              (match-let ([(Def _ _ _ dinfo _) d])
                dinfo))
            updated-defs)]
       [blocks
         (append*
           (map
             (lambda (d)
               (match-let ([(Def _ _ _ _ bs) d])
                 bs))
             updated-defs))]
       [num-root-spills
         (foldl
           +
           0
           (map
             (lambda (info)
               (cdr (assv 'num-root-spills info)))
             def-infos))])
     (X86Program
       (append
         (list
           (cons 'num-root-spills num-root-spills))
         info)
       blocks))]))

;;;;; COMPILER PASSES

(define (print-res r)
  (println r)
  r)

;; Define the compiler passes to be used by interp-tests and the grader
;; Note that your compiler file (the file that defines the passes)
;; must be named "compiler.rkt"
(define compiler-passes
  `(
     ;; Uncomment the following passes as you finish them.
     
     ; passes including interpreters
     ; ("shrink" ,shrink ,interp-Lfun-prime ,type-check-Lfun)
     ; ("partial evaluation" ,pe-Lint ,interp-Lfun-prime ,type-check-Lfun)
     ; ("uniquify" ,uniquify ,interp-Lfun-prime ,type-check-Lfun)
     ; ("uncover get!" ,uncover-get! ,interp-Lfun-prime ,type-check-Lfun)
     ; ("reveal functions" ,reveal-functions ,interp-Lfun-prime ,type-check-Lfun)
     ; ("limit functions" ,limit-functions ,interp-Lfun-prime ,type-check-Lfun)
     ; ("expose allocation" ,expose-allocation ,interp-Lfun-prime ,type-check-Lfun)
     ; ("remove complex opera*" ,remove-complex-opera* ,interp-Lfun-prime ,type-check-Lfun)
     ; ("explicate control" ,explicate-control ,interp-Cfun ,type-check-Cfun)
     ; ("instruction selection" ,select-instructions ,interp-pseudo-x86-3)
     ; ("uncover live" ,uncover-live ,interp-pseudo-x86-3)
     ; ("build interference" ,build-interference ,interp-pseudo-x86-3)
     ; ("allocate registers" ,allocate-registers ,interp-x86-3)
     ; ("patch instructions" ,patch-instructions ,interp-x86-3)
     ; ("prelude and conclusion" ,prelude-and-conclusion #f)

     ; passes exluding interpreters
     ("shrink" ,shrink #f ,type-check-Lthread)
     ("partial evaluation" ,pe-Lint #f ,type-check-Lthread)
     ("uniquify" ,uniquify #f ,type-check-Lthread)
     ("uncover get!" ,uncover-get! #f ,type-check-Lthread)
     ("reveal functions" ,reveal-functions #f ,type-check-Lthread)
     ("limit functions" ,limit-functions #f ,type-check-Lthread)
     ("expose allocation" ,expose-allocation #f ,type-check-Lthread)
     ("remove complex opera*" ,remove-complex-opera* #f ,type-check-Lthread)
     ("explicate control" ,explicate-control #f ,type-check-Cthread)
     ("instruction selection" ,select-instructions #f)
     ("uncover live" ,uncover-live #f)
     ("build interference" ,build-interference #f)
     ("allocate registers" ,allocate-registers #f)
     ("patch instructions" ,patch-instructions #f)
     ("prelude and conclusion" ,prelude-and-conclusion #f)
     ))