;;;;;; syntax-expander.scm - Syntax Expander module.    -*- Mode: Scheme -*-
;;;;;; Author: Erik Silkensen <silkense@colorado.edu>
;;;;;; Version: 3 Dec 2009

#lang scheme

(require "parser.scm")
(require "environment.scm")
(require "core.scm")

;;; This module provides the syntax expander. The parse tree of a 
;;; program is passed to expand, which returns a parse tree with all 
;;; macros expanded. This module aims to implement the syntax-rules
;;; macro system as described in R5RS.

;;; Before the parse tree of a program is expanded, a ``null-syntax'' is
;;; loaded which defines several of the core builtin macros for Scheme.

;;; There are generally two cases to consider when expanding a piece of
;;; Scheme code: (1) the code is a special form that alters the
;;; syntactic environment; or (2) the code is a macro call. (Note: we
;;; also have to mind the QUOTE special form that is never expanded.)

;;; The special forms that may alter the syntactic environment are
;;; DEFINE-SYNTAX, LET-SYNTAX, LETREC-SYNTAX, and LAMBDA. DEFINE-SYNTAX,
;;; LET-SYNTAX, and LETREC-SYNTAX all bind an identifier to a pair of
;;; the SYNTAX-RULES code defining a macro and the environment in which
;;; it was defined. LAMBDA expansion proceeds by first binding each of
;;; its parameters to fresh identifiers, and then expanding its body in
;;; the extended environment.

;;; A macro call is a list where the first element is an identifier that
;;; has been bound to a SYNTAX-RULES code and a definition environment.
;;; Macro expansion proceeds by first matching the macro call to one of
;;; the patterns specified in the SYNTAX-RULES code. A successful match
;;; yields a substitution mapping pattern variables to Scheme code from
;;; the parse tree. The substitution is used, along with the syntactic
;;; environments, to rewrite the macro call and produce an environment
;;; in which the result is expanded.

(provide expand expand-datum make-global-environment)

(define-syntax aif
  (syntax-rules ()
    ((_ name test true false)
     (let ((name test))
       (if name true false)))
    ((_ name pred test true false)
     (let ((name test))
       (if (pred name) true false)))))

(define (expand ast [dbg #f])
  (set! *gensym-id* 0)
  (let* ((global-env (make-global-environment))
         (exp (expand-all ast global-env)))
    (when dbg (for-each (lambda (x) (pretty-print (datum->syntax x))) exp))
    exp))

(define (expand-all ast env)
  (let loop ((ast ast) (exp-ast '()))
    (if (null? ast)
        (reverse exp-ast)
        (aif exp datum? (expand-datum (car ast) env)
             (loop (cdr ast) (cons exp exp-ast))
             (loop (cdr ast) exp-ast)))))

(define (expand-datum datum env)
  (cond ((define-syntax-datum? datum) (expand-define-syntax datum env))
        ((lambda-datum? datum) (expand-lambda datum env))
        ((quote-datum? datum) datum)
        ((macro-call? datum env)
         (let* ((val (value-datum datum))
                (mac (lookup (value-datum (car val)) env)))
           (expand-macro (car mac) datum (cdr mac) env)))
        ((any-list-datum? datum) (expand-list datum env))
        ((pair-datum? datum) (expand-pair datum env))
        ((symbol-datum? datum) (expand-symbol datum env))
        (else datum)))

(define (expand-define-syntax datum env)
  (let* ((id (value-datum (define-variable datum)))
         (mac (define-value datum))
         (lits (map value-datum (value-datum (syntax-rules-literals mac))))
         (vars (filter (lambda (x) (free? x env)) lits))
         (vals (map (lambda (x) 'literal) vars)))
    (bind! id (cons mac (extend vars vals env)) env)
    'defined))

(define (expand-lambda datum env)
  (let* ((params (lambda-parameters datum))
         (vars (cond ((pair-datum? params)
                      (let ((val (value-datum params)))
                        (list (value-datum (car val))
                              (value-datum (cdr val)))))
                     ((dotted-list-datum? params)
                      (let ((val (flatten-dotted-datum params)))
                        (map value-datum (value-datum val))))
                     ((list-datum? params)
                      (map value-datum (value-datum params)))
                     (else (list (value-datum params)))))
         (vals (map (lambda (id) (make-fresh id env)) vars)))
    (let* ((env (extend vars vals env))
           (params-expanded (expand-datum params env))
           (exp (cons (car-datum datum)
                      (cons params-expanded
                            (let loop ((body (lambda-body datum))
                                       (expanded '()))
                              (if (null? body)
                                  (reverse expanded)
                                  (loop (cdr body)
                                        (cons (expand-datum (car body) env)
                                              expanded))))))))
      (make-list-datum exp (position-datum datum)))))

(define (expand-list datum env)
  (let ((exp (map (lambda (dat)
                    (expand-datum dat env)) 
                  (value-datum datum))))
    (make-datum (name-datum datum) exp (position-datum datum))))

(define (expand-pair datum env)
  (let* ((kar (expand-datum (car-datum datum) env))
         (kdr (expand-datum (cdr-datum datum) env)))
  (make-pair-datum (cons kar kdr) (position-datum datum))))

(define (expand-symbol datum env)
  (aif fid string? (lookup (value-datum datum) env)
       (make-symbol-datum fid (position-datum datum))
       datum))

(define (expand-macro mac datum def-env use-env)
  (let loop ((rules (syntax-rules-clauses mac)))
    (if (null? rules)
        (syntax-error datum "no match")
        (let ((pat (car-datum (car rules)))
              (tem (car-datum (cdr-datum (car rules)))))
          (aif sub (match pat datum def-env use-env)
               (let ((rev-sub (reverse sub))
                     (new-env (make-environment)))
                 (printf "Match result:~n  ")
                 (dbg-sub rev-sub)
                 (rewrite tem rev-sub def-env use-env new-env))
               (loop (cdr rules)))))))

(define (match pat dat def-env use-env)
  (define (match-any pat dat sub)
    (cond ((and (not pat) (not dat)) sub)
          ((and (null-datum? pat) (null-datum? dat)) sub)
          ((null-datum? pat) #f)
          ((symbol-datum? pat) (match-symbol pat dat sub))
          ((pair-datum? pat) (match-pair pat dat sub))
          ((dotted-list-datum? pat) (match-dotted-list pat dat sub))
          ((elided-list-datum? pat) (match-elided-list pat dat sub))
          ((list-datum? pat) (match-list pat dat sub))
          (else (and (equal?-datum pat dat) (unify pat dat sub)))))
  (define (match-symbol pat dat sub)
    (or (and (literal? pat def-env)
             (free? (value-datum dat) use-env)
             (equal?-datum pat dat)
             (unify pat 'literal sub))
        (and (variable? pat use-env)
             (not (literal? pat def-env))
             (unify pat dat sub))))
  (define (match-rec pat dat sub rec)
    (let ((sub (match-any (car-datum pat) (car-datum dat) sub)))
      (and sub (rec (cdr-datum pat) (cdr-datum dat) sub))))
  (define (match-pair pat dat sub)
    (and (or (pair-datum? dat)
             (any-list-datum? dat))
         (not (null-datum? dat))
         (match-rec pat dat sub match-any)))
  (define (match-dotted-list pat dat sub)
    (if (pair-datum? pat)
        (match-pair pat dat sub)
        (match-rec pat dat sub match-dotted-list)))
  (define (match-elided-list pat dat sub)
    (and (any-list-datum? dat)
         (let ((vars (extract-variables pat use-env))
               (sub (if (list2-datum? pat)
                        (let ((p (penult-datum pat)))
                          (let loop ((dat dat) (tub '()))
                            (if (null-datum? dat)
                                (append (amerge (reverse tub)) sub)
                                (let ((tub (match-any p (car-datum dat) tub)))
                                  (and tub (loop (cdr-datum dat) tub))))))
                        (match-rec pat dat sub match-elided-list))))
           (let loop ((vars vars) (sub sub))  ;; insert NULL bindings 
             (if (null? vars)                 ;; for any variable not
                 sub                          ;; used in the <...>
                 (aif var (lambda (x) (subst x sub)) (car vars)
                      (loop (cdr vars) sub)
                      (loop (cdr vars) (unify var '() sub))))))))
  (define (match-list pat dat sub)
    (cond ((not (any-list-datum? dat)) #f)
          ((and (null-datum? pat) (null-datum? dat)) sub)
          ((or (null-datum? pat) (null-datum? dat)) #f)
          (else (match-rec pat dat sub match-list))))
  (let ((chg (lambda (datum)
               (make-datum (name-datum datum)
                           (cons #f (cdr (value-datum datum)))
                           (position-datum datum))))
        (sub (unify (car-datum dat) 'keyword '())))
    (match-any (chg pat) (chg dat) sub)))

(define (rewrite dat sub def-env use-env new-env)
  (define (rewrite-any dat eid)
    (cond ((quote-datum? dat)
           (let ((rew (list (car-datum dat)
                            (substitute (quote-datum dat) sub new-env))))
             (make-list-datum rew (position-datum dat))))
          ((symbol-datum? dat) (rewrite-symbol dat eid))
          ((pair-datum? dat) (rewrite-pair dat eid))
          ((dotted-list-datum? dat) (rewrite-dotted-list dat eid))
          ((list-datum? dat) (rewrite-list dat eid))
          (else dat)))
  (define (rewrite-symbol dat eid)
    (let ((id (subst dat sub)))
      (cond ((datum? id) id)
            ((list? id) (elided-ref id eid))
            ((eq? 'keyword id) dat) ;; see note in R5RS macro docs
            (else
             (let* ((id (value-datum dat))
                    (fid (make-fresh id use-env))
                    (did (or (lookup id def-env) id))
                    (fdat (make-symbol-datum fid (position-datum dat))))
               (save-for-use! id)
               (set! sub (unify dat fdat sub))
               (set! new-env (bind fid did new-env))
               (bind! id (or (lookup id use-env) id) new-env)
               fdat)))))
  (define (rewrite-pair dat eid)
    (let* ((kar (rewrite-any (car-datum dat) eid))
           (kdr (rewrite-any (cdr-datum dat) eid)))
      (cons-datum kar kdr)))
  (define (rewrite-dotted-list dat eid)
    (let rec ((seq (value-datum dat)) (rew '()))
      (cond ((list1? seq)
             (if (pair-datum? (car seq))
                 (let* ((dat (rewrite-pair (car seq) eid))
                        (beg (reverse (cons (car-datum dat) rew)))
                        (rew (append beg (value-datum (cdr-datum dat)))))
                   (make-list-datum rew (position-datum (car rew))))
                 (syntax-error (car seq) "expected dotted list")))
            ((and (list2? seq) (elided-datum? (car-datum (cadr seq))))
             (let* ((tew (rewrite-elided-datum (car seq) eid))
                    (beg (reverse (append tew rew)))
                    (end (rewrite-any (cdr-datum (cadr seq)) eid))
                    (rew (append beg (value-datum end))))
               (make-list-datum rew (position-datum (car rew)))))
            ((elided-datum? (cadr seq))
             (let ((tew (rewrite-elided-datum (car seq) eid)))
               (rec (cddr seq) (append tew rew))))
            (else (rec (cdr seq) (cons (rewrite-any (car seq) eid) rew))))))
  (define (rewrite-list dat eid)
    (let rec ((seq (value-datum dat)) (rew '()))
      (cond ((null? seq)
             (make-list-datum (reverse rew) (position-datum dat)))
            ((or (null? (cdr seq)) (not (elided-datum? (cadr seq))))
             (rec (cdr seq) (cons (rewrite-any (car seq) eid) rew)))
            (else (let ((tew (rewrite-elided-datum (car seq) eid)))
                    (rec (cddr seq) (append tew rew)))))))
  (define (rewrite-elided-datum dat eid-past)
    (let* ((eids (map (lambda (val)
                        (or (and (datum? val) 1)
                            (and (list? val) (length val))
                            (syntax-error dat "bad elision")))
                      (map (lambda (dat)
                             (subst dat sub))
                           (extract-variables dat sub subst))))
           (stop (apply max (cons 0 eids))))
      (let loop ((eid 0) (rew '()))
        (if (= eid stop)
            rew
            (let ((eid-pres (append eid-past (list eid))))
              (loop (+ eid 1) (cons (rewrite-any dat eid-pres) rew)))))))
  (define *saved-ids* '())
  (define (save-for-use! id)
    (unless (member id *saved-ids*)
            (set! *saved-ids* (cons id *saved-ids*))))
  (let* ((rew (rewrite-any dat '()))
         (exp (expand-datum rew new-env)))
    (expand-datum exp (extend *saved-ids* *saved-ids* use-env))))

;; Substitute and rewrite are almost identical. Rewrite uses substitute to
;; rewrite QUOTE forms in macro templates. The difference is that rewrite
;; normally rewrites symbols, maintaining hygiene, while substitute only
;; substitututes the template symbol with its value in the sub (or env).

;; TODO: Can we refactor these two procedures, extracting the different
;; symbol procedures? The only problem seems to be maintaining the envs.

(define (substitute dat sub env)
  (define (substitute-any dat eid)
    (cond ((symbol-datum? dat) (substitute-symbol dat eid))
          ((pair-datum? dat) (substitute-pair dat eid))
          ((dotted-list-datum? dat) (substitute-dotted-list dat eid))
          ((list-datum? dat) (substitute-list dat eid))
          (else dat)))
  (define (substitute-symbol dat eid)
    (let ((id (subst dat sub)))
      (cond ((datum? id) id)
            ((list? id) (elided-ref id eid))
            ((eq? 'keyword id) dat) ;; set note in R5RS macro docs
            (else (let ((id (lookup (value-datum dat) env)))
                    (cond ((string? id)
                           (make-symbol-datum id (position-datum dat)))
                          ((datum? id) id)
                          (else dat)))))))
  (define (substitute-pair dat eid)
    (let* ((kar (substitute-any (car-datum dat) eid))
           (kdr (substitute-any (cdr-datum dat) eid)))
      (cons-datum kar kdr)))
(define (substitute-dotted-list dat eid)
    (let rec ((seq (value-datum dat)) (rew '()))
      (cond ((list1? seq)
             (if (pair-datum? (car seq))
                 (let* ((dat (substitute-pair (car seq) eid))
                        (beg (reverse (cons (car-datum dat) rew)))
                        (rew (append beg (value-datum (cdr-datum dat)))))
                   (make-list-datum rew (position-datum (car rew))))
                 (syntax-error (car seq) "expected dotted list")))
            ((and (list2? seq) (elided-datum? (car-datum (cadr seq))))
             (let* ((tew (substitute-elided-datum (car seq) eid))
                    (beg (reverse (append tew rew)))
                    (end (substitute-any (cdr-datum (cadr seq)) eid))
                    (rew (append beg (value-datum end))))
               (make-list-datum rew (position-datum (car rew)))))
            ((elided-datum? (cadr seq))
             (let ((tew (substitute-elided-datum (car seq) eid)))
               (rec (cddr seq) (append tew rew))))
            (else (rec (cdr seq) (cons (substitute-any (car seq) eid) rew))))))
  (define (substitute-list dat eid)
    (let rec ((seq (value-datum dat)) (rew '()))
      (cond ((null? seq)
             (make-list-datum (reverse rew) (position-datum dat)))
            ((or (null? (cdr seq)) (not (elided-datum? (cadr seq))))
             (rec (cdr seq) (cons (substitute-any (car seq) eid) rew)))
            (else (let ((tew (substitute-elided-datum (car seq) eid)))
                    (rec (cddr seq) (append tew rew)))))))
  (define (substitute-elided-datum dat eid-past)
    (let* ((eids (map (lambda (val)
                        (or (and (datum? val) 1)
                            (and (list? val) (length val))
                            (syntax-error dat "bad elision")))
                      (map (lambda (dat)
                             (subst dat sub))
                           (extract-variables dat sub subst))))
           (stop (apply max (cons 0 eids))))
      (let loop ((eid 0) (rew '()))
        (if (= eid stop)
            rew
            (let ((eid-pres (append eid-past (list eid))))
              (loop (+ eid 1) (cons (substitute-any dat eid-pres) rew)))))))
  (substitute-any dat '()))

(define *gensym-id* 0)
(define (make-fresh id env)
  (let ((sids (lookup* id env)))
    (let loop ((tag *gensym-id*))
      (aif fid (lambda (x) (member x sids)) (format "_gs~a_" tag)
           (loop (+ tag 1))
           (begin (set! *gensym-id* (+ tag 1)) fid)))))

(define (make-global-environment)
  (let ((env (make-environment)))
    (with-input-from-file "null-syntax.scm"
      (lambda () (expand-all (parse) env)))
    env))

(define (literal? datum env)
  (and (symbol-datum? datum)
       (eq? 'literal (lookup (value-datum datum) env))))

(define (variable? datum env)
  (and (symbol-datum? datum)
       (not (eq? 'literal (lookup (value-datum datum) env)))))

(define (extract-variables datum env [variable? variable?])
  (let loop ((dat datum) (vars '()))
    (cond ((and (symbol-datum? dat)
                (not (elided-datum? dat)))
           (if (variable? dat env)
               (cons dat vars)
               vars))
          ((pair-datum? dat)
           (append vars
                   (loop (car-datum dat) '())
                   (loop (cdr-datum dat) '())))
          ((any-list-datum? dat)
           (let ((tars (map (lambda (x) (loop x '())) (value-datum dat))))
             (apply append (cons vars tars))))
          (else vars))))

(define (elided-ref seq path)
  (if (list1? path)
      (list-ref seq (car path))
      (elided-ref (list-ref seq (car path)) (cdr path))))

(define (elided-datum? datum)
  (and (symbol-datum? datum)
       (string=? "..." (value-datum datum))))

(define (elided-list-datum? datum)
  (and (list-datum? datum)
       (let ((val (value-datum datum)))
         (and (> (length val) 1)
              (list1? (filter elided-datum? val))
              (elided-datum? (last val))))))

(define (list1? x) (= 1 (length x)))
(define (list2? x) (= 2 (length x)))
(define (penult x) (list-ref x (- (length x) 2)))
(define (list1-datum? x) (list1? (value-datum x)))
(define (list2-datum? x) (list2? (value-datum x)))
(define (penult-datum x) (penult (value-datum x)))

(define (acons key val alist) (cons (cons key val) alist))
(define (unify x y z) (acons (string->symbol (value-datum x)) y z))
(define (subst x z)
  (let ((y (assq (string->symbol (value-datum x)) z)))
    (if y (cdr y) y)))

(define (amerge alist)
  (define (allq key alist)
    (let loop ((x alist) (y '()) (z '()))
      (if (null? x)
          (cons (reverse y) (reverse z))
          (aif a (lambda (x) (eq? (car x) key)) (car x)
               (loop (cdr x) (cons (cdr a) y) z)
               (loop (cdr x) y (cons a z))))))
  (let loop ((x alist) (y '()))
    (if (null? x)
        (reverse y)
        (let* ((a (car x))
               (as (allq (car a) (cdr x))))
          (if (null? (car as))
              (loop (cdr as) (cons a y))
              (let ((v (cons (cdr a) (car as))))
                (loop (cdr as) (acons (car a) v y))))))))

(define (bound-macro? name env)
  (let ((val (lookup name env)))
    (and (pair? val)
         (syntax-rules-datum? (car val))
         (environment? (cdr val)))))

(define (macro-call? datum env)
  (and (or (pair-datum? datum) (any-list-datum? datum))
       (not (null-datum? datum))
       (symbol-datum? (car-datum datum))
       (bound-macro? (value-datum (car-datum datum)) env)))

(define (dbg-sub sub)
  (define (dbg-rec lis)
    (cond ((datum? lis) (datum->syntax lis))
          ((list? lis) (map dbg-rec lis))
          (else lis)))
  (for-each
   (lambda (x)
     (printf "(~a " (car x))
     (display (dbg-rec (cdr x)))
     (display ") "))
   sub)
  (newline))