#lang scheme/base

(require scheme/list)
(require (file "utilities.scm"))
(require (file "primitives.scm"))

(provide map-variables)

(define (special-constant? sym)
  (memq sym '(#%f #%nil #%undef)))

(define (lambda-value? form)
  (and (pair? form) (eq? (first form) '#%lambda)))

(define (lookup fst lst)
  (ormap (lambda (x)
           (and (eq? (first x) fst)
                (second x)))
         lst))

(define (from-list lst (end-fn (lambda () #f)))
  (lambda ()
    (if (pair? lst)
        (let ([next (first lst)])
          (set! lst (cdr lst))
          next)
        (end-fn))))

(define-syntax values->list
  (syntax-rules ()
    [(values->list expr)
     (call-with-values (lambda () expr) list)]))

(define n-global-variables   (length global-variables))
(define n-instance-variables (length instance-variables))

(define (map-variables lambda/template-form (capture-map '()))
  (let ([bind       (fourth lambda/template-form)]
        [next-t-var (from-list transient-variables (lambda () (error "Out of transient vars")))]
        [next-g-var (from-list global-variables)]
        [next-i-var (from-list instance-variables)]
        [g-var-idx  n-global-variables]
        [i-var-idx  n-instance-variables]
        [t-vars     0]
        [gvar-map   '()]
        [ivar-map   '()]
        [var-map    '()]
        [exprs      '()])

    (define (extra-g-var)
      (let ([mvar (add-expr #f `(#%vector-ref-immed #%globals
                                  ,@(values->list (quotient/remainder g-var-idx 256))))])
        (set! g-var-idx (+ g-var-idx 1))
        mvar))

    (define (extra-i-var)
      (let ([mvar (add-expr #f `(#%vector-ref-immed #%inst
                                  ,@(values->list (quotient/remainder i-var-idx 256))))])
        (set! i-var-idx (+ i-var-idx 1))
        mvar))

    (define (add-g-var value)
      (cond
        [(special-constant? value) value]
        [(or (eq? value #f) (equal? value '(quote #f))) '#%f]
        [(equal? value '(quote ())) '#%nil]
        [else
          (let ([value (cond [(and (pair? value) (eq? (first value) 'quote)) (second value)]
                             [(symbol? value) (error (string-append "Undefined symbol: " (symbol->string value)))]
                             [else value])])
            (or (lookup value gvar-map)
                (let ([g-var (or (next-g-var) (extra-g-var))])
                  (set! gvar-map (cons (list value g-var) gvar-map))
                  g-var)))]))

    (define (add-i-var source)
      (or (and (special-constant? source) source)
          (lookup source ivar-map)
          (let ([i-var (or (next-i-var) (extra-i-var))])
            (set! ivar-map (cons (list source i-var) ivar-map))
            i-var)))

    (define (add-var var mvar)
      (when var
        (set! var-map
              (cons (list var mvar) 
                    (filter (lambda (x) (not (eq? (first x) var)))
                            var-map))))
      mvar)

    (define (add-expr var val)
      (let ([tvar (next-t-var)])
        (set! t-vars (+ 1 t-vars))
        (set! exprs (cons `(#%set! ,tvar ,val) exprs))
        (add-var var tvar)))

    (for ([bound-var (in-list (second bind))])
      (add-var bound-var '#%undef))

    (for ([free-var (in-list (free-variables bind))])
      (let ([capt (lookup free-var capture-map)])
        (when capt (add-var free-var (add-i-var capt)))))

    (let iter ([bind-exprs (cddr bind)])
      (cond
       [(null? bind-exprs)
          (void)]
       [(>= t-vars 120)
          (write-string "Too many expressions; splitting function.\n" (current-error-port))
          (let ([newval (map-variables `(#%lambda () () (#%bind () ,@bind-exprs)) var-map)])
            (add-expr #f
             `(#%tail-call ,(add-g-var newval) #%argv #%kw-args #%kw-vals #%ctx #%k)))]
       [else
          (let* ([expr (car bind-exprs)]
                 [setexpr? (and (pair? expr) (eq? (first expr) '#%set!))]
                 [var (if setexpr? (second expr) #f)]
                 [val (if setexpr? (third expr) expr)])
            (cond [(lambda-value? val)
                    (let ([newval (map-variables val var-map)])
                      (if (eq? (first newval) '#%lambda)
                          (add-var var (add-g-var newval))
                          (add-expr var `(#%make-lambda ,(add-g-var newval)))))]
                  [(literal-value? val)
                    (add-var var (add-g-var val))]
                  [(not (symbol? val))
                    (add-expr var (map-form val
                                    #:variable (lambda (recurse kind form)
                                                 (or (and (machine-variable? form) form)
                                                     (lookup form var-map)
                                                     (add-g-var form)))
                                    #:literal  (lambda (recurse kind form)
                                                 (add-g-var form))))]
                  [else 
                    (add-var var (or (and (machine-variable? val) val)
                                     (lookup val var-map)
                                     (add-g-var val)))]))
          (iter (cdr bind-exprs))]))

    (set! bind `(#%bind () ,@(reverse exprs)))

    `(,(if (null? ivar-map) '#%lambda '#%template)
       ,(map first (reverse gvar-map))
       ,(map first (reverse ivar-map))
       ,bind)))

; vim:set sw=2 expandtab: