#lang scheme/base
(require (for-template scheme/base)
         syntax/kerncase
         mzlib/list
         syntax/toplevel
         mzlib/plt-match
         syntax/stx
         "util.ss")
(provide free-vars)  

;; free-vars: syntax -> (listof identifier)
;; Find the free variables in an expression
(define (free-vars stx)  
  (kernel-syntax-case 
      stx (transformer?)
    [(begin be ...)
     (free-vars* (syntax->list #'(be ...)))]
    [(begin0 be ...)
     (free-vars* (syntax->list #'(be ...)))]
    [(set! v ve)
     (union (free-vars #'v)
            (free-vars #'ve))]
    [(let-values ([(v ...) ve] ...) be ...)
     (union (free-vars* (syntax->list #'(ve ...)))
            (set-diff (free-vars* (syntax->list #'(be ...)))
                      (apply append (map syntax->list (syntax->list #'((v ...) ...))))))]
    [(letrec-values ([(v ...) ve] ...) be ...)
     (set-diff (union (free-vars* (syntax->list #'(ve ...)))
                      (free-vars* (syntax->list #'(be ...))))
               (apply append (map syntax->list (syntax->list #'((v ...) ...)))))]
    [(#%plain-lambda formals be ...)
     (set-diff (free-vars* (syntax->list #'(be ...)))
               (formals-list #'formals))]
    [(case-lambda [formals be ...] ...)
     (apply union*
            (map (lambda (fs bes)
                   (set-diff (free-vars* (syntax->list bes))
                             (formals-list fs)))
                 (syntax->list #'(formals ...))
                 (syntax->list #'((be ...) ...))))]
    [(if te ce ae)
     (free-vars* (syntax->list #'(te ce ae)))]
    [(quote datum)
     empty]
    [(quote-syntax datum)
     empty]
    [(with-continuation-mark ke me be)
     (free-vars* (syntax->list #'(ke me be)))]
    [(#%plain-app e ...)
     (free-vars* (syntax->list #'(e ...)))]
    [(#%top . v)
     #;(printf "Not including top ~S in freevars~n" (syntax->datum #'v))
     empty]
    [(#%variable-reference . id)
     (let ([i-bdg (identifier-binding #'id)])
       (cond
         [(eqv? 'lexical i-bdg)
          (list #'id)]
         [(not i-bdg) 
          (list #'id)]
         [else 
          #;(printf "Not including var-reference ~S with binding ~S in freevars~n" (syntax->datum #'id) i-bdg)
          empty]))]
    [id (identifier? #'id)
        (let ([i-bdg (identifier-binding #'id)])
          #;(printf "ID ~S with binding ~S in freevars~n" (syntax->datum #'id) i-bdg)
          (cond
            [(eqv? 'lexical i-bdg)
             (list #'id)]
            [(not i-bdg) 
             (list #'id)]
            [else 
             #;(printf "Not including id ~S with binding ~S in freevars~n" (syntax->datum #'id) i-bdg)
             empty]))]
    [(letrec-syntaxes+values ([(sv ...) se] ...)
       ([(vv ...) ve] ...)
       be ...)
     (free-vars #'(letrec-values ([(vv ...) ve] ...) be ...))]
    [(#%expression d)
     (free-vars #'d)]
    [_
     (raise-syntax-error 'freevars "Dropped through:" stx)]))  

;; free-vars*: (listof expr) -> (listof identifier)
;; union the free variables that occur in several expressions
(define (free-vars* exprs)
  (foldl
   (lambda (expr acc) (union (free-vars expr) acc))
   empty exprs))

;; union: (listof identifier) (listof identifier) -> (listof identifier)
;; produce the set-theoretic union of two lists
(define (union l1 l2)
  (cond
    [(null? l1) l2]
    [else (insert (car l1) (union (cdr l1) l2))]))

(define (union* . ll)
  (foldl union
         empty
         ll))

;; insert: symbol (listof identifier) -> (listof symbol)
;; insert a symbol into a list without creating a duplicate
(define (insert sym into)
  (unless (identifier? sym)
    (raise-syntax-error 'insert "Not identifier" sym))
  (cond
    [(null? into) (list sym)]
    [(free-identifier=? sym (car into)) into]
    [else (cons (car into) (insert sym (cdr into)))]))

;; set-diff: (listof identifier) (listof identifier) -> (listof identifier)
;; produce the set-theoretic difference of two lists
(define (set-diff s1 s2)
  (cond
    [(null? s2) s1]
    [else (set-diff (sans s1 (car s2)) (cdr s2))]))

;; sans: (listof identifier) symbol -> (listof identifier)
;; produce the list sans the symbol
(define (sans s elt)
  (unless (identifier? elt)
    (raise-syntax-error 'sans "Not identifier" elt))
  (cond
    [(null? s) empty]
    [(free-identifier=? (car s) elt)
     (cdr s)] ;; if we maintain the no-dupe invariant then we don't need to recur
    [else (cons (car s)
                (sans (cdr s) elt))]))