#! /usr/bin/mzscheme
#lang scheme

(define (trace fn . args)
  (let ([x (apply fn args)])
    (pretty-print (list fn 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 (variable-value? form)
  (and (symbol? form)
       (not (eq? form '%undef))))

(define (special-value? var)
  (and (memq var '(%nil %self %argv %ctx %k)) #t))

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

(define (simple-value? form)
  (or (variable-value? form)
      (literal-value? form)))

; A value-form is any simple form which can appear on the right-hand side of a (set! ...).
; If there are any side-effect they occur before the variable is updated.
(define (value-form? form)
  (define value-ops
   '(%bind %if %apply %call/cc %make-lambda
     %make-box %unbox %cons %car %cdr %values))
  (or (simple-value? form) (memq (first form) value-ops)))

; A pure-form is any simple form known to be free of side effects.
; Creation of a new object is not counted as a side-effect.
; Pure-forms are a subset of value-forms.
(define (pure-form? form)
  (define pure-ops
   '(%if %make-lambda %make-box %unbox %cons %car %cdr %values))
  (or (simple-value? form) (memq (first form) pure-ops)))

; A statement-form is any simple form which has, or may have, side-effects.
(define (statement-form? form)
  (define statement-ops
   '(%set! %set-box! %set-car! %set-cdr! %apply %call/cc %tail-call))
  (and (pair? form) (memq (first form) statement-ops)))

(define (primitive-form? form)
  (define primitives
   '(%make-box %set-box! %unbox
     %cons %set-car! %car %set-cdr! %cdr
     %if %make-lambda %values))
  (and (pair? form) (memq (first form) primitives)))

(define (bind-form? form)
  (and (pair? form) (eq? (first form) '%bind)))

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

  (cond
    [(variable-value? form) (variable-fn recurse 'variable form)]
    [(literal-value?  form) (literal-fn  recurse 'literal  form)]
    [else
     (let ([handler (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]
                      [else         (if (primitive-form? form)
                                        primitive-fn
                                        other-fn)])])
       (apply handler recurse form))]))

; Like map-form, but intended for boolean results. (Just different defaults.)
(define (search-form form
                     #:merge-with [merge-fn      ormap]
                     #:base-value [base-value    #f]
                     #:bind       [bind-fn       (lambda (recurse op vars . subforms)
                                                   (merge-fn recurse subforms))]
                     #:lambda     [lambda-fn     (lambda (recurse op g-vars i-vars bind)
                                                   (recurse bind))]
                     #:set        [set-fn        (lambda (recurse op var value)
                                                   (recurse value))]

                     #:primitive  [primitive-fn  (lambda (recurse op . simple-values)
                                                   (merge-fn 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) base-value)]
                     #:variable   [variable-fn   simple-fn]
                     #:literal    [literal-fn    simple-fn]

                     #:other      [other-fn      (lambda (recurse . form)
                                                   (error "Unsimplified form:" form))])
  (map-form form
            #: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))

(define (simplify-form form)
  (define (same-form recurse . form) form)
  (map-form form
            #:bind      same-form
            #:lambda    same-form
            #:set       same-form
            #:primitive same-form
            #:simple    (lambda (recurse kind form) form)
            #:literal   (lambda (recurse kind form)
                          (if (and (pair? form)
                                   (eq? (first form) 'quote)
                                   (eq? (second form) '()))
                              '%nil
                              form))
            #:other     (lambda (recurse op . others)
                          (case op
                            [(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))))]
                            [(call/cc values make-box set-box! unbox
                              cons set-car! car set-cdr! cdr)
                             (simplify-primitive (case (first form)
                                                   [(call/cc)  '%call/cc]
                                                   [(values)   '%values]
                                                   [(make-box) '%make-box]
                                                   [(set-box!) '%set-box!]
                                                   [(unbox)    '%unbox]
                                                   [(cons)     '%cons]
                                                   [(set-car!) '%set-car!]
                                                   [(car)      '%car]
                                                   [(set-cdr!) '%set-cdr!]
                                                   [(cdr)      '%cdr])
                                                 (cdr form))]
                            [else (simplify-funcall form)]))))

(define (form-sets? form variable [call-may-set? #t])
  (search-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))
               #:call      (lambda _ call-may-set?)))

(define (form-uses? form variable [call-may-use? #t] [descend? #t])
  (search-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)))
               #:call      (lambda (recurse op . simple-values)
                             (or call-may-use? (ormap recurse simple-values)))
               #:variable  (lambda (recurse op var) (eq? var variable))))

(define (form-captures? form variable [input? #t] [output? #t])
  (search-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))))

(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)]
                          [(and (pair? subform) (eq? (first subform) '%values))
                           ; Requires at least one value; ignores extras.
                           (if (null? (cdr subform))
                               (error "Attempted to set variable to void in:" form)
                               `((set! ,(second form) ,(second subform))))]
                          [(value-form? subform) `((set! ,(second form) ,subform))]
                          [else (error "Attempted to set variable to non-value in:" form)]))
                      '()
                      (cddr value-form))))
        `(%set! ,(second form) ,value-form))))

(define (simplify-primitive simple-op value-forms)
  (define (value->binding value-form)
    (let ([simple-value-form (simplify-form value-form)])
      (if (simple-value? simple-value-form)
          (list simple-value-form #f)
          (let ([tmp (gensym)])
            (list tmp (simplify-set! `(set! ,tmp ,simple-value-form)))))))

  (define bindings (map value->binding value-forms))

  (simplify-form
    `(let ,(map first (filter second bindings))
       ,@(filter-map second bindings)
       (,simple-op ,@(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 (simplify-binding binding)
    (if (pair? binding)
      (list (first binding) (simplify-form (second binding)))
      (list binding)))
  (define bindings (map simplify-binding (second form)))
  (define bodyexprs (cddr form))

  (define (has-value? binding) (pair? (cdr binding)))
  (define vars (map first bindings))
  (define (bound-var? var) (and (memq var vars) #t))

  (flatten-binds
    ; If the value of any binding refers to one of the variable names being bound...
    (if (ormap (lambda (value) (ormap bound-var? (free-variables value)))
               (map second (filter has-value? bindings)))
       ; ...then bind the values to temps first, before binding the real names.
       (let ([temp-bindings (map (lambda (binding)
                                   (let ([tmp (gensym)])
                                     (list tmp
                                           (simplify-form `(set! ,tmp ,(second binding)))
                                           `(%set! ,(first binding) ,tmp))))
                                 (filter has-value? bindings))])
         `(%bind ,(map first temp-bindings)
               ,@(map second temp-bindings)
               (%bind ,vars
                 ,@(map third temp-bindings)
                 ,@(map simplify-form bodyexprs))))
      ; Otherwise, just bind the real names directly.
      `(%bind ,vars
         ,@(map (lambda (binding)
                  (simplify-set! `(set! ,@binding)))
                (filter has-value? 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 (cddr 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 (cddr 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-values (cond-expr true-expr false-expr)
                 (apply values (cdr form)))
  (let ([true-form  (simplify-form true-expr)]
        [false-form (simplify-form false-expr)]
        [cond-val   (gensym)])
    (simplify-form
      (if (and (simple-value? true-form)
               (simple-value? false-form))
          `(let ([,cond-val ,cond-expr])
             (%if ,cond-val ,true-form ,false-form))
          (let ([next-fn  (gensym)]
                [true-fn  (gensym)]
                [false-fn (gensym)])
            `(let ([,cond-val ,cond-expr]
                   [,true-fn  (lambda () ,true-form)]
                   [,false-fn (lambda () ,false-form)])
               (let ([,next-fn (%if ,cond-val ,true-fn ,false-fn)])
                 (%apply ,next-fn %nil))))))))

; (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 form)
  (define argv (gensym))
  `(%bind (,@(second form) ,argv)
     ,@(foldr (lambda (subform after)
                (cond
                  [(pair? after)
                   (cons subform after)]
                  [(simple-value? subform)
                   `((%set! ,argv (%cons ,subform %nil))
                     (%tail-call ,k ,argv #f #f))]
                  [(eq? (first subform) '%apply)
                   `((%tail-call ,(second subform) ,(third subform) ,ctx ,k))]
                  [(eq? (first subform) '%call/cc)
                   `((%set! ,argv (%cons %k %nil))
                     (%tail-call ,(second subform) ,argv ,ctx %k))]
                  [(eq? (first subform) '%values)
                   `((%set! ,argv %nil)
                     ,@(map (lambda (sv) `(%set! ,argv (%cons ,sv ,argv)))
                            (reverse (cdr subform)))
                     (%tail-call ,k ,argv #f #f))]
                  [(value-form? subform)
                   `((%set! ,argv ,subform)
                     (%set! ,argv (%cons ,argv %nil))
                     (%tail-call ,k ,argv #f #f))]
                  [(eq? (first subform) '%tail-call)
                   `(,subform)]
                  [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 (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)))
   (narrow-binds
     `(%lambda () ()
        ,((compose (lambda (bind) (transform-to-cps ctx bind))
                   (lambda (bind) (add-return ctx k bind))
                   flatten-binds)
          `(%bind (,ctx ,k)
             (%set! ,ctx %ctx)
             (%set! ,k %k)
             ,(simplify-form
                `(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)
                (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)])
                      `(%set! ,dest ,(foldl (lambda (var temp-value)
                                              (define temp-bind (fourth temp-value))
                                              (if (form-captures? temp-value var)
                                                  (narrow-binds
                                                    `(%lambda ,g-vars ,i-vars
                                                       (%bind (,@(second temp-bind) ,var)
                                                         ,@(cddr temp-bind))))
                                                  temp-value))
                                            value
                                            extra-binds)))
                    subform))
              (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 bind)
  (define captured-input?  (ormap (lambda (f) (form-captures-input?  f var)) (cddr bind)))
  (define captured-output? (ormap (lambda (f) (form-captures-output? f var)) (cddr bind)))
  (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 bind))
      #f))
  (and (not (special-value? var))
       (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
      [(statement-form? subform)
         (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 rename-vars)
    (define (needs-rename? var) (memq var rename-vars))
    (define (make-binding-unique var bind)
      (let* ([prefix (string-append (symbol->string var) "->g")]
             [unique-var (gensym prefix)])
        (subst-var var unique-var bind)))
    (foldr make-binding-unique bind (filter needs-rename? (second bind))))

  (map-form form
            #:bind   (lambda (recurse op bound-vars . original-subforms)
                       (define rename-vars
                          (remove-duplicates
                            (append (free-variables `(,op ,bound-vars ,@original-subforms))
                                    bound-vars)))
                       (define (form->list subform)
                         (if (bind-form? subform)
                             (let ([unique-form (make-bindings-unique
                                                  (recurse subform)
                                                  rename-vars)])
                               (set! bound-vars (append (second unique-form) bound-vars))
                               (cddr unique-form))
                             (list subform)))
                       (let ([subforms (append-map form->list original-subforms)])
                         `(%bind ,bound-vars ,@subforms)))
            #:lambda (lambda (recurse . form) 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?
                                   (variable-value? form)
                                   (not (memq form '(%nil %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)
  (define (bind-fn recurse op vars . subforms)
    (let* ([reduced-forms (map recurse subforms)]
           [ref-vars (remove-duplicates (append-map free-variables reduced-forms))])
      (define (referenced? var) (and (memq var ref-vars) #t))
      `(%bind ,(filter referenced? vars)
         ,@reduced-forms)))
  (map-form form #:bind bind-fn))

(define (value-used? variable forms)
  (cond
    [(null? forms) #f]
    [(form-uses? (first forms) variable #f #t) #t]
    [(form-sets? (first forms) variable #f) #f]
    [else (value-used? variable (cdr forms))]))

; Don't set variables which won't be accessed later.
(define (reduce-set! form)
  (define (bind-fn recurse op vars . subforms)
    (define (prepend-if-used subform after)
      (if (and (pair? subform)
               (eq? (first subform) '%set!)
               (or (memq (second subform) vars)
                   (error "Setting unbound var:" subform))
               (not (value-used? (second subform) after)))
          after
          (cons subform after)))
    `(%bind ,vars
       ,@(foldr prepend-if-used '() (map recurse subforms))))
  (narrow-binds
    (map-form form #:bind bind-fn)))

(define (propogate-value variable value invalidates? forms)
  (if (null? forms)
      forms
      (let* ([form (car forms)]
             [after (cdr forms)]
             [new-form (case (first form)
                         [(%set!) (if (eq? (third form) variable)
                                      `(%set! ,(second form) ,value)
                                      form)]
                         [else form])])
        (if (or (and (eq? (first (car forms)) '%set!)
                     (eq? (second (car forms)) variable))
                (invalidates? new-form))
            (cons new-form after)
            (cons new-form (propogate-value variable value invalidates? after))))))

; Simple values (literals, variables) can replace arguments as well as %set! values.
(define (propogate-simple-value variable value invalidates? forms)
  (if (null? forms)
      forms
      (let* ([form (car forms)]
             [after (cdr forms)]
             [new-form (case (first form)
                         [(%set!)
                          (let ([set-value (if (eq? (third form) variable) value (third form))])
                            (if (simple-value? set-value)
                                `(%set! ,(second form) ,set-value)
                                `(%set! ,(second form)
                                        (,(first set-value)
                                         ,@(subst variable value (cdr set-value))))))]
                         [else `(,(first form) ,@(subst variable value (cdr form)))])])
        (if (or (and (eq? (first (car forms)) '%set!)
                     (eq? (second (car forms)) variable))
                (invalidates? new-form))
            (cons new-form after)
            (cons new-form (propogate-simple-value variable value invalidates? after))))))

; When value of var2 is known, change (%set! var1 var2) to (%set! var1 value).
; Known values are:
;   literals, always
;   var, until (%set! var ...)
;   (%unbox var), until (%set-box! var ...) or (%set! var)
;   (%car var), until (%set-car! var) or (%set! var)
;   (%cdr var), until (%set-cdr! var) or (%set! var)
(define (propogate-set! form)
  (define (bind-fn recurse op vars . subforms)
    (define (prepend subform after)
      (if (eq? (first subform) '%set!)
          (let ([var (second subform)]
                [value (third subform)])
            (cons
             subform
             (cond
               [(simple-value? value)
                (propogate-simple-value var value
                                        (lambda (form)
                                          (and (eq? (first form) '%set!)
                                               (eq? (second form) value)))
                                        after)]
               [(eq? (first value) '%unbox)
                (let ([box-var (second value)])
                  (propogate-value var value
                                   (lambda (form)
                                     (or (and (eq? (first form) '%set!)
                                              (eq? (second form) box-var))
                                         (and (eq? (first form) '%set-box!)
                                              (eq? (second form) box-var))))
                                   after))]
               [(eq? (first value) '%car)
                (let ([pair-var (second value)])
                  (propogate-value var value
                                   (lambda (form)
                                     (or (and (eq? (first form) '%set!)
                                              (eq? (second form) pair-var))
                                         (and (eq? (first form) '%set-car!)
                                              (eq? (second form) pair-var))))
                                   after))]
               [(eq? (first value) '%cdr)
                (let ([pair-var (second value)])
                  (propogate-value var value
                                   (lambda (form)
                                     (or (and (eq? (first form) '%set!)
                                              (eq? (second form) pair-var))
                                         (and (eq? (first form) '%set-cdr!)
                                              (eq? (second form) pair-var))))
                                   after))]
              [else after])))
          (cons subform after)))
    `(%bind ,vars
       ,@(foldr prepend '() (map recurse subforms))))
  (map-form form #:bind bind-fn))

(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 (machine-var? var)
  (or (special-value? var)
      (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 ([value (if (and (pair? value) (eq? (first value) 'quote))
                       (second value)
                       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 (machine-var? form) form (add-g-var form)))
                 #:literal  (lambda (recurse kind form)
                              (if (eq? form '%nil) form (add-g-var form)))))

    `(,(if (null? i-vars) '%lambda '%template) ,g-vars ,i-vars
       ,bind)))

(define (variable->code var)
  (or (and (eq? var '%nil) #x00)
      (let ([index (find var global-vars)])
        (and index (+ #x01 index)))
      (let ([index (find var instance-vars)])
        (and index (+ #x40 index)))
      (let ([index (find var frame-vars)])
        (and index (+ #x80 index)))
      (let ([index (find var '(%self %argv %ctx %k))])
        (and index (+ #xfc index)))
      (error "No bytecode for variable:" var)))

(define (statement->code form)
  (case (first form)
    [(%set!) (let ([out (variable->code (second form))]
                   [value (third form)])
               (if (machine-var? value)
                   (list #x00 out #x01 (variable->code value) form)
                   (case (first value)
                     [(%unbox)       (list #x00 out #x02 (variable->code (second value)) form)]
                     [(%car)         (list #x00 out #x03 (variable->code (second value)) form)]
                     [(%cdr)         (list #x00 out #x04 (variable->code (second value)) form)]
                     [(%make-lambda) (list #x00 out #x1b (variable->code (second value)) form)]
                     [else (error "Unknown statement type:" form)])))]
    [(%set-box!) (list #x50 (variable->code (second form)) (variable->code (third form)) #x00 form)]
    [(%set-car!) (list #x51 (variable->code (second form)) (variable->code (third form)) #x00 form)]
    [(%set-cdr!) (list #x52 (variable->code (second form)) (variable->code (third form)) #x00 form)]
    [else (error "Unknown statement type:" form)]))

(define (simplify-function lambda-form)
  ((compose
            promote-free-vars
            promote-shared-vars
            simplify-lambda
            )
   lambda-form))

(define (optimize-function simple-lambda-form)
  ((compose
            reduce-variables
            reduce-set!
            propogate-set!
            )
   simple-lambda-form))

(define (compile-function lambda-form)
  ((compose pretty-print
            map-variables
            optimize-function
            simplify-function
            )
   lambda-form))

(compile-function `(lambda () ,(read)))

; vim:set sw=2 expandtab: