#! /usr/bin/mzscheme
#lang scheme

(define (trace fn . args)
  (let ([x (apply fn args)])
    (pretty-print x)
    x))

(define (subst old new lst)
  (foldr (lambda (x rst)
           (cons (if (eq? x old)
                     new
                     x)
                 rst))
         '()
         lst))

(define (find x lst)
  (let/cc return
    (for ([i (in-naturals 0)]
          [y (in-list lst)])
      (when (eq? y x) (return i)))
    #f))

(define (simplify-form form)
  (if (pair? form)
      (case (car form)
        [(let) (simplify-let form)]
        [(let*) (simplify-let* form)]
        [(letrec) (simplify-letrec form)]
        [(if) (simplify-if form)]
        [(lambda) (simplify-lambda form)]
        [(begin) (simplify-form `(let () ,@(cdr form)))]
        [(set!) (simplify-set! form)]
        [(let/cc) (simplify-form `(call/cc (lambda (,(second form)) ,@(cddr form))))]
        [(car cdr cons call/cc)
         (simplify-primitive (case (first form)
                               [(make-box) '%make-box]
                               [(set-box!) '%set-box!]
                               [(unbox)    '%unbox]
                               [(cons)     '%cons]
                               [(set-car!) '%set-car!]
                               [(car)      '%car]
                               [(set-cdr!) '%set-cdr!]
                               [(cdr)      '%cdr]
                               [(call/cc)  '%call/cc]
                               )
                             (cdr form))]
        [(quote) (if (eq? (second form) '()) '%nil form)]
        [(%bind %if %set! %lambda
          %tail-call %apply %call/cc
          %cons %set-car! %car %set-cdr! %cdr
          %make-box %set-box! %unbox)
         form]
        [else (simplify-funcall form)])
      (if (eq? form '()) '%nil form)))

(define (simple-value? form)
  (or (not (pair? form))
      (eq? (first form) 'quote)
      (eq? (first form) '%template)))

(define (map-form form
                  #:bind       [bind-fn       (lambda (recurse op vars . subforms)
                                                `(,op ,vars ,@(map recurse subforms)))]
                  #:lambda     [lambda-fn     (lambda (recurse op g-vars i-vars bind)
                                                `(,op ,g-vars ,i-vars ,(recurse bind)))]
                  #:set        [set-fn        (lambda (recurse op var value)
                                                `(,op ,var ,(recurse value)))]

                  #:primitive  [primitive-fn  (lambda (recurse op . simple-values)
                                                `(,op ,@(map recurse simple-values)))]
                  #:call       [call-fn       primitive-fn]
                  #:apply      [apply-fn      call-fn]
                  #:call/cc    [call/cc-fn    call-fn]
                  #:tail-call  [tail-call-fn  call-fn]

                  #:simple     [simple-fn     (lambda (recurse kind form) form)]
                  #:variable   [variable-fn   simple-fn]
                  #:literal    [literal-fn    simple-fn]

                  #:other      [other-fn      (lambda (recurse . form)
                                                (error "Unsimplified form:" form))])
  (define (recurse subform)
    (map-form subform
              #:bind      bind-fn
              #:lambda    lambda-fn
              #:set       set-fn
              #:primitive primitive-fn
              #:call      call-fn
              #:apply     apply-fn
              #:call/cc   call/cc-fn
              #:tail-call tail-call-fn
              #:simple    simple-fn
              #:variable  variable-fn
              #:literal   literal-fn
              #:other     other-fn))

  (if (simple-value? form)
      (if (and (symbol? form) (not (memq form '(%nil %undef))))
          (variable-fn recurse 'variable form)
          (literal-fn  recurse 'literal  form))
      (apply (case (first form)
               [(%bind)      bind-fn]
               [(%lambda)    lambda-fn]
               [(%set!)      set-fn]
               [(%apply)     apply-fn]
               [(%call/cc)   call/cc-fn]
               [(%tail-call) tail-call-fn]

               [(%if %cons %set-car! %car %set-cdr! %cdr %make-box %set-box! %unbox %make-lambda)
                             primitive-fn]

               [else         other-fn])
             recurse form)))

(define (form-sets? form variable [call-may-set? #t])
  (map-form (simplify-form form)
            #:bind      (lambda (recurse op vars . subforms)
                          (and (not (memq variable vars))
                               (ormap recurse subforms)))
            #:lambda    (lambda _ #f)
            #:set       (lambda (recurse op var complex-value) (eq? var variable))
            #:primitive (lambda _ #f)
            #:call      (lambda _ call-may-set?)
            #:simple    (lambda _ #f)))

(define (form-uses? form variable [call-may-use? #t] [descend? #t])
  (map-form (simplify-form form)
            #:bind      (lambda (recurse op vars . subforms)
                          (and (not (memq variable vars))
                               (ormap recurse subforms)))
            #:lambda    (lambda (recurse op g-vars i-vars bind) (and descend? (recurse bind)))
            #:set       (lambda (recurse op var complex-value) (recurse complex-value))
            #:primitive (lambda (recurse op . simple-values) (ormap recurse simple-values))
            #:call      (lambda (recurse op . simple-values)
                          (or call-may-use? (ormap recurse simple-values)))
            #:simple    (lambda _ #f)
            #:variable  (lambda (recurse op var) (eq? var variable))))

(define (form-captures? form variable [input? #t] [output? #t])
  (map-form (simplify-form form)
            #:bind      (lambda (recurse op vars . subforms)
                          (and (not (memq variable vars))
                               (ormap recurse subforms)))
            #:lambda    (lambda (recurse op g-vars i-vars bind)
                          (and (memq variable (free-variables bind input? output?)) #t))
            #:set       (lambda (recurse op var complex-value) (recurse complex-value))
            #:primitive (lambda _ #f)
            #:simple    (lambda _ #f)))

(define (form-captures-input? form var)
  (form-captures? form var #t #f))

(define (form-captures-output? form var)
  (form-captures? form var #f #t))

(define (simplify-set! form)
  (let ([value-form (simplify-form (third form))])
    (if (and (pair? value-form) (eq? (first value-form) '%bind))
        (simplify-form
          `(let ,(second value-form)
             ,@(foldr (lambda (subform after)
                        (cond
                          [(pair? after) (cons subform after)]
                          [(or (not (pair? subform))
                               (memq (first subform) '(%apply %call/cc %car %cdr %cons %bind %if %unbox quote)))
                           `((set! ,(second form) ,subform))]
                          [(and (pair? subform) (eq? (first subform) '%tail-call))
                           `(,subform)] ; The %set! wouldn't be executed anyway.
                          [else
                           `(,subform
                             (%set! ,(second form) %void))]))
                      '()
                      (cddr value-form))))
        `(%set! ,(second form) ,value-form))))

(define (simplify-primitive new-id value-forms)
  (define bindings (map (lambda (vf)
                          (let ([simple-vf (simplify-form vf)])
                            (if (simple-value? simple-vf)
                                (list simple-vf #f)
                                (let ([tmp (gensym)])
                                  (list tmp (simplify-set! `(set! ,tmp ,simple-vf)))))))
                        value-forms))
  (simplify-form
    `(let ,(map first (filter second bindings))
       ,@(filter-map second bindings)
       (,new-id ,@(map first bindings)))))

; (let ([var expr]...) bodyexpr...)  ; first eval exprs, then assign bindings in parallel
;    => (%bind (tmp...)
;         (%set! tmp ,(simplify-form expr))...
;         (%bind (var...)
;           (%set! var tmp)...
;           bodyexpr...))

(define (simplify-let form)
  (define bindings (second form))
  (define bodyexprs (cdr (cdr form)))
  (flatten-binds
    (cond
      [(not (pair? bindings))
       `(%bind () ,@(map simplify-form bodyexprs))]
      [(not (pair? (cdr bindings)))
       (let ([binding (first bindings)])
         `(%bind (,(if (pair? binding) (first binding) binding))
            ,@(if (pair? binding) `(,(simplify-set! `(set! ,(first binding)
                                                           ,(second binding))))
                                  '())
            ,@(map simplify-form bodyexprs)))]
      [else
       (let ([vars (map (lambda (x) (if (pair? x) (first x) x)) bindings)]
             [temp-bindings (append-map (lambda (x)
                                          (if (pair? x)
                                              (let ([tmp (gensym)])
                                                `((,tmp
                                                   ,(simplify-form `(set! ,tmp ,(second x)))
                                                   (%set! ,(first x) ,tmp))))
                                              '()))
                                        bindings)])
         `(%bind ,(map first temp-bindings)
               ,@(map second temp-bindings)
               (%bind ,vars
                 ,@(map third temp-bindings)
                 ,@(map simplify-form bodyexprs))))])))

; (let* ...)   ; eval exprs & bind variables serially
;    => (let ([var-0 expr-0])
;         (let ([var-1 expr-1])
;           (...
;             bodyexprs...)))

(define (simplify-let* form)
   (define bindings (second form))
   (define bodyexprs (cdr (cdr form)))
   (define (add-binding bind bodyexpr)
      `(let (,bind) ,bodyexpr))
   (simplify-form (foldr add-binding `(begin ,@bodyexprs) bindings)))

; (letrec ...) ; init bindings to undefined, then assign values in series
;    => (let (var...)
;         (set! var expr)...
;         bodyexprs)

(define (simplify-letrec form)
  (define bindings (second form))
  (define bodyexprs (cdr (cdr form)))
  (simplify-form
    `(let (,@(map (lambda (x) (if (pair? x) (first x) x)) bindings))
       ,@(append-map
           (lambda (x)
             (if (pair? x)
                 `((set! ,(first x) ,(second x)))
                 '()))
           bindings)
       ,@bodyexprs)))

(define (simplify-if form)
  (define cond-val (gensym))
  (define next-fn (gensym))
  (define true-fn (gensym))
  (define false-fn (gensym))
  (define-values (cond-expr true-expr false-expr)
                 (apply values (cdr form)))
  (simplify-form
    (if (or (pair? true-expr) (pair? false-expr))
        `(let ([,cond-val ,cond-expr]
               [,true-fn  (lambda () ,true-expr)]
               [,false-fn (lambda () ,false-expr)])
           (let ([,next-fn (%if ,cond-val ,true-fn ,false-fn)])
             (%apply ,next-fn %nil)))
        `(let ([,cond-val ,cond-expr])
           (%if ,cond-val ,(simplify-form true-expr) ,(simplify-form false-expr))))))

; (lambda (required... [optional default-expr]... . rest) bodyexpr...)
;    => (lambda argv
;         (let ([argv-temp argv])
;           (let ([required-0 (car argv-temp)])
;             (set! argv-temp (cdr argv-temp)))
;             (let ([required-1 (car argv-temp)])
;               (set! argv-temp (cdr argv-temp)))
;               (...
;                 (let (optional-0)
;                   (if (eq? argv-temp %nil)
;                       (set! optional-0 default-expr-0)
;                       (set! optional-0 (car argv-temp)))
;                   (set! argv-temp (cdr argv-temp))
;                   (let (optional-1)
;                     (if (eq? argv-temp %nil)
;                         (set! optional-1 default-expr-1)
;                         (set! optional-1 (car argv-temp)))
;                     (set! argv-temp (cdr argv-temp))
;                     (...
;                       (let ([rest argv-temp])
;                         bodyexpr...)...)))...)))

(define (split-arglist arglist)
   (define (split-optional arglist)
      (if (pair? arglist)
          (let-values ([(opt rst) (split-optional (cdr arglist))])
            (values (cons (car arglist) opt) rst)) 
          (values '() arglist)))
   (if (pair? arglist)
       (if (pair? (car arglist))
         (let-values ([(opt rst) (split-optional arglist)])
            (values '() opt rst))
         (let-values ([(req opt rst) (split-arglist (cdr arglist))])
            (values (cons (car arglist) req) opt rst)))
       (values '() '() #f)))

(define (add-return ctx k rval form)
  (define k-argv (gensym))
  `(%bind (,@(second form) ,k-argv)
     ,@(foldr (lambda (subform after)
                (cond
                  [(pair? after)
                   (cons subform after)]
                  [(and (pair? subform)
                        (eq? (first subform) '%set!)
                        (eq? (second subform) rval)
                        (pair? (third subform))
                        (eq? (first (third subform)) '%apply))
                   (let ([fn (second (third subform))]
                         [argv (third (third subform))])
                     `((%tail-call ,fn ,argv ,ctx ,k)))]
                  [(and (pair? subform)
                        (eq? (first subform) '%set!)
                        (eq? (second subform) rval)
                        (pair? (third subform))
                        (eq? (first (third subform)) '%call/cc))
                   (let ([fn (second (third subform))])
                     `((%set! ,k-argv (%cons %k %nil))
                       (%tail-call ,fn ,k-argv ,ctx %k)))]
                  [(and (pair? subform)
                        (eq? (first subform) '%tail-call))
                   `(,subform)]
                  [(and (pair? subform)
                        (eq? (first subform) '%apply))
                   `((%tail-call ,(second subform)
                                 ,(third subform)
                                 ,ctx
                                 ,k))]
                  [(and (pair? subform)
                        (eq? (first subform) '%set!)
                        (eq? (second subform) rval)
                        (eq? (third subform) '%void))
                   `((%tail-call ,k %nil #f #f))]
                  [(and (pair? subform)
                        (eq? (first subform) '%set!)
                        (eq? (second subform) rval))
                   `(,subform
                     (%set! ,rval (%cons ,rval %nil))
                     (%tail-call ,k ,rval #f #f))]
                  [else
                   `(,subform
                     (%tail-call ,k %nil #f #f))]))
               '()
               (cddr form))))

(define (simplify-lambda form)
   (define arglist (car (cdr form)))
   (define bodyexprs (cdr (cdr form)))

   (define-values (requireds optionals rest) (split-arglist arglist))
   (define argv-temp (gensym))
   (define ctx (gensym))
   (define k (gensym))
   (define rval (gensym))

   (define (add-req req inner) `(let ([,req (car ,argv-temp)])
                                  (set! ,argv-temp (cdr ,argv-temp))
                                  ,inner))
   (define (add-opt opt-list inner) `(let (,(car opt-list))
                                       (if (pair? ,argv-temp)
                                           (begin
                                             (set! ,(first opt-list) (car ,argv-temp))
                                             (set! ,argv-temp (cdr ,argv-temp)))
                                           (set! ,(first opt-list) ,(second opt-list)))
                                       ,inner))
   (define rest+bodyexprs (if rest `(let ([,rest ,argv-temp]) ,@bodyexprs)
                                   `(begin ,@bodyexprs)))
   `(%lambda () ()
      ,((compose (lambda (x) (transform-to-cps ctx x))
                 (lambda (x) (add-return ctx k rval x))
                 flatten-binds)
        `(%bind (,rval ,ctx ,k)
           (%set! ,ctx %ctx)
           (%set! ,k %k)
           ,(simplify-form
              `(set! ,rval (let ([,argv-temp %argv])
                             ,(foldr add-req
                                     (foldr add-opt
                                            rest+bodyexprs
                                            optionals)
                                     requireds))))))))

(define (narrow-binds simple-lambda-form)
  (define bind (fourth simple-lambda-form))

  (define (at-top-level? var)
    (or (ormap (lambda (x) (form-sets? x var #f)) (cddr bind))
        (ormap (lambda (x) (form-uses? x var #f #f)) (cddr bind))))

  (define (captured-twice? var)
    (let/cc return
      (foldl (lambda (subform once?)
               (if (form-captures? subform var)
                   (if once? (return #t) #t)
                   once?))
             (at-top-level? var)
             (cddr bind))
      #f))

  (define extra-binds
    (filter-not captured-twice?
                (filter-not at-top-level?
                            (second bind))))

  `(%lambda ,(second simple-lambda-form) ,(third simple-lambda-form)
     (%bind ,(remove* extra-binds (second bind))
       ,@(map (lambda (subform)
                (foldl (lambda (var subform)
                         (if (and (pair? subform)
                                  (eq? (first subform) '%set!)
                                  (pair? (third subform))
                                  (eq? (first (third subform)) '%lambda))
                             (let* ([dest (second subform)]
                                    [value (third subform)]
                                    [g-vars (second value)]
                                    [i-vars (third value)]
                                    [bind (fourth value)])
                               `(%set! ,dest ,(narrow-binds
                                                `(%lambda ,g-vars ,i-vars
                                                   ,(if (form-captures? value var)
                                                        `(%bind (,@(second bind) ,var)
                                                           ,@(cddr bind))
                                                        bind)))))
                             subform))
                         subform
                         extra-binds))
              (cddr bind)))))

(define (promote-to-box variable form)
  (map-form form
            #:bind      (lambda (recurse op vars . subforms)
                          (flatten-binds
                            `(%bind ,(subst variable variable vars)
                               ,@(if (memq variable vars)
                                     `((%set! ,variable (%make-box %undef)))
                                     '())
                               ,@(map recurse subforms))))
            #:set       (lambda (recurse op var value)
                          (let ([new-value (recurse value)])
                            (if (eq? var variable)
                                (if (simple-value? new-value)
                                    `(%set-box! ,variable ,new-value)
                                    (let ([tmp (gensym)])
                                      `(%bind (,tmp)
                                         ,(simplify-set! `(set! ,tmp ,new-value))
                                         (%set-box! ,variable ,tmp))))
                                (simplify-set! `(set! ,var ,new-value)))))
            #:primitive (lambda (recurse op . simple-values)
                          (let ([new-args (map recurse simple-values)])
                            ;; if any new-arg is not simple, must bind to a temp first
                            (let ([temps (map (lambda (x)
                                                (if (simple-value? x)
                                                    (list x #f)
                                                    (let ([tmp (gensym)])
                                                      (list tmp `(%set! ,tmp ,x)))))
                                              new-args)])
                              (if (ormap second temps)
                                  `(%bind ,(map first (filter second temps))
                                     ,@(filter-map second temps)
                                     (,op ,@(map first temps)))
                                  `(,op ,@new-args)))))
            #:variable  (lambda (recurse op var)
                          (if (eq? var variable) `(%unbox ,variable) var))))

; form needs to be flattened (%bind ...)
(define (is-shared-var? var form)
  (define captured-input?  (ormap (lambda (f) (form-captures-input?  f var)) (cddr form)))
  (define captured-output? (ormap (lambda (f) (form-captures-output? f var)) (cddr form)))
  (define (set-after-first-use?)
    (let/cc return
      (foldr (lambda (subform set-after?)
               (if (or set-after? (form-sets? subform var captured-output?))
                   (if (form-uses? subform var captured-input?)
                       (return #t)
                       #t)
                   #f))
             #f
             (cddr form))
      #f))
  (and (or captured-input?
           captured-output?)
       (set-after-first-use?)))

(define (promote-shared-vars simple-lambda-form)
  (define bind (fourth simple-lambda-form))
  `(%lambda ,(second simple-lambda-form) ,(third simple-lambda-form)
     ,(foldl (lambda (var frm)
               (if (is-shared-var? var frm)
                   (promote-to-box var frm)
                   frm))
             bind
             (second bind))))

(define (promote-free-vars simple-lambda-form)
  (define bind (fourth simple-lambda-form))
  `(%lambda ,(second simple-lambda-form) ,(third simple-lambda-form)
     ,(foldl promote-to-box bind (free-variables bind))))

; <= (%bind (var...)
;      @before
;      (%apply x y)
;      @after))
; => (%bind (var... k)
;      @before
;      (%set! k (lambda _ @after))
;      (%tail-call x y ctx k)))

; <= (%bind (var...)
;      @before
;      (%set! v (%apply x y))
;      @after))
; => (%bind (var... k)
;      @before
;      (%set! k (lambda (x)
;                 (%set! v x)
;                 @after))
;      (%tail-call x y ctx k)))

; <= (%bind (var...)
;      @before
;      (call/cc l)
;      @after)
; => (%bind (var... k k2)
;      @before
;      (%set! k (lambda _ @after))
;      (%set! k-argv (%cons k %nil))
;      (%tail-call l k-argv ctx k))

(define (transform-to-cps ctx bind)
  (define (cps-prepend subform after)
    (cond
      ; (%set! v (%apply x y))
      [(and (pair? subform)
            (eq? (first subform) '%set!)
            (pair? (third subform))
            (eq? (first (third subform)) '%apply))
         (let ([k (gensym)]
               [x (gensym)])
           `((%bind (,k ,x)
               (%set! ,k ,(simplify-form
                            `(lambda (,x . ,(gensym))
                               (%set! ,(second subform) ,x)
                               ,@after)))
               (%tail-call ,(second (third subform))
                           ,(third (third subform))
                           ,ctx
                           ,k))))]
      ; (%apply x y)
      [(and (pair? subform)
            (eq? (first subform) '%apply))
         (let ([k (gensym)])
           `((%bind (,k)
               (%set! ,k ,(simplify-form
                            `(lambda ,(gensym)
                               ,@after)))
               (%tail-call ,(second subform)
                           ,(third subform)
                           ,ctx
                           ,k))))]
      ; (%set! v (%call/cc x))
      [(and (pair? subform)
            (eq? (first subform) '%set!)
            (pair? (third subform))
            (eq? (first (third subform)) '%call/cc))
         (let ([k (gensym)]
               [k-argv (gensym)]
               [x (gensym)])
           `((%bind (,k ,k-argv)
               (%set! ,k ,(simplify-form
                            `(lambda (,x . ,(gensym))
                               (%set! ,(second subform) ,x)
                               ,@after)))
               (%set! ,k-argv (%cons ,k %nil))
               (%tail-call ,(second (third subform))
                           ,k-argv
                           ,ctx
                           ,k))))]
      ; (%call/cc x)
      [(and (pair? subform)
            (eq? (first subform) '%call/cc))
         (let ([k (gensym)]
               [k-argv (gensym)])
           `((%bind (,k ,k-argv)
               (%set! ,k ,(simplify-form
                            `(lambda ,(gensym)
                               ,@after)))
               (%set! ,k-argv (%cons ,k %nil))
               (%tail-call ,(second subform)
                           ,k-argv
                           ,ctx
                           ,k))))]
      ; keep all other forms with side-effects as-is
      [(and (pair? subform)
            (memq (first subform) '(%set! %set-box! %set-car! %set-cdr! %tail-call)))
         (cons subform after)]
      ; discard any form without side-effects
      [else after]))
  (flatten-binds
    `(%bind ,(second bind)
       ,@(foldr cps-prepend '() (cddr bind)))))

; (fn-expr arg-expr...)
;    => (let ([fn-var fn-expr] arg-var... argv)
;         (set! fn-var fn-expr)
;         (set! arg-var arg-expr)...
;         (set! argv %nil)
;         (set! argv (cons arg-var argv))... [reversed]
;         (%apply fn-var argv))

(define (simplify-funcall form)
  (define fn-expr (car form))
  (define arg-exprs (cdr form))
  (define fn-var (gensym))
  (define arg-vars (map (lambda (x) (gensym)) arg-exprs))
  (define argv (gensym))
  (simplify-form
    `(let (,fn-var ,@arg-vars ,argv)
       (set! ,fn-var ,fn-expr)
       ,@(map (lambda (x y) `(set! ,x ,y)) arg-vars arg-exprs)
       (%set! ,argv %nil)
       ,@(reverse (map (lambda (x) `(%set! ,argv (%cons ,x ,argv))) arg-vars))
       (%apply ,fn-var ,argv))))

(define (subst-var old-var new-var form)
  (map-form form
            #:bind     (lambda (recurse op vars . subforms)
                         `(%bind ,(subst old-var new-var vars) ,@(map recurse subforms)))
            #:set      (lambda (recurse op var value)
                         `(,op ,(if (eq? var old-var) new-var var) ,(recurse value)))
            #:variable (lambda (recurse op var)
                         (if (eq? var old-var) new-var var))))

(define (flatten-binds form)
  (define (make-bindings-unique bind)
    (foldr (lambda (var bind)
             (subst-var var (gensym) bind))
           bind
           (second bind)))

  (if (and (pair? form) (eq? (car form) '%bind))
      (let* ([bound-vars (second form)]
             [subforms (append-map (lambda (subform)
                                     (if (and (pair? subform) (eq? (car subform) '%bind))
                                         (let ([unique-form (make-bindings-unique
                                                              (flatten-binds subform))])
                                           (set! bound-vars (append (second unique-form) bound-vars))
                                           (cddr unique-form))
                                         (list subform)))
                                   (cddr form))])
        `(%bind ,bound-vars ,@subforms))
      form))

(define (free-variables form [input? #t] [output? #t])
  (map-form form
            #:bind      (lambda (recurse op vars . subforms)
                          (remove-duplicates (remove* vars (append-map recurse subforms))))
            #:lambda    (lambda (recurse op g-vars i-vars bind)
                          (recurse bind))
            #:set       (lambda (recurse op var value)
                          (let ([value-free (recurse value)])
                            (if output?
                                (cons var value-free)
                                value-free)))
            #:primitive (lambda (recurse op . simple-values)
                          (remove-duplicates (append-map recurse simple-values)))
            #:simple    (lambda (recurse kind form)
                          (if (and input?
                                   (symbol? form)
                                   (not (memq form '(%nil %undef %self %argv %ctx %k))))
                              (list form)
                              '()))))

(define (free-input-variables form)
  (free-variables form #t #f))

(define (free-output-variables form)
  (free-variables form #f #t))

; Don't bind variables which aren't referenced.
(define (reduce-variables form)
  (map-form form
            #:bind (lambda (recurse op vars . subforms)
                     (let ([ref-vars (remove-duplicates (append-map free-variables subforms))])
                       `(%bind ,(filter (lambda (x) (memq x ref-vars)) vars)
                          ,@(map recurse subforms))))))

; Don't set variables which won't be accessed later.
(define (reduce-set! form)
  (map-form form
            #:bind (lambda (recurse op vars . subforms)
                     (define (prepend-if-used subform after)
                       (if (and (pair? subform)
                                (eq? (first subform) '%set!)
                                (memq (second subform) vars)
                                (not (memq (second subform)
                                           (append-map free-input-variables after))))
                           after
                           (cons subform after)))
                     `(%bind ,vars
                        ,@(foldr prepend-if-used '() (map recurse subforms))))))

(define (simplify-toplevel-lambda form)
  (promote-free-vars
    (promote-shared-vars
      (narrow-binds
        (simplify-lambda form)))))

(define (optimize-simplified-lambda form)
  (reduce-variables
    (reduce-set!
      form)))

(define frame-vars
  (for/list ([i (in-range 0 120)])
    (string->uninterned-symbol (string-append "%f" (number->string i)))))

(define instance-vars
  (for/list ([i (in-range 0 64)])
    (string->uninterned-symbol (string-append "%i" (number->string i)))))

(define global-vars
  (for/list ([i (in-range 1 64)])
    (string->uninterned-symbol (string-append "%g" (number->string i)))))

(define (frame-var? var) (and (memq var frame-vars) #t))
(define (instance-var? var) (and (memq var instance-vars) #t))
(define (frame/instance-var? var) (or (frame-var? var) (instance-var? var)))

(define (global-var? var) (and (memq var global-vars) #t))

(define (special-var? var)
  (or (and (memq var '(%nil %self %argv %ctx %k)) #t)
      (frame/instance-var? var)
      (global-var? var)))

(define (map-variables lambda/template-form)
  (let ([bind (fourth lambda/template-form)]
        [g-vars '()]
        [unused-g-vars global-vars]
        [i-vars '()])
    (define (add-g-var value)
      (let/cc return
        (for ([g-var (in-list global-vars)]
              [val (in-list g-vars)])
          (when (eq? value val) (return g-var)))
        (let ([g-var (first unused-g-vars)])
          (set! unused-g-vars (cdr unused-g-vars))
          (set! g-vars (append g-vars (list value)))
          g-var)))

    (for ([free-var (in-list (filter frame/instance-var? (free-variables bind)))]
          [inst-var (in-list instance-vars)])
      (set! i-vars (append i-vars (list free-var)))
      (set! bind (subst-var free-var inst-var bind)))
    (for ([bound-var (in-list (second bind))]
          [frame-var (in-list frame-vars)])
      (set! bind (subst-var bound-var frame-var bind)))
    (set! bind (map-form bind
                 #:lambda   (lambda (recurse op inner-g-vars i-vars bind)
                              `(%make-lambda ,((compose add-g-var map-variables)
                                               `(%template ,inner-g-vars ,i-vars ,bind))))
                 #:variable (lambda (recurse kind form)
                              (if (special-var? form) form (add-g-var form)))
                 #:literal  (lambda (recurse kind form)
                              (if (eq? form '%nil) form (add-g-var form)))))
    `(,(if (pair? i-vars) '%template '%lambda) ,g-vars ,i-vars
       ,bind)))

((compose
   pretty-print
   map-variables
   optimize-simplified-lambda
   simplify-toplevel-lambda)
 `(lambda () ,(read)))

; vim:set sw=2 expandtab: