(define-module (scmvm language scheme)
  #:use-module (scmvm assembler)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-26)
  #:use-module (ice-9 match)
  #:use-module ((rnrs base)
                #:version (6)
                #:select (assert))
  #:export (desugar-prgm
            cps-convert-prgm
            ir-convert))

;; 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 (and (pair? top) (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>)
;;          | (define-then! <var> <aexp> <cexp>)
;; <aexp> ::= (lambda (<var> ...) exp)
;;          | <num> | <var> | #t | #f
;; 
;; We choose a hybrid transformation based on https://matt.might.net/articles/cps-conversion/
;; Admittedly this is a little black magic to me, but it's useful

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

(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))]
    [ ('define v e) (T-k `(define-then! ,v ,e) k)]
    [ ('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))]
    [ ('define v e) (T-c `(define-then! ,v ,e) c)]
    [ ('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-convert-prgm prgm)
  (T-c `(begin ,@prgm) 'ktail))

(define-cps-loop T*-k T-k)

(define (ir-convert prgm)
  (cps-convert-prgm (desugar-prgm prgm)))

 ;; Denotation

(define *globals* (make-parameter '()))
(define *predefined* (make-parameter '()))

(define (meaning e r)
  (match e
    [(? constant?) (meaning-constant e)]
    [(? symbol?) (meaning-reference e r)]
    [('lambda (vars ...) body)
     (meaning-abstraction vars body r)]
    [('set-then! var e k)
     (meaning-assignment var e k r)]
    [('define-then! var e k)
     (meaning-definition var e k r)]
    [('if e k1 k2)
     (meaning-alternative e k1 k2 r)]
    [(f e* ... k)
     (meaning-application f e* k r)]))

(define (meaning-reference v r)
  (match (locate-variable v r)
    [`(local ,i)      (+local-reference+ i)]
    [`(global ,i)     (+global-reference+ i)]
    [`(predefined ,p) (+predefined+ p)]
    [_ (static-error "Reference to undefined variable" v)]))

(define (meaning-constant c)
  (+constant+ c))

(define (meaning-abstraction vars body r)
  (let* ([arity (length vars)]
         [r0 (r-extend r vars)]
         [m+ (meaning body r0)])
    (append-meanings
     (+fix-closure+ arity)
     m+
     (+return+ arity))))

(define (meaning-assignment var e k r)
  (let* ([m (meaning e r)]
         [ma (match (locate-variable var r)
               [`(local ,i) (+set!-local+ i)]
               [`(global ,i) (+set!-global+ i)]
               [`(predefined ,_) (static-error "Assignment to predefined variable" var)] 
               [_ (static-error "Assignment to undefined variable" var)])]
        [mk (meaning k r)])
    (append-meanings m ma mk)))

(define (meaning-definition var e k r)
  (let* ([m (meaning e r)]
         [mv (match (locate-variable var r)
               [`(local ,_) (static-error "Definition conflicts local variable" var)]
               [`(global ,_) (static-error "Redefinition of global variable" var)]
               [`(predefined ,_) (static-error "Redefinition of predefined" var)]
               [#f (+global-definition+ var)])]
         [mk (meaning k r)])
    (append-meanings m mv mk)))

(define (meaning-alternative e k1 k2 r)
  (let* ([jump-false-label (gensym "jump-false-")]
         [endif-label (gensym "endif-")]
         [m (meaning e r)]
         [mk1 (meaning k1 r)]
         [mk2 (meaning k2 r)])
    (append-meanings
     m
     (+branch+ jump-false-label) mk1 (+goto+ endif-label)
     (+label+ jump-false-label) mk2 (+label+ endif-label))))

(define (meaning-application f e* k r)
  (let* ([arity (length e*)]
         [mf (meaning f r)]
         [m* (meaning* e* r)]
         [mk (meaning k r)])
    (append-meanings
     mf 
     m* (+frame-allocate+ arity)
     (+frame-push+ arity) (+function-invoke+) (+frame-pop+ arity)
     mk)))

(define (meaning* e* r)
  (if (pair? e*)
      (let ([m (meaning (car e*) r)]
            [m* (meaning* (cdr e*) r)])
        (append-meanings m m*))
      '()))

(define (locate-variable v r)
  (cond
   [(list-index (lambda (v0) (eq? v v0)) r) =>
    (lambda (i) `(local ,i))]
   [(list-index (lambda (v0) (eq? v v0)) (*globals*)) =>
    (lambda (i) `(global ,i))]
   [(list-index (lambda (v0) (eq? v v0)) (*predefined*)) =>
    (lambda (p) `(predefined ,p))]
   [else #f]))

(define (constant? x)
  (or (number? x)
      (boolean? x)
      (and (pair? x) (eq? 'quote (car x)))))

(define (drop-environment vars r)
  (let ([n (length vars)])
    (assert (equal? vars (take n r)))
    (drop n r)))

(define append-meanings append)

(define (r-extend r vars)
  (append vars r))

(define (global-extend! vars)
  (*globals* (append vars (*globals*))))

(define (static-error . args)
  `((+error+ ,@args)))

(define-syntax define-combinator
  (syntax-rules ()
    [(_ (name args ...))
     (define (name args ...)
       `((name ,@(list args ...))))]))

(define-combinator (+predefined+ i))
(define-combinator (+global-reference+ i))
(define-combinator (+local-reference+ i))
(define-combinator (+constant+ c))
(define-combinator (+fix-closure+ arity))
(define-combinator (+return+ arity))
(define-combinator (+set!-global+ i))
(define-combinator (+set!-local+ i))
(define-combinator (+alternative+ i))
(define-combinator (+global-definition+ v))
(define-combinator (+branch+ label))
(define-combinator (+goto+ label))
(define-combinator (+label+ name))
(define-combinator (+frame-allocate+ size))
(define-combinator (+frame-push+ size))
(define-combinator (+function-invoke+))
(define-combinator (+frame-pop+ size))

(define (std-predefined)
  '(cons car cdr eq? pair? null? symbol? = + - * /))

(define* (scheme-compile prgm #:key (globals '()))
  (parameterize ([*globals* globals]
                 [*predefined* (std-predefined)])
    (meaning (ir-convert prgm) '())))