#lang scheme/base

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

(provide simplify-lambda
         promote-free-variables)

(define (simplify-form form)
  (define (same-form recurse . form) form)
  (define (simplify-complex-form 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))))]
      [(value-list) (simplify-value-list form)]
      [(values)     (simplify-primitive '%values  (cdr form))]
      [(apply)      (simplify-apply (second form) (cddr form))]
      [(call/cc)    (simplify-primitive '%call/cc (cdr form))]
      [(call-with-values)
       (simplify-form
         `(apply ,(third form)
                 (value-list (,(second form)))))]
      [else
       (let ([primitive (memf (lambda (x) (eq? (third x) (first form)))
                              all-primitives)])
         (if primitive
             (simplify-primitive (first (first primitive))
                                 (cdr form))
             (simplify-apply (first form) (append (cdr form) '(%nil)))))]))
  (map-form form
            #:bind       same-form
            #:lambda     same-form
            #:set        same-form
            #:value-list same-form
            #:primitive  same-form
            #:simple     (lambda (recurse kind form) form)
            #:literal    (lambda (recurse kind form)
                           (if (equal? form '(quote ())) '%nil form))
            #:other      simplify-complex-form))

(define (simplify-set! form)
  (let ([variable (second form)]
        [value-form (simplify-form (third form))])
    (match value-form
      [`(%bind ,bound-vars . ,subforms)
       (if (memq variable bound-vars)
           (let ([tmp (gensym)])
             `(%bind (,tmp)
                ; guaranteed not to cause unbounded recursion: tmp is unique
                ,(simplify-set! `(set! ,tmp ,value-form))
                (%set! ,variable ,tmp)))
           `(%bind ,bound-vars
              ,@(foldr (lambda (subform after)
                         (if (pair? after)
                             (cons subform after)
                             (list (simplify-set! `(set! ,variable ,subform)))))
                       '()
                       subforms)))]
      [`(%values ,first-val . ,other-vals)
       `(%set! ,variable ,first-val)]
      [`(%values)
       (error "Attempted to set variable to void in:" form)]
      [(? value-form?)
       `(%set! ,variable ,value-form)]
      [else
       (error "Attempted to set variable to void in:" form)])))

(define (simplify-value-list form)
  (let ([values-form (simplify-form (second form))])
    (match values-form
      [`(%bind ,bound-vars . ,subforms)
       `(%bind ,bound-vars
          ,@(foldr (lambda (subform after)
                     (if (pair? after)
                         (cons subform after)
                         (list (simplify-value-list `(value-list ,subform)))))
                   '()
                   subforms))]
      [`(%values . ,simple-vals)
       ; (%value-list (%values ...)) => (list ...)
       (let ([tmp (gensym)])
         `(%bind (,tmp)
            (%set! ,tmp %nil)
            ,@(map (lambda (x) (simplify-set! `(set! ,tmp (cons ,x ,tmp))))
                   (reverse simple-vals))
            ,tmp))]
      [(or `(%apply _ _)
           `(%call/cc _))
       `(%value-list ,values-form)]
      [(? value-form?)
       (simplify-value-list `(value-list (values ,values-form)))]
      [_ '%nil])))

(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 bound-var? (free-variables `(%bind () ,@(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-set! `(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)))
;                     ; TODO: Handle keyword arguments here...
;                     (set! argv-temp (cdr argv-temp))
;                     (...
;                       (let ([rest argv-temp])
;                         bodyexpr...)...)))...)))

(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)))))
            #:value-list (lambda (recurse op values-form)
                           `(,op ,(recurse values-form)))
            #: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))))

(define (is-shared-var? var forms)
  (define (set-after-first-capture?)
    (let/cc return
      (foldr (lambda (form set-after?)
               (if (or set-after? (form-sets? form var #f))
                   (if (form-captures-input? form var)
                       (return #t)
                       #t)
                   #f))
             #f
             forms)
      #f))
  (or (ormap (lambda (f) (form-captures-output? f var)) forms)
      (set-after-first-capture?)))

(define (promote-shared-variables nested-bind)
  (define flat-bind (flatten-binds nested-bind))
  (foldl (lambda (var frm)
           (if (is-shared-var? var (cddr frm))
               (promote-to-box var frm)
               frm))
         flat-bind
         (second flat-bind)))

(define (promote-free-variables form)
  (foldl promote-to-box form (free-variables form)))

(define (narrow-binds+promote nested-bind)
  (define flat-bind (flatten-binds nested-bind))

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

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

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

  (promote-shared-variables
    `(%bind ,(remove* extra-bindings (second flat-bind))
       ,@(map (lambda (subform)
                (match subform
                  [`(%set! ,var (%lambda ,g-vars ,i-vars ,bind))
                   `(%set! ,var (%lambda ,g-vars ,i-vars
                                  ,(narrow-binds+promote
                                     (let ([free-vars (free-variables bind)])
                                       (define (free-var? v) (memq v free-vars))
                                       `(%bind (,@(second bind) ,@(filter free-var? extra-bindings))
                                          ,@(cddr bind))))))]
                  [_ subform]))
              (cddr flat-bind)))))

(define (split-arglist arglist)
  (match arglist
    [`(,(? symbol? reqs) ... ,(and opts (list (? symbol?) _)) ... . ,(? symbol? rst))
     (values reqs opts rst)]
    [`(,(? symbol? reqs) ... ,(and opts (list (? symbol?) _)) ...)
     (values reqs opts #f)]
    [_ (error "Invalid argument list:" arglist)]))

(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)
                (if (pair? after)
                    (cons subform after)
                    (match subform
                      [(? simple-value?)
                       `((%set! ,argv (%cons ,subform %nil))
                         (%tail-call ,k ,argv #f #f))]
                      [`(%apply ,x ,y)
                       `((%tail-call ,x ,y ,ctx ,k))]
                      [`(%call/cc ,x)
                       `((%set! ,argv (%cons %k %nil))
                         (%tail-call ,x ,argv ,ctx %k))]
                      [`(%values . ,simple-vals)
                       `((%set! ,argv %nil)
                         ,@(map (lambda (sv) `(%set! ,argv (%cons ,sv ,argv)))
                                (reverse simple-vals))
                         (%tail-call ,k ,argv #f #f))]
                      [(? value-form?)
                       `(,(simplify-set! `(set! ,argv ,subform))
                         (%set! ,argv (%cons ,argv %nil))
                         (%tail-call ,k ,argv #f #f))]
                      [`(%tail-call . ,_)
                       `(,subform)]
                      [_
                       `(,subform
                         (%tail-call ,k %nil #f #f))])))
               '()
               (cddr flat-bind))))

(define (transform-to-cps ctx nested-bind)
  (define flat-bind (flatten-binds nested-bind))
  (define (cps-prepend subform after)
    (match subform
      [`(%set! ,v (%value-list (%apply ,x ,y)))
       (let ([k (gensym)])
         `((%bind (,k)
             (%set! ,k ,(simplify-form
                          `(lambda ,v
                             ,@after)))
             (%tail-call ,x ,y ,ctx ,k))))]
      [`(%set! ,v (%apply ,x ,y))
       (let ([k (gensym)])
         `((%bind (,k)
             (%set! ,k ,(simplify-form
                          `(lambda (,v . ,(gensym))
                             ,@after)))
             (%tail-call ,x ,y ,ctx ,k))))]
      [(or `(%value-list (%apply ,x ,y))
           `(%apply ,x ,y))
       (let ([k (gensym)])
         `((%bind (,k)
             (%set! ,k ,(simplify-form
                          `(lambda ,(gensym)
                             ,@after)))
             (%tail-call ,x ,y ,ctx ,k))))]
      [`(%set! ,v (%value-list (%call/cc ,x)))
       (let ([k (gensym)]
             [k-argv (gensym)])
         `((%bind (,k ,k-argv)
             (%set! ,k ,(simplify-form
                          `(lambda ,v
                             ,@after)))
             (%set! ,k-argv (%cons ,k %nil))
             (%tail-call ,x ,k-argv ,ctx ,k))))]
      [`(%set! ,v (%call/cc ,x))
       (let ([k (gensym)]
             [k-argv (gensym)])
         `((%bind (,k ,k-argv)
             (%set! ,k ,(simplify-form
                          `(lambda (,v . ,(gensym))
                             ,@after)))
             (%set! ,k-argv (%cons ,k %nil))
             (%tail-call ,x ,k-argv ,ctx ,k))))]
      [(or `(%value-list (%call/cc ,x))
           `(%call/cc ,x))
       (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 ,x ,k-argv ,ctx ,k))))]
      ; keep all other forms with side-effects as-is
      [(? statement-form?) (cons subform after)]
      ; discard any form without side-effects
      [_ after]))
  `(%bind ,(second flat-bind)
     ,@(foldr cps-prepend '() (cddr flat-bind))))

(define (simplify-lambda form)
  (define arglist   (cadr form))
  (define bodyexprs (cddr 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)))

  `(%lambda () ()
     ,((compose narrow-binds+promote
                (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-apply fn-expr arg-exprs)
  (define fn-var (gensym))
  (define argv (gensym))
  (define arguments
    (foldr (lambda (expr args)
             (if (null? args)
                 (cons (list argv expr) args)
                 (if (literal-value? expr)
                     (cons (list expr #f) args)
                     (cons (list (gensym) expr) args))))
           '()
           arg-exprs))
  (simplify-form
    `(let ([,fn-var ,fn-expr] ,@(filter second arguments))
       ,@(map (lambda (x) `(%set! ,argv (%cons ,x ,argv)))
              (map first (reverse (drop-right arguments 1))))
       (%apply ,fn-var ,argv))))

; vim:set sw=2 expandtab: