; 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 <syntax>
      (syntax e scopes)
      syntax?
      (e syntax-e)
      (scopes syntax-scopes))


    (define *next-scope* 0)


    (define (scope)
      (set! *next-scope* (+ 1 *next-scope*))
      *next-scope*)


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


    (define (set . l)
      (list.foldl
        (lambda (acc x)
          (map.insert acc x #t))
        (map.empty -)
        l))


    (define *core-scopes* (set (scope)))


    ; stamp turns an stree into a syntax expression by stamping all of the
    ; libvars with the core scope.
    (define (stamp tree)
      (let loop ((tree tree))
        (cond
          ((ir.libvar? tree)
            (syntax tree *core-scopes*))
          ((list? tree)
            (map loop tree))
          ((constant? tree) tree)
          (else (error "unexpected form in stamp" tree)))))


    (define (add-scope tree . scopes)
      (cond
        ((syntax? tree)
          (syntax (syntax-e tree) (map.union (syntax-scopes tree) (apply set scopes))))
        ((list? tree)
          (map (lambda (x) (apply add-scope x scopes)) tree))
        ((constant? tree) tree)
        (else (error "unexpected form in add-scope"))))


    (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))


    ; *environment* holds the mappings for expansion.
    ; The type of *environment* is a bit complicated, so it deserves
    ; documentation. It is a map keyed by library variables. The values are
    ; association lists, keyed by scope sets. The value of the inner
    ; association list can be either a macro transformer, represented as a
    ; procedure, or a gensym for a variable. Variables are not resolved yet,
    ; but we put them in the environment so that we can track accurately
    ; whether a symbol resolves to a macro or a variable.
    (define *environment* *empty-libvar-map*)


    (define (best-mapping scopes candidate-mappings)
      (cdr
        (list.foldl
          (lambda (best x)
            (define key (car x))
            (if (and (map.subset? key scopes)
                     (> (map.size key) (map.size (car best))))
              x
              best))
          (cons (set) #f)
          candidate-mappings)))


    (define (lookup stx)
      (define candidate-mappings (map.lookup *environment* (syntax-e stx) #f))
      (and candidate-mappings
           (best-mapping (syntax-scopes stx) candidate-mappings)))


    (define (lookup-macro stx)
      (define mapping (lookup stx))
      (and (not (gensym.gensym? mapping))
           mapping))


    (define (insert stx val)
      (define name (syntax-e stx))
      (set! *environment*
        (map.insert *environment* name
          (cons (cons (syntax-scopes stx) val)
                (map.lookup *environment* name '())))))


    ; expand expands all macros in a syntax expression.
    (define (expand tree)
      (cond
        ((and (pair? tree)
              (syntax? (car tree))
              (lookup-macro (car tree))) => (lambda (transformer)
          (transformer tree)))
        ((null? tree)
          (error "nil by itself is an error"))
        ((list? tree)
          (ir.make-apply
            (expand (car tree))
            (list
              (list.foldr
                (lambda (x acc)
                  (ir.make-call-builtin 'cons
                     (list (expand x) acc)))
                (ir.make-const '())
                (cdr tree)))))
        ((syntax? tree)
          (when (lookup-macro tree)
            (error "macro in wrong context" tree))
          ; Variables are unresolved for now. In order to support
          ; mutually-recursive procedures, or even procedures defined out of
          ; order, we need to wait until all the macros are expanded to
          ; resolve variables.
          tree)
        ((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 (gensym.gensym? expr)
             (syntax? 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 (syntax? expr)
             (gensym.gensym? expr)
             (ir.const? expr)
             (ir.void? expr))
          acc)
        (else (error "unexpected form in fold-ir1" expr))))


    ; unstamp resolves the variables.
    (define (unstamp stx)
      (cond
        ((syntax? stx)
          (or (lookup stx)
              (error "unbound variable" stx)))
        (else
          (map-ir1 unstamp stx))))


    (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)


    (define (builtin-define expr)
      (unless (and (list? expr)
                   (= (length expr) 3)
                   (syntax? (cadr expr)))
        (error "invalid form in builtin-define"))
      (let-values (((var body) (apply values (cdr expr))))
        (insert var (gensym.gen))
        (make-toplevel-define var (expand body))))


    (define (builtin-lambda expr)
      (unless (and (list? expr)
                   (>= (length expr) 3)
                   (syntax? (cadr expr)))
        (error "invalid form in builtin-lambda"))
      (let ((var (cadr expr))
            (body (cddr expr))
            (var-scope (scope))
            (body-scope (scope)))
        (define new-var (add-scope var var-scope))
        (define expanded-body (map expand (add-scope body var-scope body-scope)))
        (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"))
        (insert new-var (gensym.gen))
        (ir.make-lambda (list new-var) (rewrite-body expanded-body))))



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


    (define (gensym=? x y)
      (and (gensym.gensym? x)
           (gensym.gensym? y)
           (= (gensym.gensym->int x) (gensym.gensym->int y))))


    (define (free-identifier=? a b)
      (and (syntax? a)
           (syntax? b)
           (let ((binding-a (lookup a))
                 (binding-b (lookup b)))
             (or (and binding-a
                      binding-b
                      (or (gensym=? binding-a binding-b)
                          ; If not a gensym, then a procedure.
                          (eq? binding-a binding-b)))
                 (and (not binding-a)
                      (not binding-b)
                      (string=? (ir.libvar-var (syntax-e a)) (ir.libvar-var (syntax-e b))))))))


    (define (identifier=? a b)
      (and (syntax? a)
           (syntax? b)
           (string=? (ir.libvar-var (syntax-e a)) (ir.libvar-var (syntax-e b)))))


    (define *empty-string-map* (map.empty cmp-strings))


    ; TODO: handle ellipsis
    (define (pattern-bindings ellipsis literals pattern object)
      (cond
        ((list.any (lambda (lit) (identifier=? pattern lit)) literals)
          (and (free-identifier=? object pattern)
               *empty-string-map*))
        ((and (syntax? pattern)
              (string=? "_" (ir.libvar-var (syntax-e pattern))))
          *empty-string-map*)
        ((syntax? pattern)
          (map.singleton cmp-strings (ir.libvar-var (syntax-e pattern)) object))
        ((null? pattern)
          (and (null? object)
               *empty-string-map*))
        ((pair? pattern)
          (and (pair? object)
            (let ((pbindings (pattern-bindings ellipsis literals (car pattern) (car object)))
                  (p*bindings (pattern-bindings ellipsis literals (cdr pattern) (cdr object))))
              (and pbindings
                   p*bindings
                   (map.union pbindings p*bindings)))))
        ((constant? pattern)
          (and (equal? object pattern)
               *empty-string-map*))
        (else (error "unexpected form in pattern-bindings" pattern))))


    (define (expand-template ellipsis ellipsis-nesting bindings template)
      (cond
        ((null? template)
          '())
        ((pair? template)
          (cons (expand-template ellipsis ellipsis-nesting bindings (car template))
                (expand-template ellipsis ellipsis-nesting bindings (cdr template))))
        ((and (syntax? template)
              (map.lookup bindings (ir.libvar-var (syntax-e template)) #f)) => (lambda (x) x))
        (else template)))


    (define (builtin-syntax-rules expr)
      (error "syntax-rules in wrong context"))


    (define (real-syntax-rules expr)
      (define-values (ellipsis literals rules)
        (cond
          ((and (list? expr)
                (>= (length expr) 3)
                (syntax? (cadr expr))
                (list? (list-ref expr 2))
                (list.all syntax? (list-ref expr 2)))
            (values (ir.libvar-var (syntax-e (cadr expr))) (list-ref expr 2) (list-tail expr 3)))
          ((and (list? expr)
                (>= (length expr) 2)
                (list? (cadr expr))
                (list.all syntax? (cadr expr)))
            (values "..." (cadr expr) (cddr expr)))
          (else (error "invalid form in builtin-syntax-rules"))))
      (unless (list.all (lambda (rule) (and (list? rule)
                                            (= 2 (length rule))))
                        rules)
        (error "invalid form in builtin-syntax-rules"))
      (lambda (arg)
        (call/cc (lambda (return)
          (for-each
            (lambda (rule)
              (define bindings (pattern-bindings ellipsis literals (car rule) arg))
              (when bindings
                (return (expand (expand-template ellipsis (vector) bindings (add-scope (cadr rule) (scope)))))))
            rules)
          (error "builtin-syntax-rules: no match")))))


    (define (builtin-define-syntax expr)
      (unless (and (list? expr)
                   (= (length expr) 3)
                   (syntax? (cadr expr))
                   (pair? (list-ref expr 2))
                   (syntax? (car (list-ref expr 2)))
                   (eq? builtin-syntax-rules (lookup-macro (car (list-ref expr 2)))))
        (error "invalid form in builtin-define"))
      (let-values (((var transformer-spec) (apply values (cdr expr))))
        (define transformer (real-syntax-rules transformer-spec))
        (insert var transformer)
        ir.*void*))


    (define *builtins-environment*
      (list.foldl
        (lambda (acc x)
          (map.insert acc (ir.make-libvar "(csc builtins)" (car x)) (list (cons *core-scopes* (cdr x)))))
        *empty-libvar-map*
        (list (cons "define" builtin-define)
              (cons "define-syntax" builtin-define-syntax)
              (cons "lambda" builtin-lambda)
              (cons "syntax-rules" builtin-syntax-rules)
              (cons "call-builtin" builtin-call-builtin))))


    (define (list-builtins)
      (append
        (map
          (lambda (bindings)
            (map (lambda (x) (string->symbol (ir.libvar-var (syntax-e (cdr x)))))
              (cdr bindings)))
          (map.map->list *builtins-environment*))))


    (define (expand-program prog)
      (set! *environment* *builtins-environment*)
      (unstamp
        (rewrite-body
          (map
            (lambda (x)
              (expand (stamp x)))
            prog))))))