(define-library (csc cps)
  (export
    closure-convert
    ir1->ir2)
  (import (scheme base)
          (only (csc gensym)
            gensym
            gensym->int)
          (only (csc hash-map)
            insert
            key-not-found-error?
            lookup
            make-map
            map->alist
            merge)
          (only (csc ir1)
            %call
            %call-builtin
            %define-syntax
            %if
            %lambda
            %letrec
            %lexical-ref
            %lexical-set
            %library-define
            %library-ref
            %sequence
            call?
            constant?
            if?
            lambda?
            letrec-gensyms
            letrec-names
            letrec-values
            letrec?
            lexical-ref-gensym
            lexical-ref?
            lexical-set?
            library-define?
            library-ref?
            make-call
            make-call-builtin
            make-constant
            make-if
            make-lambda
            make-letrec
            make-lexical-ref
            make-lexical-set
            make-library-define
            make-library-ref
            make-sequence
            sequence?)
          (only (csc ir2)
            %apply
            %branch
            %fix
            %primitive
            *globals*
            branch-atom
            closure-arguments
            closure-body
            closure-name
            make-apply
            make-atom
            make-branch
            make-call-closure
            make-closure
            make-fix
            make-kargs
            make-klabel
            make-ktail
            make-primitive
            make-variable)
          (only (csc loop)
            loop
            return)
          (only (csc match)
            define-match-record-type
            match))
  (begin


    (define (new-ref)
      (make-lexical-ref 'generated-symbol (gensym)))


    ; Converts an IR1 expression to an equivalent expression where every
    ; procedure takes exactly one argument.
    (define (argument-conversion expr)
      (match expr
        (_ when (or (constant? expr)
                    (lexical-ref? expr)
                    (library-ref? expr))
          expr)
        ((% %lexical-set ref arg)
          (make-lexical-set ref (argument-conversion arg)))
        ((% %library-define ref arg)
          (make-library-define ref (argument-conversion arg)))
        ((% %define-syntax _ _) expr)
        ((% %if test consequent alternate)
          (make-if
            (argument-conversion test)
            (argument-conversion consequent)
            (argument-conversion alternate)))
        ((% %call proc args)
          (define argvec (new-ref))
          (define nargs (length args))
          (make-call
            (make-lambda (list argvec) #f
              (make-sequence
                (loop for arg in args
                      for i from 2
                      with expr = (make-sequence
                                    (make-call-builtin 'poke (list (make-constant 0) argvec (make-constant 0)))
                                    (make-call-builtin 'poke (list (make-constant nargs) argvec (make-constant 1))))
                      do (set! expr (make-sequence
                                      expr
                                      (make-call-builtin 'poke (list (argument-conversion arg) argvec (make-constant i)))))
                      finally (return expr))
                (make-call (argument-conversion proc) (list argvec))))
            (list (make-call-builtin 'alloc (list (make-constant (+ 2 nargs)))))))
        ((% %call-builtin op args)
          (make-call-builtin op (map argument-conversion args)))
        ((% %sequence head tail)
          (make-sequence
            (argument-conversion head)
            (argument-conversion tail)))
        ((% %lambda args rest body) when rest
          (define argvec (new-ref))
          (define nargs (length args))
          (make-lambda (list argvec) #f
            (make-if (make-call-builtin 'int<? (list (make-call-builtin 'peek (list argvec (make-constant 1)))
                                                     (make-constant nargs)))
              (make-call (make-library-ref 'wrong-number-of-arguments '(csc based)) (list argvec))
              (loop for arg in (reverse args)
                    for i downfrom (+ 1 nargs)
                    with expr = (make-call (make-lambda (list rest) #f
                                             (argument-conversion body))
                                           (list
                                             (argument-conversion
                                               (make-call (make-library-ref 'vector->list '(csc based))
                                                          (list argvec (make-constant nargs))))))
                    ; I'm relying on beta reduction here.
                    do (set! expr (make-call (make-lambda (list arg) #f
                                               expr)
                                             (list (make-call-builtin 'peek (list argvec (make-constant i))))))
                    finally (return expr)))))
        ((% %lambda args _ body)
          (define argvec (new-ref))
          (define nargs (length args))
          (make-lambda (list argvec) #f
            (make-if (make-call-builtin 'int=? (list (make-call-builtin 'peek (list argvec (make-constant 1)))
                                                     (make-constant nargs)))
              (loop for arg in (reverse args)
                    for i downfrom (+ 1 nargs)
                    with expr = (argument-conversion body)
                    do (set! expr (make-call (make-lambda (list arg) #f
                                               expr)
                                             (list
                                               (make-call-builtin 'peek (list argvec (make-constant i))))))
                    finally (return expr))
              (make-call (make-library-ref 'wrong-number-of-arguments '(csc based)) (list argvec)))))
        ((% %letrec in-order? names gensyms exprs body)
          (make-letrec in-order? names gensyms (map argument-conversion exprs) (argument-conversion body)))
        (_ (error "Unexpected form in argument-conversion" expr))))


    ; Update is a CPS expression that is used internally as part of
    ; CPS conversion.
    ; Update expressions are then removed by box-conversion.
    (define-match-record-type <update>
      (make-update ref atom continuation)
      update?
     %update
      (ref update-ref)
      (atom update-atom)
      (continuation update-continuation))


    (define (collect-functions-and-variables expr)
      (let ((names (letrec-names expr))
            (gensyms (letrec-gensyms expr))
            (vals (letrec-values expr)))
        (loop for name in names
              for gensym in gensyms
              for value in vals
              if (lambda? value)
                collect (match value
                          ((% %lambda args _ body)
                            (define continuation (new-ref))
                            (make-closure
                              (make-lexical-ref name gensym)
                              (cons continuation args)
                              (to-cps
                                body
                                (lambda (z)
                                  (make-apply continuation (list z)))))))
                        into functions
              else
                collect (make-lexical-ref name gensym) into variable-names
                and collect value into variable-values
              finally (return (values functions variable-names variable-values)))))


    ; Converts the given IR1 expression that has undergone argument conversion
    ; into an IR2 expression in continuation passing style.
    ; The resulting expression will include <update> forms.
    (define (to-cps expr continuation)
      (match expr
        (_ when (or (constant? expr)
                    (lexical-ref? expr))
          (continuation expr))
        ((% %library-ref . _)
          (unless (library-ref? expr)
            (error "wtf"))
          (define temp (new-ref))
          (make-primitive 'peek (list *globals* expr) (list temp)
            (continuation temp)))
        ((% %lexical-set ref arg)
          (to-cps
            arg
            (lambda (val)
              (make-update ref val (continuation (make-constant #f))))))
        ((% %library-define ref arg)
          (to-cps
            arg
            (lambda (val)
              (make-update ref val (continuation (make-constant #f))))))
        ((% %define-syntax _ _)
          ; no-op
          (continuation (make-constant #f)))
        ((% %if test consequent alternate)
          (to-cps
            test
            (lambda (val)
              (define continuation-ref (new-ref))
              (define result-ref (new-ref))
              (make-fix
                (list (make-closure continuation-ref (list result-ref)
                        (continuation result-ref)))
                (make-branch val
                  (to-cps
                    consequent
                    (lambda (result)
                      (make-apply continuation-ref (list result))))
                  (to-cps
                    alternate
                    (lambda (result)
                      (make-apply continuation-ref (list result)))))))))
        ((% %call proc (arg))
          (define return-address (new-ref))
          (define result (new-ref))
          (make-fix
            (list (make-closure return-address (list result) (continuation result)))
            (to-cps
              proc
              (lambda (f)
                (to-cps
                  arg
                  (lambda (v)
                    (make-apply f (list return-address v))))))))
        ((% %call-builtin op args)
          (define returns-value? (not (symbol=? op 'poke)))
          (loop for arg in (reverse args)
                with expr = (lambda (vals)
                              (if returns-value?
                                (let ((result (new-ref)))
                                  (make-primitive op (reverse vals) (list result)
                                    (continuation result)))
                                (make-primitive op (reverse vals) '()
                                  (continuation (make-constant #f)))))
                do (set! expr (let ((e* expr)  ; make copies to avoid modifying the expr in the closure.
                                    (arg* arg))
                                (lambda (vals)
                                  (to-cps arg*
                                    (lambda (val)
                                      (e* (cons val vals)))))))
                finally (return (expr '()))))
        ((% %sequence head tail)
          (to-cps
            head
            (lambda (x)
              (to-cps
                tail
                continuation))))
        ((% %lambda (arg) _ body)
          (define f (new-ref))
          (define k (new-ref))
          (make-fix
            (list
              (make-closure f (list k arg)
                (to-cps
                  body
                  (lambda (ret)
                    (make-apply k (list ret))))))
            (continuation f)))
        ((% %letrec _ _ _ _ body)
          (define-values (functions variable-names variable-values) (collect-functions-and-variables expr))
          (if (null? variable-names)
            (make-fix functions
              (to-cps body continuation))
            (let ((new-expr (loop for var in (reverse variable-names)
                                  for val in (reverse variable-values)
                                  with new-body = (to-cps body continuation)
                                  do (set! new-body (to-cps val (lambda (x)
                                                                  (make-update var x
                                                                    new-body))))
                                  finally (return new-body))))
              (unless (null? functions)
                (set! new-expr (make-fix functions new-expr)))
              (loop for var in variable-names
                    do (set! new-expr (make-primitive 'alloc (list (make-constant 1)) (list var)
                                        new-expr))
                    finally (return new-expr)))))
        (_ (error "unexpected type in to-cps" expr))))


    (define (make-ref-map)
      (make-map
        (lambda (ref)
          (gensym->int (lexical-ref-gensym ref)))
        (lambda (x y)
          (< (gensym->int (lexical-ref-gensym x))
                         (gensym->int (lexical-ref-gensym y))))))


    (define (get-boxed expr)
      (match expr
        ((% %update ref _ continuation)
          (define m (get-boxed continuation))
          (when (lexical-ref? ref)
            (set! m (insert m ref #t)))
          m)
        ((% %primitive _ _ _ continuation)
          (get-boxed continuation))
        ((% %branch _ true false)
          (merge
            (get-boxed true)
            (get-boxed false)))
        ((% %apply proc args)
          (make-ref-map))
        ((% %fix funs body)
          (loop with m = (get-boxed body)
                for fun in funs
                do (set! m (merge m (get-boxed (closure-body fun))))
                finally (return m)))
        (_ (error "Unexpected form in get-boxed" expr))))


    ; Rewrites the given expression to have no more <update> forms.
    (define (box-conversion expr)
      (define boxed-refs (get-boxed expr))
      (define (boxed? ref)
        (and (lexical-ref? ref)
             (guard (e ((key-not-found-error? e) #f))
               (lookup boxed-refs ref))))
      (define (convert-arg-list args)
        (define boxed-args (loop for arg in args
                                 if (boxed? arg)
                                   collect arg))
        (define vars (loop for x in boxed-args
                           collect (new-ref)))
        (define new-args (loop with v* = vars
                               for arg in args
                               collect (if (boxed? arg)
                                         (car v*)
                                         arg)
                               if (boxed? arg)
                                 do (set! v* (cdr v*))))
        (values new-args boxed-args vars))
      (let convert ((expr expr))
        (match expr
          ((% %update ref atom continuation) when (library-ref? ref)
            (make-primitive 'poke (list atom *globals* ref) '()
              (convert continuation)))
          ((% %update ref atom continuation) when (lexical-ref? ref)
            (make-primitive 'poke (list atom ref (make-constant 0)) '() (convert continuation)))
          ((% %primitive op args res continuation)
            ; Note that no reference in res can be boxed.
            (define-values (new-args boxed-args vars) (convert-arg-list args))
            (define new-expr (make-primitive op new-args res (convert continuation)))
            (loop for arg in boxed-args
                  for var in vars
                  do (set! new-expr (make-primitive 'peek (list arg (make-constant 0)) (list var)
                                      new-expr))
                  finally (return new-expr)))
          ((% %branch atom true false) when (boxed? atom)
            (define temp (new-ref))
            (make-primitive 'peek (list atom (make-constant 0)) (list temp)
              (make-branch temp
                (convert true)
                (convert false))))
          ((% %branch atom true false)
            (make-branch atom
              (convert true)
              (convert false)))
          ((% %apply proc args)
            (define-values (new-params boxed-params vars) (convert-arg-list (cons proc args)))
            (define new-expr (make-apply (car new-params) (cdr new-params)))
            (loop for p in boxed-params
                  for var in vars
                  do (set! new-expr (make-primitive 'peek (list p (make-constant 0)) (list var)
                                      new-expr))
                  finally (return new-expr)))
          ((% %fix funs body)
            (define-values (new-names boxed-names temp-names) (convert-arg-list (loop for fun in funs
                                                                                      collect (closure-name fun))))
            (define new-funs (loop for fun in funs
                                   for new-name in new-names
                                   collect (let-values (((new-args boxed-args temp-args) (convert-arg-list (closure-arguments fun))))
                                             (make-closure
                                               new-name
                                               new-args
                                               (let ((new-expr (convert (closure-body fun))))
                                                 (loop for arg in boxed-args
                                                       for var in temp-args
                                                       do (set! new-expr (make-primitive 'alloc (list (make-constant 1)) (list arg)
                                                                           (make-primitive 'poke (list var arg (make-constant 0)) '()
                                                                             new-expr)))
                                                       finally (return new-expr)))))))
            (define new-body (convert body))
            (loop for name in boxed-names
                  for var in temp-names
                  do (set! new-body (make-primitive 'poke (list var name (make-constant 0)) '()
                                      new-body)))
            (define new-expr (make-fix new-funs new-body))
            (loop for name in boxed-names
                  do (set! new-expr (make-primitive 'alloc (list (make-constant 1)) (list name)
                                      new-expr))
                  finally (return new-expr)))
          (_ (error "Unexpected form in box-conversion" expr)))))


    (define (ir1->ir2 expr continuation)
      (box-conversion
        (to-cps
          (argument-conversion expr)
          continuation)))


    (define (free-vars-expr expr bound-vars)
      (define (free? ref)
        (and (lexical-ref? ref)
             (not (guard (e ((key-not-found-error? e) #f))
                    (lookup bound-vars ref)))))
      (match expr
        ((% %primitive _ args res continuation)
          (loop for r in res
                if (lexical-ref? r)
                  do (set! bound-vars (insert bound-vars r #t)))
          (loop with m = (free-vars-expr continuation bound-vars)
                for arg in args
                if (free? arg)
                  do (set! m (insert m arg #t))
                finally (return m)))
        ((% %branch atom true false)
          (define m (merge (free-vars-expr true bound-vars) (free-vars-expr false bound-vars)))
          (if (free? atom)
            (set! m (insert m atom #t)))
          m)
        ((% %apply proc args)
          (define m (make-ref-map))
          (if (free? proc)
            (set! m (insert m proc #t)))
          (loop for arg in args
                if (free? arg)
                  do (set! m (insert m arg #t))
                finally (return m)))
        ((% %fix funs body)
          (loop for fun in funs
                for name = (closure-name fun)
                if (lexical-ref? name)
                  do (set! bound-vars (insert bound-vars name #t)))
          (define m (free-vars-expr body bound-vars))
          (loop for fun in funs
                do (set! m (merge m (free-vars-closure fun bound-vars)))
                finally (return m)))
        (_ (error "Unexpected form in free-vars-expr" expr))))


    (define (free-vars-closure fun bound-vars)
      (define name (closure-name fun))
      (when (lexical-ref? name)
        (set! bound-vars (insert bound-vars name #t)))
      (loop for arg in (closure-arguments fun)
            do (set! bound-vars (insert bound-vars arg #t)))
      (free-vars-expr (closure-body fun) bound-vars))


    ; Returns a list of the free variables in a closure.
    (define (free-vars expr)
      (define m (free-vars-closure expr (make-ref-map)))
      (map car (map->alist m)))


    (define (translate-ref ref env)
      (if (lexical-ref? ref)
        (guard (e ((key-not-found-error? e) (error "Undefined symbol in closure-convert" ref)))
          (lookup env ref))
        ref))


    ; converts a CPS expression into an equivalent expression with no
    ; free variables.
    (define (closure-convert expr)
      (let convert ((expr expr)
                    (env (make-ref-map)))
        (define (translate ref)
          (translate-ref ref env))
        (match expr
          ((% %primitive op args res continuation)
            (loop for r in res
                  do (set! env (insert env r (make-variable (gensym)))))
            (make-primitive op (map translate args) (map translate res) (convert continuation env)))
          ((% %branch atom true false)
            (make-branch (translate atom) (convert true env) (convert false env)))
          ((% %apply proc args)
            (let ((p (translate proc))
                  (fn (make-variable (gensym))))
              (make-primitive 'peek (list p (make-constant 0)) (list fn)
                (make-apply fn (cons p (map translate args))))))
          ((% %fix functions body)
            (define frees (map free-vars functions))
            (define fn-ptrs (loop for fun in functions
                                  collect (make-variable (gensym))))
            (define converted-functions (loop for fun in functions
                                              for fn-ptr in fn-ptrs
                                              for free-list in frees
                                              for env* = env
                                              for name = (closure-name fun)
                                              for closure = (make-variable (gensym))
                                              if (lexical-ref? name)
                                                do (set! env* (insert env* name closure))
                                              do (loop for arg in (closure-arguments fun)
                                                       do (set! env* (insert env* arg (make-variable (gensym)))))
                                                 (loop for var in free-list
                                                       do (set! env* (insert env* var (make-variable (gensym)))))
                                              collect (let ((new-body (convert (closure-body fun) env*)))
                                                        (loop for var in free-list
                                                              for i from 0
                                                              do (set! new-body (make-primitive 'peek (list closure (make-constant i)) (list (translate-ref var env*))
                                                                                  new-body)))
                                                        (make-closure
                                                          fn-ptr
                                                          (cons closure (map (lambda (x) (translate-ref x env*)) (closure-arguments fun)))
                                                          new-body))))
            (loop for fun in functions
                  for name = (closure-name fun)
                  if (lexical-ref? name)
                    do (set! env (insert env name (make-variable (gensym)))))
            (let ((new-body (convert body env)))
              ; Build the closures.
              (loop for fun in functions
                    for free-list in frees
                    for ptr in fn-ptrs
                    for closure = (translate (closure-name fun))
                    do (loop for var in free-list
                             for i from 1
                             do (set! new-body (make-primitive 'poke (list (translate var) closure (make-constant i)) '()
                                                 new-body)))
                       (set! new-body (make-primitive 'poke (list ptr closure (make-constant 0)) '()
                                        new-body)))
              ; Allocate the closures.
              (loop for fun in functions
                    for free-list in frees
                    for closure = (translate (closure-name fun))
                    do (set! new-body (make-primitive 'alloc (list (make-constant (+ 1 (length free-list)))) (list closure)
                                        new-body)))
              (make-fix converted-functions new-body)))
          (_ (error "Unexpected form in closure-convert" expr)))))))