Implement conversion to tail-call form (CPS).

Also add some basic optimizations concerning binding and setting
variables which will never be used.

compiler.ss

@@ -1,6 +1,11 @@
 #! /usr/bin/mzscheme
 #lang scheme
 
+(define (trace fn . args)
+  (let ([x (apply fn args)])
+    (pretty-print x)
+    x))
+
 (define (simplify-form form)
   (if (pair? form)
     (case (car form)
@@ -24,10 +29,13 @@
   (let ([value-form (simplify-form (third form))])
     (if (and (pair? value-form) (eq? (first value-form) '%bind))
         `(%bind ,(second value-form)
-           ,@(foldr (lambda (x s)
+           ,@(foldr (lambda (subform after)
-                      (if (pair? s)
+                      (if (pair? after)
-                          (cons x s)
+                          (cons subform after)
-                          (list (simplify-set! `(set! ,(second form) ,x)))))
+                          (if (or (not (pair? subform))
+                                  (memq (first subform) '(%apply %car %cdr %cons %bind)))
+                              (list (simplify-set! `(set! ,(second form) ,subform)))
+                              (list subform))))
                     '()
                     (cddr value-form)))
         `(%set! ,(second form) ,value-form))))
@@ -77,7 +85,7 @@
                   `(,(simplify-form `(set! ,(first vars) ,(second (first bindings)))))
                   '())
               ,@(map simplify-form bodyexprs)))
-       `(%bind () ,@bodyexprs)))
+       `(%bind () ,@(map simplify-form bodyexprs))))
 
 ; (let* ...)   ; eval exprs & bind variables serially
 ;    => (let ([var-0 expr-0])
@@ -100,15 +108,14 @@
 (define (simplify-letrec form)
   (define bindings (second form))
   (define bodyexprs (cdr (cdr form)))
-  (simplify-form
   `(%bind (,@(map (lambda (x) (if (pair? x) (first x) x)) bindings))
      ,@(append-map
          (lambda (x)
            (if (pair? x)
-               `((set! ,(first x) ,(simplify-form (second x))))
+               `(,(simplify-form `(set! ,(first x) ,(simplify-form (second x)))))
                '()))
          bindings)
-     ,@(map simplify-form bodyexprs))))
+     ,@(map simplify-form bodyexprs)))
 
 (define (simplify-if form)
   (define cond-val (gensym))
@@ -146,9 +153,6 @@
 ;                       (let ([rest argv-temp])
 ;                         bodyexpr...)...)))...)))
 
-(define (simplify-lambda form)
-   (define arglist (car (cdr form)))
-   (define bodyexprs (cdr (cdr form)))
 (define (split-arglist arglist)
    (define (split-optional arglist)
       (if (pair? arglist)
@@ -163,6 +167,42 @@
             (values (cons (car arglist) req) opt rst)))
        (values '() '() #f)))
 
+(define (add-tail-call k rval form)
+  `(%bind ,(second form)
+     ,@(foldr (lambda (subform after)
+                (if (pair? after)
+                    (cons subform after)
+                    (if (and (pair? subform)
+                             (eq? (first subform) '%set!)
+                             (eq? (second subform) rval)
+                             (pair? (third subform))
+                             (eq? (first (third subform)) '%apply))
+                        (let ([fn (second (third subform))]
+                              [argv (third (third subform))])
+                          `((%tail-call ,fn ,argv ,k)))
+                        (if (and (pair? subform)
+                                 (eq? (first subform) '%tail-call))
+                            (list subform)
+                            (if (and (pair? subform)
+                                     (eq? (first subform) '%apply))
+                                `((%tail-call ,(second subform)
+                                              ,(third subform)
+                                              ,k))
+                                (if (and (pair? subform)
+                                         (eq? (first subform) '%set!)
+                                         (eq? (second subform) rval))
+                                    `(,subform
+                                      (%set! ,rval (%cons ,rval %nil))
+                                      (%tail-call ,k ,rval #f))
+                                    `(,subform
+                                      (%tail-call ,k %nil #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 k (gensym))
@@ -180,25 +220,9 @@
                                        ,inner))
    (define rest+bodyexprs (if rest `(let ([,rest ,argv-temp]) ,@bodyexprs)
                                    `(begin ,@bodyexprs)))
-   (define (add-tail-call k rval bind)
-     `(%bind ,(second bind)
-        ,@(foldr (lambda (subform after)
-                   (if (pair? after)
-                       (cons subform after)
-                       (if (and (pair? subform)
-                                (eq? (first subform) '%set!)
-                                (eq? (second subform) rval)
-                                (pair? (third subform))
-                                (eq? (first (third subform)) '%apply))
-                           (let ([fn (second (third subform))]
-                                 [argv (third (third subform))])
-                             `((%tail-call ,fn ,argv ,k)))
-                           `((%set! ,rval (%cons ,rval %nil))
-                             (%tail-call ,k ,rval #f)))))
-                  '()
-                  (cddr bind))))
    `(%lambda 
-      ,(add-tail-call k rval
+      ,(transform-to-cps
+         (add-tail-call k rval
            (flatten-binds
              `(%bind (,rval ,k)
                 (%set! ,k %k)
@@ -209,7 +233,59 @@
                                             (foldr add-opt
                                                    rest+bodyexprs
                                                    optionals)
-                                          requireds))))))))))
+                                            requireds)))))))))))
+; <= (%bind (var...)
+;      @before
+;      (%apply x y)
+;      @after))
+; => (%bind (var... k)
+;      @before
+;      (%set k (lambda _ @after))
+;      (%tail-call x y 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 k)))
+
+(define (transform-to-cps form)
+  (flatten-binds
+    `(%bind ,(second form)
+       ,@(foldr (lambda (subform after)
+                  (cond
+                    [(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)
+                                             (%set! ,(second subform) ,x)
+                                             ,@after)))
+                             (%tail-call ,(second (third subform))
+                                         ,(third (third subform))
+                                         ,k))))]
+                    [(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)
+                                         ,k))))]
+                    [else (cons subform after)]))
+                '()
+                (cddr form)))))
 
 ; (fn-expr arg-expr...)
 ;    => (let ([fn-var fn-expr] arg-var... argv)
@@ -225,13 +301,12 @@
   (define fn-var (gensym))
   (define arg-vars (map (lambda (x) (gensym)) arg-exprs))
   (define argv (gensym))
-  (simplify-form
   `(%bind (,fn-var ,@arg-vars ,argv)
      ,(simplify-form `(set! ,fn-var ,fn-expr))
      ,@(map (lambda (x y) (simplify-form `(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))))
+     (%apply ,fn-var ,argv)))
 
 (define (subst-var old-var new-var form)
   (define (recurse form)
@@ -273,24 +348,94 @@
     (case (car form)
       [(%bind)
        (let* ([bound-vars (second form)]
-              [subforms (append-map (lambda (form)
+              [subforms (append-map (lambda (new-form)
-                                      (let ([new-form (flatten-binds form)])
                                       (if (and (pair? new-form) (eq? (car new-form) '%bind))
                                         (begin
                                           (set! bound-vars (append bound-vars (second new-form)))
                                           (cddr new-form))
-                                          (list new-form))))
+                                        (list new-form)))
-                                    (cddr form))])
+                                    (map flatten-binds (cddr form)))])
          `(%bind ,bound-vars ,@subforms))]
       [(%if %tail-call %apply %set! %lambda %cons %car %cdr) form]
       [else (error "Unsimplified form:" form)])
     form))
 
-(define (compile form)
+(define (compile-lambda lambda-form)
   (flatten-binds
       (make-bindings-unique
-        (simplify-form `(lambda () ,form)))))
+        (simplify-form lambda-form))))
 
-(pretty-print (compile (read)))
+(define (compile form)
+  (compile-lambda `(lambda () ,form)))
+
+(define (free-variables form [input? #t] [output? #t])
+  (define (recurse form) (free-variables form input? output?))
+  (if (pair? form)
+    (case (car form)
+      [(%bind)
+       (remove* (second form)
+                (remove-duplicates (append-map recurse (cddr form))))]
+      [(%set!) (if output?
+                   (cons (second form) (recurse (third form)))
+                   (recurse (third form)))]
+      [(%if %tail-call %apply %lambda %cons %car %cdr)
+       (remove-duplicates (append-map recurse (cdr form)))]
+      [else
+       (error "Unsimplified form:" form)])
+    (if (and input?
+             (symbol? 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)
+  (if (pair? form)
+    (case (car form)
+      [(%bind)
+       (let ([ref-vars (remove-duplicates (append-map free-variables (cddr form)))])
+         `(%bind ,(filter (lambda (x) (memq x ref-vars)) (second form))
+            ,@(map reduce-variables (cddr form))))]
+      [(%if %tail-call %apply %set! %lambda %cons %car %cdr)
+         `(,(first form) ,@(map reduce-variables (cdr form)))]
+      [else (error "Unsimplified form:" form)])
+    form))
+
+; Don't set variables which won't be accessed later.
+(define (reduce-set! form)
+  (if (pair? form)
+    (case (car form)
+      [(%bind)
+       (let ([free-vars (free-variables form)])
+         `(%bind ,(second form)
+            ,@(foldr (lambda (subform after)
+                       (if (and (pair? subform)
+                                (eq? (first subform) '%set!)
+                                (not (memq (second subform) free-vars))
+                                (not (memq (second subform)
+                                           (append-map free-input-variables after))))
+                           after
+                           (cons subform after)))
+                     '()
+                     (map reduce-set! (cddr form)))))]
+      [(%if %tail-call %apply %set! %lambda %cons %car %cdr)
+         `(,(first form) ,@(map reduce-set! (cdr form)))]
+      [else (error "Unsimplified form:" form)])
+    form))
+
+(define (optimize form)
+  (reduce-variables
+    (reduce-set!
+      form)))
+
+;(void (trace optimize (trace compile-lambda (trace simplify-form `(lambda () ,(read))))))
+;(pretty-print (compile (read)))
+(pretty-print (optimize (compile (read))))
 
 ; vim:set sw=2 expandtab: