#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 (reverse-args new-op args)
    (simplify-form
      (let ([a (gensym)] [b (gensym)])
        `(let ([,a ,(first args)]
               [,b ,(second args)])
           (,new-op ,b ,a)))))

  (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        `(eq? ,@(cdr form)))]
      [(fix>)          (reverse-args 'fix<           (cdr form))]
      [(fix<=)         (reverse-args 'fix>=          (cdr form))]
      [(float>)        (reverse-args 'float<         (cdr form))]
      [(float<=)       (reverse-args 'float>=        (cdr form))]
      [(byte-string>)  (reverse-args 'byte-string<   (cdr form))]
      [(byte-string<=) (reverse-args 'byte-string>=  (cdr form))]
      [(value-list)    (simplify-value-list form)]
      [(values)        (simplify-primitive '#%values  (cdr form))]
      [(list)          (simplify-form `(value-list (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)))))]
      [(and)
       (simplify-form
         (cond
           [(null? (cdr form)) '#t]
           [(null? (cddr form)) (second form)]
           [else (let ([x (gensym)])
                   `(let ([,x ,(second form)])
                      (if ,x (and ,@(cddr form)) ,x)))]))]
      [(or)
       (simplify-form
         (cond
           [(null? (cdr form)) '#f]
           [(null? (cddr form)) (second form)]
           [else (let ([x (gensym)])
                   `(let ([,x ,(second form)])
                      (if ,x ,x (or ,@(cddr form)))))]))]
      [(cond)
       (simplify-form
         (match (cdr form)
           [`() '(values)]
           [`([else . ,forms] . ,_) `(begin ,@forms)]
           [`([,cond-expr . ,forms] . ,rst) `(if ,cond-expr (begin ,@forms) (cond ,@rst))]
           [_ (error "Malformed (cond) form.")]))]
      [(when)   (simplify-form `(if ,(second form) (begin ,@(cddr form)) (values)))]
      [(unless) (simplify-form `(if ,(second form) (values) (begin ,@(cddr 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 (body->forms body)
  (let iter ([body body]
             [bindings '()])
    (match body
      ['()
       (if (null? bindings)
           '()
           `((letrec ,(reverse bindings))))]
      [`((define (,(? symbol? var) . ,arglist) . ,body) . ,rst)
       (iter rst (cons (list var `(lambda ,arglist ,@body)) bindings))]
      [`((define ,(? symbol? var) ,expr) . ,rst)
       (iter rst (cons (list var expr) bindings))]
      [`((define . ,_) . ,_)
       (error "Unrecognized define-form:" (first body))]
      [`(,form . ,rst)
       (if (null? bindings)
           (cons form (iter rst '()))
           `((letrec ,(reverse bindings)
               ,@(cons form (iter rst '())))))])))

(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) '#%nil]
      [`(#%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 (body->forms (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 #%nil #%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))
                   (define (free-var? v) (free-variable? v bind))
                   (define local-binds (filter free-var? extra-bindings))
                   (if (null? local-binds)
                       subform
                       (begin
                         (set! extra-bindings (remove* local-binds extra-bindings))
                         `(#%set! ,var (#%lambda ,g-vars ,i-vars
                                        ,(narrow-binds+promote
                                           `(#%bind (,@(second bind) ,@local-binds)
                                              ,@(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 #%nil #%nil #f #f))]
                      [`(#%apply  . ,sv)
                       `((#%tail-call ,@sv ,ctx ,k))]
                      [`(#%call/cc ,x)
                       `((#%set! ,argv (#%cons ,k #%nil))
                         (#%tail-call ,x ,argv #%nil #%nil ,ctx ,k))]
                      [`(#%values . ,simple-vals)
                       `((#%set! ,argv #%nil)
                         ,@(map (lambda (sv) `(#%set! ,argv (#%cons ,sv ,argv)))
                                (reverse simple-vals))
                         (#%tail-call ,k ,argv #%nil #%nil #f #f))]
                      [(? value-form?)
                       `(,(simplify-set! `(set! ,argv ,subform))
                         (#%set! ,argv (#%cons ,argv #%nil))
                         (#%tail-call ,k ,argv #%nil #%nil #f #f))]
                      [`(#%tail-call . ,_)
                       `(,subform)]
                      [_
                       `(,subform
                         (#%tail-call ,k #%nil #%nil #%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 . ,sv)))
       (let ([k (gensym)])
         `((#%bind (,k)
             (#%set! ,k ,(simplify-form
                          `(lambda ,v
                             ,@after)))
             (#%tail-call ,@sv ,ctx ,k))))]
      [`(#%set! ,v (#%apply . ,sv))
       (let ([k (gensym)])
         `((#%bind (,k)
             (#%set! ,k ,(simplify-form
                          `(lambda (,v . ,(gensym))
                             ,@after)))
             (#%tail-call ,@sv ,ctx ,k))))]
      [(or `(#%value-list (#%apply . ,sv))
           `(#%apply . ,sv))
       (let ([k (gensym)])
         `((#%bind (,k)
             (#%set! ,k ,(simplify-form
                          `(lambda ,(gensym)
                             ,@after)))
             (#%tail-call ,@sv ,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 #%nil #%nil ,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 #%nil #%nil ,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 #%nil #%nil ,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 kw-vals (gensym))

  (define-values (bindings plain-args keywords)
    (let iter ([arg-exprs arg-exprs])
      (match arg-exprs
        [`(,base-expr)
         (values `([,argv ,base-expr]) '() '())]
        [`(,(? keyword? kw) ,expr . ,rst)
         (let-values ([(bnd args kws) (iter rst)]
                      [(x) (gensym)])
           (values (cons `[,x ,expr] bnd) args (cons (cons kw x) kws)))]
        [`(,expr . ,(and rst `(,_ . ,_)))
         (let-values ([(bnd args kws) (iter rst)]
                      [(x) (gensym)])
           (values (cons `[,x ,expr] bnd) (cons x args) kws))])))

  (define sorted-kws (sort keywords keyword<? #:key car))

  (simplify-form
    `(let ([,fn-var ,fn-expr] ,@bindings ,kw-vals)
       ,@(map (lambda (x) `(#%set! ,argv (#%cons ,x ,argv)))
              (reverse plain-args))
       ,(simplify-form `(set! ,kw-vals (list ,@(map cdr sorted-kws))))
       (#%apply ,fn-var
                ,argv
                ,(if (null? sorted-kws)
                     '#%nil
                     `',(map car sorted-kws))
                ,kw-vals))))

; vim:set sw=2 expandtab: