(define-library (csc cps)
  (export ir->cps)
  (import (scheme base)
          (prefix (csc gensym) gensym.)
          (prefix (csc ir) ir.)
          (prefix (csc list) list.)
          (prefix (csc map) map.))
  (begin


    (define-record-type <update>
      (make-update var body k)
      update?
      (var update-var)
      (body update-body)
      (k update-k))


    (define (returns-value? op) (not (memq op '(poke exit))))


    (define (to-cps expr k)
      (cond
        ((or (ir.const? expr)
             (ir.void? expr)
             (gensym.gensym? expr))
          (k expr))
        ((ir.lambda? expr)
          (unless (= (length (ir.lambda-vars expr)) 1)
            (error "invalid lambda form: all lambdas at this point should take one argument"))
          (let ((arg (car (ir.lambda-vars expr)))
                (f (gensym.gen))
                (cont (gensym.gen)))
            (ir.make-letrec
              (list
                (cons f
                  (ir.make-lambda (list cont arg)
                    (to-cps (ir.lambda-body expr)
                      (lambda (ret)
                        (ir.make-apply cont (list ret)))))))
              (k f))))
        ((ir.letrec? expr)
          (ir.make-letrec
            (map
              (lambda (func)
                (define fname (car func))
                (define args (ir.lambda-vars (cdr func)))
                (define arg (if (= (length args) 1)
                              (car args)
                              (error "invalid lambda form: blah blah blah")))
                (define body (ir.lambda-body (cdr func)))
                (define cont (gensym.gen))
                (cons fname
                  (ir.make-lambda (list cont arg)
                    (to-cps body
                      (lambda (ret)
                        (ir.make-apply cont (list ret)))))))
              (ir.letrec-funcs expr))
            (to-cps (ir.letrec-body expr) k)))
        ((ir.apply? expr)
          (unless (= (length (ir.apply-args expr)) 1)
            (error "invalid apply form: all function applications at this point should have one argument"))
          (let ((arg (car (ir.apply-args expr)))
                (return-address (gensym.gen))
                (result (gensym.gen)))
            (ir.make-letrec
              (list (cons return-address (ir.make-lambda (list result) (k result))))
              (to-cps (ir.apply-func expr)
                (lambda (f)
                  (to-cps arg
                    (lambda (v)
                      (ir.make-apply f (list return-address v)))))))))
        ((ir.sequence? expr)
          (to-cps (ir.sequence-head expr)
            (lambda (ignored)
              (to-cps (ir.sequence-tail expr) k))))
        ((ir.set? expr)
          (to-cps (ir.set-body expr)
            (lambda (val)
              ; We're compiling to an intermediate form <update>,
              ; which will get removed later by convert-boxes.
              (make-update (ir.set-var expr) val (k ir.*void*)))))
        ((ir.call-builtin? expr)
          (let loop ((args (ir.call-builtin-args expr))
                     (args* '()))
            (if (null? args)
              (if (returns-value? (ir.call-builtin-name expr))
                (let ((result (gensym.gen)))
                  (ir.make-primop (ir.call-builtin-name expr) (reverse args*) (list result) (list (k result))))
                (ir.make-primop (ir.call-builtin-name expr) (reverse args*) '() (list (k ir.*void*))))
              (to-cps (car args)
                (lambda (arg)
                  (loop (cdr args) (cons arg args*)))))))
        (else (error "invalid form in to-cps" expr))))


    (define (fold-ir2 f acc expr)
      (cond
        ((or (ir.const? expr)
             (ir.void? expr)
             (gensym.gensym? expr))
          acc)
        ((ir.apply? expr)
          (list.foldl f acc (cons (ir.apply-func expr) (ir.apply-args expr))))
        ((ir.letrec? expr)
          (f
            (list.foldl
              (lambda (acc x)
                (list.foldl f
                  (f (f acc (car x)) (ir.lambda-body (cdr x)))
                  (ir.lambda-vars (cdr x))))
              acc
              (ir.letrec-funcs expr))
            (ir.letrec-body expr)))
        ((ir.primop? expr)
          (list.foldl f
            (list.foldl f
              (list.foldl f acc (ir.primop-ks expr))
              (ir.primop-vals expr))
            (ir.primop-args expr)))
        (else (error "invalid form in fold-ir2" expr))))


    (define (cmp-gensyms x y)
      (- (gensym.gensym->int x) (gensym.gensym->int y)))


    (define *empty-gensym-map* (map.empty cmp-gensyms))


    (define (get-boxed expr)
      (cond
        ((update? expr)
          (map.insert (get-boxed (update-k expr))
            (update-var expr) #t))
        (else
          (fold-ir2
            (lambda (acc x) (map.union acc (get-boxed x)))
            *empty-gensym-map*
            expr))))


    (define (unbox expr boxed-vars)
      (define (boxed? var)
        (and (gensym.gensym? var)
             (map.lookup boxed-vars var #f)))
      (cond
        ((update? expr)
          (ir.make-primop 'poke (list (update-var expr) (ir.make-const 0) (update-body expr)) '()
            (list (unbox (update-k expr) boxed-vars))))
        ((ir.apply? expr)
          (let loop ((args (ir.apply-args expr))
                     (args* '()))
            (if (null? args)
              (if (boxed? (ir.apply-func expr))
                (let ((temp (gensym.gen)))
                  (ir.make-primop 'peek (list (ir.apply-func expr) (ir.make-const 0)) (list temp)
                    (ir.make-apply temp (reverse args*))))
                (ir.make-apply (ir.apply-func expr) (reverse args*)))
              (let ((arg (car args)))
                (if (boxed? arg)
                  (let ((temp (gensym.gen)))
                    (ir.make-primop 'peek (list arg (ir.make-const 0)) (list temp)
                      (loop (cdr args) (cons temp args*))))
                  (loop (cdr args) (cons arg args*)))))))
        ((ir.letrec? expr)
          (let loop ((funcs (ir.letrec-funcs expr))
                     (funcs* '())
                     (name-translations '()))
            (if (null? funcs)
              (list.foldl
                (lambda (acc x)
                  (ir.make-primop 'poke (list (car x) (ir.make-const 0) (cdr x)) '()
                    (list acc)))
                (unbox (ir.letrec-body expr) boxed-vars)
                name-translations)
              (let ((name (caar funcs))
                    (new-func
                      (let loop ((args (ir.lambda-vars (cdar funcs)))
                                 (args* '())
                                 (arg-mapping '()))
                        (if (null? args)
                          (ir.make-lambda (reverse args*)
                            (list.foldl
                              (lambda (acc x)
                                (ir.make-primop 'alloc (list (ir.make-const 1)) (list (car x))
                                  (list
                                    (ir.make-primop 'poke (list (car x) (ir.make-const 0) (cdr x)) '()
                                      (list acc)))))
                              (unbox (ir.lambda-body (cdar funcs)) boxed-vars)
                              arg-mapping))
                          (let ((arg (car args)))
                            (if (boxed? arg)
                              (let ((temp (gensym.gen)))
                                (loop (cdr args)
                                  (cons temp args*)
                                  (cons (cons arg temp) arg-mapping)))
                              (loop (cdr args) (cons arg args*) arg-mapping)))))))
                (if (boxed? name)
                  (let ((new-name (gensym.gen)))
                    (ir.make-primop 'alloc (list (ir.make-const 1)) (list name)
                      (list
                        (loop (cdr funcs)
                          (cons (cons new-name new-func) funcs*)
                          (cons (cons name new-name) name-translations)))))
                  (loop (cdr funcs)
                    (cons (cons name new-func) funcs*)
                    name-translations))))))
        ((ir.primop? expr)
          (let loop ((args (ir.primop-args expr))
                     (args* '()))
            (if (null? args)
              (ir.make-primop (ir.primop-name expr) (reverse args*) (ir.primop-vals expr)
                (map (lambda (x) (unbox x boxed-vars)) (ir.primop-ks expr)))
              (if (boxed? (car args))
                (let ((temp (gensym.gen)))
                  (ir.make-primop 'peek (list (car args) (ir.make-const 0)) (list temp)
                    (list (loop (cdr args) (cons temp args*)))))
                (loop (cdr args) (cons (car args) args*))))))
        ((or (ir.const? expr)
             (ir.void? expr))
          expr)))


    (define (convert-boxes expr)
      (unbox expr (get-boxed expr)))


    (define (free-vars expr bound-vars)
      (cond
        ((gensym.gensym? expr)
          (if (map.lookup bound-vars expr #f)
            *empty-gensym-map*
            (map.singleton cmp-gensyms expr #t)))
        ((ir.letrec? expr)
          (map.union
            (list.foldl
              (lambda (acc x)
                (map.union acc
                  (free-vars-fun x bound-vars)))
              *empty-gensym-map*
              (ir.letrec-funcs expr))
            (free-vars (ir.letrec-body expr) bound-vars)))
        (else
          (fold-ir2
            (lambda (acc x)
              (map.union acc (free-vars x bound-vars)))
            *empty-gensym-map*
            expr))))


    (define (free-vars-fun fixfun bound-vars)
      (define name (car fixfun))
      (define fun (cdr fixfun))
      (define bound-vars*
        (list.foldl
          (lambda (acc x)
            (map.insert acc x #t))
          (map.singleton cmp-gensyms name #t)
          (ir.lambda-vars fun)))
      (free-vars (ir.lambda-body fun) bound-vars*))


    (define (convert-closures expr env)
      (define (translate var)
        (map.lookup env var var))
      (cond
        ((ir.apply? expr)
          (let ((p (translate (ir.apply-func expr)))
                (temp (gensym.gen)))
            (ir.make-primop 'peek (list p (ir.make-const 0)) (list temp)
              (list
                (ir.make-apply temp
                  (cons p (map translate (ir.apply-args expr))))))))
        ((ir.letrec? expr)
          (let* ((functions (ir.letrec-funcs expr))
                 (frees (map (lambda (x) (free-vars-fun x *empty-gensym-map*)) functions))
                 (fn-ptrs (map (lambda (x) (ir.make-label (gensym.gen))) functions))
                 (new-funcs
                   (map
                     (lambda (fixfun free-vars new-name)
                       (define name (car fixfun))
                       (define fun (cdr fixfun))
                       (define env*
                         (list.foldl
                           (lambda (acc x) (map.insert acc x (gensym.gen)))
                           *empty-gensym-map*
                           free-vars))
                       (define closure (gensym.gen))
                       (cons new-name
                         (ir.make-lambda (map (lambda (x) (map.lookup env* x x)) (ir.lambda-vars fun))
                           (let loop ((free-vars free-vars)
                                      (i 1))
                             (if (null? free-vars)
                               (convert-closures (ir.lambda-body fun) env*)
                               (ir.make-primop 'peek (list closure (ir.make-const i)) (list (map.lookup env* (car free-vars) (car free-vars)))
                                 (list (loop (cdr free-vars) (+ 1 i)))))))))
                     functions frees fn-ptrs))
                 (new-body
                   (list.foldr
                     (lambda (func new-name free-vars acc)
                       (define closure (car func))
                       (ir.make-primop 'alloc (list (ir.make-const (+ 1 (length free-vars)))) (list closure)
                         (list
                           (ir.make-primop 'poke (list closure (ir.make-const 0) new-name) '()
                             (list
                               (let loop ((free-vars free-vars)
                                          (i 1))
                                 (if (null? free-vars)
                                   acc
                                   (ir.make-primop 'poke (list closure (ir.make-const i) (car free-vars)) '()
                                     (list (loop (cdr free-vars) (+ i 1)))))))))))
                     (convert-closures (ir.letrec-body expr) env)
                     functions fn-ptrs frees)))
            (ir.make-letrec new-funcs new-body)))
        ((ir.primop? expr)
          (ir.make-primop (ir.primop-name expr)
            (map translate (ir.primop-args expr))
            (ir.primop-vals expr)
            (map (lambda (x) (convert-closures x env)) (ir.primop-ks expr))))
        ((or (ir.const? expr)
             (ir.void? expr))
          expr)
        (else (error "unexpected form in convert-closures"))))


    (define (hoist expr)
      (define functions '())
      (define body
        (let hoist ((expr expr))
          (cond
            ((ir.letrec? expr)
              (set! functions (append (map hoist (ir.letrec-funcs expr)) functions))
              (hoist (ir.letrec-body expr)))
            ((ir.primop? expr)
              (ir.make-primop (ir.primop-name expr) (ir.primop-args expr) (ir.primop-vals expr)
                (map hoist (ir.primop-ks expr))))
            ((or (ir.apply? expr)
                 (ir.const? expr)
                 (ir.void? expr))
              expr)
            (else (error "unexpected form in hoist")))))
      (ir.make-letrec functions body))


    (define (tail x)
      (ir.make-primop 'exit (list (ir.make-const 0)) '() '()))


    (define (ir->cps expr)
      (hoist
        (convert-closures
          (convert-boxes
            (to-cps expr tail))
          *empty-gensym-map*)))))