(define-library (csc cps)
  (export
    ir1->ir2)
  (import (scheme base)
          (only (csc gensym)
            gensym
            gensym->int)
          (only (csc hash-map)
            insert
            key-not-found-error?
            lookup
            make-map
            merge)
          (only (csc ir1)
            %call
            %define-syntax
            %if
            %lambda
            %letrec
            %lexical-set
            %library-define
            %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-constant
            make-lambda
            make-lexical-ref
            make-lexical-set
            make-sequence
            sequence?)
          (only (csc ir2)
            %apply
            %branch
            %fix
            %primitive
            closure-arguments
            closure-body
            closure-name
            closure-rest
            make-apply
            make-atom
            make-branch
            make-call-closure
            make-closure
            make-fix
            make-kargs
            make-klabel
            make-ktail
            make-primitive)
          (only (csc loop)
            loop
            return)
          (only (csc match)
            define-match-record-type
            match))
  (begin


    ; 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 (new-ref)
      (make-lexical-ref 'generated-symbol (gensym)))


    (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 rest body)
                            (define continuation (new-ref))
                            (make-closure
                              (make-lexical-ref name gensym)
                              (cons continuation args)
                              rest
                              (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)))))


    (define (to-cps expr continuation)
      (match expr
        (_ (when (or (constant? expr)
                     (lexical-ref? expr)
                     (library-ref? expr)))
          (continuation expr))
        ((% %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) #f
                        (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 args)
          (define return-address (new-ref))
          (define result (new-ref))
          (make-fix
            (list (make-closure return-address (list result) #f (continuation result)))
            (to-cps
              proc
              (lambda (f)
                ; Technically the order of evaluation is unspecified.
                ; We evaluate expressions left to right.
                ;
                ; I would use the loop macro, but it mutates the loop
                ; variables which plays badly with building a lambda.
                (let loop ((args* (reverse args))
                           (exprs (lambda (vals)
                                    (make-apply f (cons return-address (reverse vals))))))
                  (if (null? args*)
                    (exprs '())
                    (loop (cdr args*)
                          (lambda (vals)
                            (to-cps
                              (car args*)
                              (lambda (val)
                                (exprs (cons val vals))))))))))))
        ((% %sequence head tail)
          (to-cps
            head
            (lambda (x)
              (to-cps
                tail
                continuation))))
        ((% %lambda args rest body)
          (define f (new-ref))
          (define k (new-ref))
          (make-fix
            (list
              (make-closure f (cons k args) rest
                (to-cps
                  body
                  (lambda (ret)
                    (make-apply k (list ret))))))
            (continuation f)))
        ((% %letrec in-order? _ _ _ body)
          (define-values (functions variable-names variable-values) (collect-functions-and-variables expr))
          (make-fix functions
            (to-cps
              ; We re-write a letrec into a corresponding lambda form.
              (if in-order?
                (loop for name in (reverse variable-names)
                      for value in (reverse variable-values)
                      for expr = (make-call
                                   (make-lambda
                                     (list name)
                                     #f
                                     body)
                                   (list value))
                               then (make-call
                                      (make-lambda
                                        (list name)
                                        #f
                                        expr)
                                      (list value))
                      finally (return expr))
                (make-call
                  (make-lambda
                    variable-names
                    #f
                    body)
                  variable-values))
              continuation)))
        (_ (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"))))


    (define (all-closure-args fun)
      (define args (closure-arguments fun))
      (define rest (closure-rest fun))
      (when rest
        (set! args (cons rest args)))
      args)


    ; Rewrites the given expression to have no more <update> forms.
    (define (box-conversion expr)
      (define boxed-refs (get-boxed expr))
      (define (boxed? ref)
        (or (library-ref? ref) ; globals are always boxed
            (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)
            (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)
            (if (boxed? atom)
              (let ((var (new-ref)))
                (make-primitive 'peek (list atom (make-constant 0)) (list var)
                                (make-branch var (convert true) (convert 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
                                   for rest = (closure-rest fun)
                                   collect (let-values (((new-args boxed-args temp-args) (convert-arg-list (all-closure-args fun))))
                                             (make-closure
                                               new-name
                                               (if rest
                                                 (cdr new-args)
                                                 new-args)
                                               (if rest
                                                 (car new-args)
                                                 #f)
                                               (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 expr))))))


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