#lang scheme/base

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

(provide simplify-function
         promote-free-variables)

(define (simplify-function lambda-form)
  ((compose promote-shared-variables
            simplify-lambda)
   lambda-form))

(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))))]
                            [(fix>) (simplify-form
                                      (let ([a (gensym)] [b (gensym)])
                                        `(let ([,a ,(second form)]
                                               [,b ,(third form)])
                                           (fix< ,b ,a))))]
                            [(fix<=) (simplify-form
                                       (let ([a (gensym)] [b (gensym)])
                                         `(let ([,a ,(second form)]
                                                [,b ,(third form)])
                                            (fix>= ,b ,a))))]
                            [(values)  (simplify-primitive '%values  (cdr form))]
                            [(call/cc) (simplify-primitive '%call/cc (cdr form))]
                            [else
                             (let ([primitive (memf (lambda (x) (eq? (third x) (first form)))
                                                    all-primitives)])
                               (if primitive
                                   (simplify-primitive (first (first primitive))
                                                       (cdr form))
                                   (simplify-funcall form)))]))))

(define (simplify-set! form)
  (let ([variable (second form)]
        [value-form (simplify-form (third form))])
    (if (and (pair? value-form) (eq? (first value-form) '%bind))
        (if (memq variable (second value-form))
            (let ([tmp (gensym)])
              `(%bind (,tmp)
                 ; guaranteed not to cause unbounded recursion: tmp is unique
                 ,(simplify-set! `(set! ,tmp ,value-form))
                 (%set! ,variable ,tmp)))
            `(%bind ,(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! ,variable ,(second subform))))]
                            [(value-form? subform)
                             (list (simplify-set! `(set! ,variable ,subform)))]
                            [else (error "Attempted to set variable to void in:" form)]))
                        '()
                        (cddr value-form))))
        `(%set! ,variable ,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))

  ; 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)))
   (cond
    [(null? arglist)       (values '() '() #f)]
    [(not (pair? arglist)) (values '() '() arglist)]
    [(pair? (car arglist)) (let-values ([(opt rst) (split-optional arglist)])
                             (values '() opt rst))]
    [else                  (let-values ([(req opt rst) (split-arglist (cdr arglist))])
                             (values (cons (car arglist) req) opt rst))]))

(define (add-return ctx k nested-bind)
  (define flat-bind (flatten-binds nested-bind))
  (define argv (gensym))
  `(%bind (,@(second flat-bind) ,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 flat-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 nested-bind)
  (define flat-bind (flatten-binds nested-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 flat-bind)
       ,@(foldr cps-prepend '() (cddr flat-bind)))))

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

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

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

; (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 (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-variable? var))
       (or captured-input?
           captured-output?)
       (set-after-first-use?)))

(define (promote-shared-variables 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-variables 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))))

; vim:set sw=2 expandtab: