#lang racket

; Input Language (simplified cps):
; e ::= (let ([x 'dat]) e)
;     | (let ([x (prim op x ...)]) e)
;     | (let ([x (lambda (x ...) e)]) e)
;     | (if x e e)
;     | (x x ...)

; Output Language (procedural):
; p ::= ((proc (x x ...) e) ...)
; e ::= (let ([x 'dat]) e)
;     | (let ([x (prim op x ...)]) e)
;     | (let ([x (make-closure x x ...)]) e)
;     | (let ([x (env-ref x nat)]) e)
;     | (if x e e)
;     | (clo-app x x ...)

; Interpreter for output language:
(define (racket-proc-eval procs)
  (with-handlers ([exn:fail? (lambda (x) x)])
                 (parameterize ([current-namespace (make-base-namespace)])
                               (namespace-require 'rnrs)
                               (namespace-require 'racket)
                               (namespace-require 'srfi/34)
                               (eval (compile
                                      `(call/cc (lambda (halt)
                                                  (define (prim op . args) (apply op args))
                                                  (define make-closure vector)
                                                  (define env-ref vector-ref)
                                                  (define (clo-app f . vs) (apply (vector-ref f 0) (cons f vs)))
                                                  ,@(map (match-lambda [`(proc (,xs ...) ,bdy) `(define ,xs ,bdy)]) procs)
                                                  (main))))))))

; (f a b c) -> (clo-app f a b c) -> (C-style-app f[0] f a b c)



(define (closure-convert simple-cps)
  ; Exp x List[Proc] -> Exp x Set[Var] x List[Proc]
  (define (bottom-up e procs)
    (match e
      [`(let ([,x ',dat]) ,e0)
       (match-define `(,e0+ ,free+ ,procs+)
                     (bottom-up e0 procs))
       `((let ([,x ',dat]) ,e0+)
         ,(set-remove free+ x)
         ,procs+)]
      [`(let ([,x (prim ,op ,xs ...)]) ,e0)
       (match-define `(,e0+ ,free+ ,procs+)
                     (bottom-up e0 procs))
       `((let ([,x (prim ,op ,@xs)]) ,e0+)
         ,(set-remove (set-union free+ (list->set xs)) x)
         ,procs+)]
      [`(let ([,x (lambda (,xs ...) ,body)]) ,e0)
       (match-define `(,e0+ ,free0+ ,procs0+)
                     (bottom-up e0 procs))
       (match-define `(,body+ ,freelam+ ,procs1+)
                     (bottom-up body procs0+))
       (define env-vars (foldl (lambda (x fr) (set-remove fr x))
                               freelam+
                               xs))
       (define ordered-env-vars (set->list env-vars))
       (define lamx (gensym 'lam))
       (define envx (gensym 'env))
       (define body++ (cdr (foldl (lambda (x count+body)
                                    (match-define (cons cnt bdy) count+body)
                                     (cons (+ 1 cnt)
                                           `(let ([,x (env-ref ,envx ,cnt)])
                                              ,bdy)))
                                  (cons 1 body+)
                                  ordered-env-vars)))
       `((let ([,x (make-closure ,lamx ,@ordered-env-vars)]) ,e0+)
         ,(set-remove (set-union free0+ env-vars) x)
         ((proc (,lamx ,envx ,@xs) ,body++) . ,procs1+))]
      [`(if ,(? symbol? x) ,e0 ,e1)
       (match-define `(,e0+ ,free0+ ,procs0+)
                     (bottom-up e0 procs))
       (match-define `(,e1+ ,free1+ ,procs1+)
                     (bottom-up e1 procs0+))
       `((if ,x ,e0+ ,e1+)
         ,(set-union free1+ free0+ (set x))
         ,procs1+)]
      [`(,(? symbol? xs) ...)
       `((clo-app ,@xs)
         ,(list->set xs)
         ,procs)]))
  (match-define `(,main-body ,free ,procs) (bottom-up simple-cps '()))
  `((proc (main) ,main-body) . ,procs))


 
(define test0 '(let ([a '3])
                 (let ([f (lambda (b) (let ([c (prim + a b)])
                                        (let ([_ (prim halt c)])
                                          (_ _))))])
                   (f a))))

; Run test: (racket-proc-eval (closure-convert test0))



