; https://web.cs.ucdavis.edu/~devanbu/teaching/260/kohlbecker.pdf

(define-library (csc macros)
  (export
    expand-program
    list-builtins)
  (import (scheme base)
          (only (scheme case-lambda) case-lambda)
          (prefix (csc gensym) gensym.)
          (prefix (csc ir) ir.)
          (prefix (csc list) list.)
          (prefix (csc map) map.))
  (import (scheme write))
  (begin


    ; Macro expansion is a function on an stree, to a simplified form called
    ; the lambda-term. An stree is one of:
    ;  - symbol
    ;  - constant (null, boolean, char, bytevector, number, string, or vector)
    ;  - list of stree (function call or macro application)
    ;
    ; A lambda-term is one of:
    ;  - symbol
    ;  - constant
    ;  - list of stree (only function call this time)
    ;  - or one of a number of special forms defined in ir.scheme.
    ;
    ; A lambda-term is the same as an stree, except there are no remaining
    ; macro applications. In other words, we are expanding the macros.


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


    (define-record-type <tsvar>
      (make-tsvar var timestamp)
      tsvar?
      (var tsvar-var)
      (timestamp tsvar-ts))


    (define (stamper ts)
      (lambda (var)
        (make-tsvar var ts)))


    (define (constant? x)
      (or (boolean? x)
          (char? x)
          (bytevector? x)
          (number? x)
          (string? x)
          (vector? x)))


    ; timestamp stamps all unstamped symbols.
    (define (timestamp tree stamp)
      (cond
        ((ir.libvar? tree)
          (stamp tree))
        ((list? tree)
          (map (lambda (x) (timestamp x stamp)) tree))
        ((constant? tree) tree)
        (else (error "unexpected form in timestamp" tree))))


    (define (lookup-macro env var)
      (define binding1 (map.lookup env var #f))
      (if (procedure? binding1)
        binding1
        ; Check for var in the global namespace.
        (let ((binding2 (map.lookup env (make-tsvar (tsvar-var var) 0) #f)))
          (and (procedure? binding2)
               binding2))))


    ; expand expands all macros in a timestamped stree.
    (define (expand tree env time)
      (cond
        ((and (pair? tree)
              (tsvar? (car tree))
              (lookup-macro env (car tree))) =>
          (lambda (transformer)
            (transformer tree env time)))
        ((null? tree)
          (error "nil by itself is an error"))
        ((list? tree)
          (ir.make-apply
            (expand (car tree) env time)
            (list
              (list.foldr
                (lambda (x acc)
                  (ir.make-call-builtin 'cons
                    (list
                      (expand x env time)
                      acc)))
                (ir.make-const '())
                (cdr tree)))))
        ((tsvar? tree)
          (or (map.lookup env tree #f)
              tree))  ; Global variables are still unresolved.
        ((constant? tree)
          (ir.make-const tree))
        (else (error "unexpected form in expand" tree))))


    (define (map-ir1 f expr)
      (cond
        ((toplevel-define? expr)
          (make-toplevel-define (f (toplevel-define-var expr)) (f (toplevel-define-body expr))))
        ((ir.lambda? expr)
          (ir.make-lambda (map f (ir.lambda-vars expr)) (f (ir.lambda-body expr))))
        ((ir.letrec? expr)
          (ir.make-letrec (map (lambda (func) (cons (f (car func)) (f (cdr func)))) (ir.letrec-funcs expr)) (f (ir.letrec-body expr))))
        ((ir.apply? expr)
          (ir.make-apply (f (ir.apply-func expr)) (map f (ir.apply-args expr))))
        ((ir.sequence? expr)
          (ir.make-sequence (f (ir.sequence-head expr)) (f (ir.sequence-tail expr))))
        ((ir.set? expr)
          (ir.make-set (f (ir.set-var expr)) (f (ir.set-body expr))))
        ((ir.call-builtin? expr)
          (ir.make-call-builtin (ir.call-builtin-name expr) (map f (ir.call-builtin-args expr))))
        ((or (tsvar? expr)
             (gensym.gensym? expr)
             (ir.const? expr))
          expr)
        ((ir.void? expr) ir.*void*)
        (else (error "unexpected form in map-ir1"))))


    (define (fold-ir1 f acc expr)
      (cond
        ((toplevel-define? expr)
          (f (f acc (toplevel-define-var expr)) (toplevel-define-body expr)))
        ((ir.lambda? expr)
          (f (list.foldl f acc (ir.lambda-vars expr)) (ir.lambda-body expr)))
        ((ir.letrec? expr)
          (f (list.foldl (lambda (acc x) (f (f acc (car x)) (cdr x))) acc (ir.letrec-funcs expr)) (ir.letrec-body expr)))
        ((ir.apply? expr)
          (list.foldl f (f acc (ir.apply-func expr)) (ir.apply-args expr)))
        ((ir.sequence? expr)
          (f (f acc (ir.sequence-head expr)) (ir.sequence-tail expr)))
        ((ir.set? expr)
          (f (f acc (ir.set-var expr)) (ir.set-body expr)))
        ((ir.call-builtin? expr)
          (list.foldl f acc (ir.call-builtin-args expr)))
        ((or (tsvar? expr)
             (gensym.gensym? expr)
             (ir.const? expr)
             (ir.void? expr))
          acc)
        (else (error "unexpected form in fold-ir1" expr))))


    (define (cmp-strings s1 s2)
      (cond
        ((string=? s1 s2)
          0)
        ((string<? s1 s2)
          -1)
        (else 1)))


    (define (cmp-libvar x y)
      (define d (cmp-strings (ir.libvar-lib x)
                             (ir.libvar-lib y)))
      (if (zero? d)
        (cmp-strings (ir.libvar-var x) (ir.libvar-var y))
        d))


    (define *empty-libvar-map* (map.empty cmp-libvar))


    (define (cmp-tsvar v1 v2)
      (define d (- (tsvar-ts v1) (tsvar-ts v2)))
      (if (zero? d)
        (cmp-libvar (tsvar-var v1) (tsvar-var v2))
        d))


    (define *empty-tsvar-map* (map.empty cmp-tsvar))


    (define (expand-body prog env time)
      (let loop ((prog prog)
                 (env env)
                 (time time))
        (if (null? prog)
          '()
          (let ((expanded (expand (car prog) env time)))
            (cond
              ((and (toplevel-define? expanded)
                    (procedure? (toplevel-define-body expanded)))
                (loop (cdr prog)
                  (map.insert env (toplevel-define-var expanded) (toplevel-define-body expanded))
                  (+ 1 time)))
              (else
                (cons expanded
                  (loop (cdr prog) env (+ 1 time)))))))))


    (define (remaining-defines? expr)
      (or (toplevel-define? expr)
        (fold-ir1 (lambda (acc x) (or acc (remaining-defines? x))) #f expr)))


    (define (rewrite-body body)
      (define funcs
        (list.map-maybe
          (lambda (x)
            (and (toplevel-define? x)
                 (ir.lambda? (toplevel-define-body x))
                 (cons (toplevel-define-var x) (toplevel-define-body x))))
          body))
      (define exprs
        (list.filter
          (lambda (x)
            (not (and (toplevel-define? x)
                      (ir.lambda? (toplevel-define-body x)))))
          body))
      (define new-body
        (list.foldr
          (lambda (x acc)
            (if (toplevel-define? x)
              (ir.make-apply (ir.make-lambda (list (toplevel-define-var x)) acc) (list ir.*void*))
              acc))
          (ir.make-letrec funcs
            (list.foldl
              (lambda (acc x)
                (ir.make-sequence acc
                  (if (toplevel-define? x)
                    (ir.make-set (toplevel-define-var x) (toplevel-define-body x))
                    x)))
              ir.*void*
              exprs))
          exprs))
      (when (remaining-defines? new-body)
        (error "definition in expression context"))
      new-body)


    ; remaining-symbols returns the timestamped symbols in tree.
    (define (remaining-symbols tree)
      (cond
        ((tsvar? tree)
          ; These free variables should be unstamped, since they refer to
          ; toplevel bindings.
          (map.singleton cmp-strings (tsvar-var tree) (gensym.gen)))
        (else
          (fold-ir1
            (lambda (acc x)
              (map.union acc (remaining-symbols x)))
            *empty-libvar-map*
            tree))))


    (define (unstamp-tree tree symbols)
      (cond
        ((tsvar? tree)
          (map.lookup symbols (tsvar-var tree)))
        (else (map-ir1 (lambda (x) (unstamp-tree x symbols)) tree))))


    ; unstamp changes the timestamped symbols in expr into gensyms.
    (define (unstamp expr)
      (unstamp-tree expr (remaining-symbols expr)))


    (define (builtin-define expr env time)
      (unless (and (list? expr)
                   (= (length expr) 3)
                   (tsvar? (cadr expr)))
        (error "invalid form in builtin-define"))
      (let-values (((var body) (apply values (cdr expr))))
        (make-toplevel-define var (expand body env (+ 1 time)))))


    (define (builtin-lambda expr env time)
      (unless (and (list? expr)
                   (>= (length expr) 3)
                   (tsvar? (cadr expr)))
        (error "invalid form in builtin-lambda"))
      (let ((var (cadr expr))
            (body (cddr expr))
            (new-var (gensym.gen)))
        (define env* (map.insert env var new-var))
        (define expanded-body (expand-body body env* (+ 1 time)))
        (define exprs
          (let loop ((body expanded-body))
            (if (or (null? body)
                    (not (toplevel-define? (car body))))
              body
              (loop (cdr body)))))
        (when (list.any remaining-defines? exprs)
          (error "out of order define"))
        (ir.make-lambda (list new-var) (rewrite-body expanded-body))))


    (define (builtin-call-builtin expr env time)
      (unless (and (list? expr)
                   (>= (length expr) 2)
                   (tsvar? (cadr expr)))
        (error "invalid form in builtin-call-builtin"))
      (let ((builtin-name (tsvar-var (cadr expr)))
            (args (cddr expr)))
        (ir.make-call-builtin (string->symbol (ir.libvar-var builtin-name))
          (map
            (lambda (arg)
              (expand arg env (+ 1 time)))
            args))))


    (define *builtins-environment*
      (list.foldl
        (lambda (acc x)
          (map.insert acc (make-tsvar (ir.make-libvar "(csc builtins)" (car x)) 0) (cdr x)))
        *empty-tsvar-map*
        (list (cons "define" builtin-define)
              (cons "lambda" builtin-lambda)
              (cons "call-builtin" builtin-call-builtin))))
    ; Damn I can't believe we didn't actually implement syntax-rules.


    (define (list-builtins)
      (map
        (lambda (x)
          (string->symbol (ir.libvar-var (tsvar-var (car x)))))
        (map.map->list *builtins-environment*)))


    (define (expand-program prog)
      (unstamp
        (rewrite-body
          (expand-body
            (map
              (lambda (x)
                (timestamp x (stamper 0)))
              prog)
            *builtins-environment*
            1))))))