(define-module (scmvm language scheme)
  #:use-module (scmvm assembler)
  #:use-module (ice-9 match))

;; Scheme compiler
;; Scheme subset we're targeting
;; <prgm> ::= <top> ...
;; <top>  ::= <def> | <exp>
;; <def>  ::= (define <var> <exp>)
;;          | (define (<var> <var> ...) <exp> ...)
;; <exp>  ::= (lambda (<var> ...) <exp> ...)
;;          | (if <exp> <exp> <exp>)
;;          | (<exp> <exp> ...)
;;          | (let ((<var> <exp>) ...) <exp> ...)
;;          | (begin <exp> ...)
;;          | <num> | <sym> | <var> | #t | #f

(define (atomic? x)
  (or (number? x)
      (symbol? x)
      (boolean? x)))

(define-syntax-rule (define-cps-loop name unit)
  (define (name v* k)
    (if (null? v*)
	(k '())
	(unit (car v*)
	      (lambda (t)
		(name (cdr v*)
		      (lambda (t*)
			(k (cons t t*)))))))))


;; Desugaring
;; Transforms to simplify the language
;; - lambdas and lets can only have 1 expression in body position
;; - define is always simple binds, function defs bind a lambda
;; <prgm> ::= <top> ...
;; <top>  ::= <def> | <exp>
;; <def>  ::= (define <var> <exp>)
;; <exp>  ::= (lambda (<var> ...) <exp>)
;;          | (if <exp> <exp> <exp>)
;;          | (<exp> <exp> ...)
;;          | (let ((<var> <exp>) ...) <exp>)
;;          | (begin <exp> ...)
;;          | <num> | <sym> | <var> | #t | #f

(define (desugar-prgm prgm)
  (map (lambda (top)
	 (if (eq? (car top) 'define)
	     (desugar-define top)
	     (desugar-exp top)))
       prgm))

(define (desugar-define def)
  (match def
    [`(define ,(name params ...) . ,e*)
     `(define ,name ,(desugar-exp `(lambda ,params ,@e*)))]
    [`(define ,name ,exp)
     `(define ,name ,(desugar-exp exp))]))

(define (desugar-exp exp)
  (match exp
    [`(lambda ,params . ,body)
     `(lambda ,params ,(desugar-body body))]
    [`(if ,exp1 ,exp2 ,exp3)
     `(if ,(desugar-exp exp1) ,(desugar-exp exp2) ,(desugar-exp exp3))]
    [`(,f . ,args)
     `(,(desugar-exp f) ,@(map desugar-exp args))]
    [`(let ,((v* e*) ...) . ,body)
     `(let (,(map (lambda (v e) `(,v ,(desugar-exp e))) v* e*))
	,(desugar-body body))]
    [`(begin . ,body) (desugar-body body)]
    [(? atomic?) exp]))

(define (desugar-body body)
  (match body
    ['() '()]
    [(e) (desugar-exp e)]
    [(e* ...) `(begin ,@(map desugar-exp e*))]))


;; CPS conversion
;; Re-structure the program into "Continuation Passing Style", where non-atomic
;; expressions must pass their continuations explicitly, changing to a very
;; "lambda-like" format
;; - begin expressions are decomposed
;; - let expressions are transformed into closed function applications
;; <prgm> ::= <top> ...
;; <top>  ::= <def> | <exp>
;; <def>  ::= (define <var> <exp>)
;; <exp>  ::= <aexp>
;;          | <cexp>
;; <cexp> ::= (<aexp> <aexp> ...)
;;          | (if <aexp> <cexp> <cexp>)
;;          | (set-then! <var> <aexp> <cexp>)
;; <aexp> ::= (lambda (<var> ...) exp)
;;          | <num> | <sym> | <var> | #t | #f
;; 
;; We choose a hybrid transformation based on https://matt.might.net/articles/cps-conversion/

(define undefined-value (make-symbol "undefined"))

(define (cps-convert expr ktail)
  ;; M : expr -> aexp
  ;; T-k : expr, (aexp -> cexp) -> cexp
  ;; T-c : expr, aexp -> cexp
  (define (M expr)
    ;; M dispatches to the appropriate transformer
    (match expr
      [('lambda (var ...) e)
       (let ([$k (gensym "$k")])
	 `(lambda (,@var ,$k) ,(T-c e $k)))]
      [(? atomic?) expr]))

  (define (T-k expr k)
    ;; T-k takes an explicit continuation and calls it when done
    ;; As an invariant, T-k cannot nest a T-c call directly
    (match expr
      [`(lambda . ,_) (k (M expr))]
      [ (? atomic?)   (k (M expr))]
      [ ('begin e) (T-k e k)]
      [ ('begin e e* ...)
        (T-k e (lambda _ (T-k `(begin ,@e*) k)))]
      [ ('let ([v* e*] ...) body)
        (T-k `((lambda (,@v*) ,body) ,@e*) k)]
      [ ('if exp1 exp2 exp3)
        (T-k exp1 (lambda ($exp1)
                    `(if ,$exp1
                         ,(T-k exp2 k)
                         ,(T-k exp3 k))))]
      [ ('set! var expr)
        (T-k expr (lambda ($expr)
                    `(set-then! ,var ,$expr ,(k undefined-value))))]
      [(f e* ...)
       (let* ([$rv (gensym "$rv")]
	      [cont `(lambda (,$rv) ,(k $rv))])
	 (T-k f (lambda ($f)
		  (T*-k e* (lambda ($e*)
			     `(,$f ,@$e* ,cont))))))]))

  (define (T-c expr c)
    ;; T-c takes a symbolic continuation, and uses it to construct CPS
    (match expr
      [`(lambda . ,_) `(,c ,(M expr))]
      [ (? atomic?)   `(,c ,(M expr))]
      [ ('begin e) (T-c e c)]
      [ ('begin e e* ...)
        (T-k e (lambda _ (T-c `(begin ,@e*) c)))]
      [ ('let ([v* e*] ...) body)
        (T-c `((lambda (,@v*) ,body) ,@e*) c)]
      [ ('if exp1 exp2 exp3)
        (let ([$k (gensym "$k")]) ;; Bind cont to avoid blow up
          `((lambda (,$k)
              ,(T-k exp1 (lambda (aexp)
                           `(if ,aexp
                                ,(T-c exp2 $k)
                                ,(T-c exp3 $k)))))
            ,c))]
      [ ('set! var expr)
        (T-k expr (lambda ($expr)
                    `(set-then! ,var ,$expr (,c ,undefined-value))))]
      [ (f e* ...)
        (T-k f (lambda ($f)
		 (T*-k e* (lambda ($e*)
			    `(,$f ,@$e* ,c)))))]))
  (define-cps-loop T*-k T-k)
  (T-k expr ktail))

(define (cps-convert-prgm prgm)
  (if (pair? prgm)
      (cons (cps-convert-top (car prgm))
            (cps-convert-prgm (cdr prgm)))
      '()))

(define (cps-convert-top top)
  (match top
    [`(define ,v ,e)
     (cps-convert e (lambda ($rv) `(define ,v ,$rv)))]
    [_
     (cps-convert top (lambda _ `(nop)))]))