rosella/libcompiler/simplifier.scm

#lang scheme/base

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

(provide simplify-lambda)

(define (simplify-form form)
  (define (same-form recurse . form) form)
  (define (reverse-args new-op args)
    (simplify-let
      (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)]
      [(keyword-lambda) (simplify-keyword-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))]
      [(values->list)   (simplify-values->list form)]
      [(values)         (simplify-primitive '#%values (cdr form))]
      [(list)           (simplify-form `(values->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) (values->list (,(second form)))))]
      [(and)
       (cond
         [(null? (cdr form)) '#t]
         [(null? (cddr form)) (simplify-form (second form))]
         [else (let ([x (gensym)])
                 (simplify-form
                  `(let ([,x ,(second form)])
                     (if ,x (and ,@(cddr form)) ,x))))])]
      [(or)
       (cond
         [(null? (cdr form)) '#f]
         [(null? (cddr form)) (simplify-form (second form))]
         [else (let ([x (gensym)])
                 (simplify-form
                  `(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
            #:values->list same-form
            #:primitive    same-form
            #:simple       (lambda (recurse kind form) form)
            #:literal      (lambda (recurse kind form) 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))]
      [`((begin . ,forms) . ,rst)
       (iter (append forms rst) bindings)]
      [`(,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-form `(set! ,tmp ,value-form))
                (#%set! ,variable ,tmp)))
           `(#%bind ,bound-vars
              ,@(foldr (lambda (subform after)
                         (if (pair? after)
                             (cons subform after)
                             (list (simplify-form `(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-values->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-form `(values->list ,subform)))))
                   '()
                   subforms))]
      [`(#%values) '#%nil]
      [`(#%values . ,simple-vals)
       ; (#%values->list (#%values ...)) => (list ...)
       (let ([tmp (gensym)])
         `(#%bind (,tmp)
            (#%set! ,tmp #%nil)
            ,@(map (lambda (x) (simplify-form `(set! ,tmp (cons ,x ,tmp))))
                   (reverse simple-vals))
            ,tmp))]
      [(or `(#%apply . ,_)
           `(#%call/cc . ,_))
       `(#%values->list ,values-form)]
      [(? value-form?)
       (simplify-form `(values->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-form `(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 name ([var expr]...) bodyexpr...)
;    => ((letrec ([name (lambda (var...)
;                         bodyexpr...)])
;          name) expr...)

(define (simplify-named-let form)
  (define (expand-binding binding)
    (if (pair? binding) binding (list binding '#%undef)))
  (define let-name (second form))
  (define bindings (map expand-binding (third form)))
  (define bodyexprs (body->forms (cdddr form)))

  (simplify-form
   `((letrec ([,let-name (lambda ,(map first bindings) ,@bodyexprs)]) ,let-name)
     ,@(map second 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-unnamed-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-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-form `(set! ,@binding)))
              (filter has-value? bindings))
       ,@(map simplify-form bodyexprs))))

(define (simplify-let form)
  (if (symbol? (second form))
      (simplify-named-let form)
      (simplify-unnamed-let form)))

; (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-boxes variables form)
  (map-form form
            #:bind         (lambda (recurse op vars . subforms)
                             (let ([unbound-vars (remove* vars variables)])
                               (if (null? unbound-vars)
                                   `(,op ,vars ,@subforms)
                                   (flatten-binds
                                     `(#%bind ,vars
                                       ,@(map (lambda (f) (promote-to-boxes unbound-vars f))
                                              subforms))))))
            #:set          (lambda (recurse op var value)
                             (let ([new-value (recurse value)])
                               (if (memq var variables)
                                   (simplify-form `(set-box! ,var ,new-value))
                                   (simplify-form `(set! ,var ,new-value)))))
            #:values->list (lambda (recurse op values-form)
                             `(,op ,(recurse values-form)))
            #:primitive    (lambda (recurse op . simple-values)
                             (simplify-primitive op (map recurse simple-values)))
            #:variable     (lambda (recurse op var)
                             (if (memq var variables)
                                 `(#%unbox ,var)
                                 var))))

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

(define (promote-shared-variables flat-bind)
  (let* ([shared-vars (filter (lambda (v) (is-shared-var? v (cddr flat-bind)))
                              (second flat-bind))])
    (flatten-binds
      `(#%bind ,(second flat-bind)
         ,@(map (lambda (v) `(#%set! ,v (#%make-box #%undef))) shared-vars)
         ,@(map (lambda (f) (promote-to-boxes shared-vars f)) (cddr flat-bind))))))

(define (narrow-binds+promote flat-bind)
  (define (at-top-level? var)
    (ormap (lambda (x) (or (form-uses? x var #f)
                           (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-vars (free-variables bind))
                   (define (free-var? v) (memq v free-vars))
                   (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 flat-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 #%undef #%undef))]
                      [`(#%apply  ,fn ,av ,kw ,kv)
                       `((#%tail-call ,fn ,av ,kw ,kv ,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 #%undef #%undef))]
                      [(? value-form?)
                       `((#%set! ,argv ,subform)
                         (#%set! ,argv (#%cons ,argv #%nil))
                         (#%tail-call ,k ,argv #%nil #%nil #%undef #%undef))]
                      [`(#%tail-call . ,_)
                       `(,subform)]
                      [_
                       `(,subform
                         (#%tail-call ,k #%nil #%nil #%nil #%undef #%undef))])))
               '()
               (cddr flat-bind))))

(define (transform-to-cps ctx flat-bind)
  (define (cps-prepend subform after)
    (match subform
      [`(#%set! ,v (#%values->list (#%apply ,fn ,av ,kw ,kv)))
       (let ([k (gensym)]
             [t (gensym)])
         `((#%bind (,k)
             (#%set! ,k ,(simplify-form
                          `(lambda ,t
                             (set! ,v ,t)
                             ,@after)))
             (#%tail-call ,fn ,av ,kw ,kv ,ctx ,k))))]
      [`(#%set! ,v (#%apply ,fn ,av ,kw ,kv))
       (let ([k (gensym)]
             [t (gensym)])
         `((#%bind (,k)
             (#%set! ,k ,(simplify-form
                          `(lambda (,t . ,(gensym))
                             (set! ,v ,t)
                             ,@after)))
             (#%tail-call ,fn ,av ,kw ,kv ,ctx ,k))))]
      [(or `(#%values->list (#%apply ,fn ,av ,kw ,kv))
           `(#%apply ,fn ,av ,kw ,kv))
       (let ([k (gensym)])
         `((#%bind (,k)
             (#%set! ,k ,(simplify-form
                          `(lambda ,(gensym)
                             ,@after)))
             (#%tail-call ,fn ,av ,kw ,kv ,ctx ,k))))]
      [`(#%set! ,v (#%values->list (#%call/cc ,x)))
       (let ([k (gensym)]
             [k-argv (gensym)]
             [t (gensym)])
         `((#%bind (,k ,k-argv)
             (#%set! ,k ,(simplify-form
                          `(lambda ,t
                             (set! ,v ,t)
                             ,@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)]
             [t (gensym)])
         `((#%bind (,k ,k-argv)
             (#%set! ,k ,(simplify-form
                          `(lambda (,t . ,(gensym))
                             (set! ,v ,t)
                             ,@after)))
             (#%set! ,k-argv (#%cons ,k #%nil))
             (#%tail-call ,x ,k-argv #%nil #%nil ,ctx ,k))))]
      [(or `(#%values->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
      [(? side-effect-form?) (cons subform after)]
      ; discard any form without side-effects
      [_ after]))
  `(#%bind ,(second flat-bind)
     ,@(foldr cps-prepend '() (cddr flat-bind))))

(define (arguments->lets arglist argv bodyexprs)
  (define-values (requireds optionals rest) (split-arglist arglist))

  (define (add-req req inner)
    `(let ([,req (car ,argv)])
       (set! ,argv (cdr ,argv))
       ,inner))

  (define (add-opt opt-list inner)
    `(let (,(first 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)))

  (foldr add-req
         (foldr add-opt
                rest+bodyexprs
                optionals)
         requireds))

(define (simplify-lambda form)
  (define arglist   (cadr form))
  (define bodyexprs (cddr form))
  (define argv      (gensym))
  (define ctx       (gensym))
  (define k         (gensym))

  `(#%lambda () ()
     ,((compose narrow-binds+promote
                flatten-binds
                (lambda (bind) (transform-to-cps ctx bind))
                (lambda (bind) (add-return ctx k bind))
                flatten-binds
                simplify-form)
       `(let ([,argv #%argv]
              [,ctx  #%ctx]
              [,k    #%k])
          ,(arguments->lets arglist argv bodyexprs)))))

(define (simplify-keyword-lambda form)
  (define arglist     (cadr form))
  (define bodyexprs   (cddr form))
  (define argv        (gensym))
  (define kw-args     (car  arglist))
  (define kw-vals     (cadr arglist))
  (define normal-args (cddr arglist))
  (define ctx         (gensym))
  (define k           (gensym))

  `(#%lambda () ()
     ,((compose narrow-binds+promote
                flatten-binds
                (lambda (bind) (transform-to-cps ctx bind))
                (lambda (bind) (add-return ctx k bind))
                flatten-binds
                simplify-form)
       `(let ([,argv    #%argv]
              [,kw-args #%kw-args]
              [,kw-vals #%kw-vals]
              [,ctx     #%ctx]
              [,k       #%k])
          ,(arguments->lets normal-args argv bodyexprs)))))

; (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))]
        [_ (error "Malformed argument list")])))

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