Improve CPS.

Goodbye soup. Thanks to "Compiling with Continuations" by Appel.

1 files changed, 119 insertions, 105 deletions

diff --git a/csc/cps.csc b/csc/cps.csc
index 0c85707..cf4a26d 100644
--- a/csc/cps.csc
+++ b/csc/cps.csc
@@ -16,15 +16,13 @@
             if-consequent
             if-test
             if?
+            lambda-arguments
             lambda-body
-            lambda-case-alternate
-            lambda-case-arguments
-            lambda-case-body
-            lambda-case-gensyms
-            lambda-case-rest
+            lambda-rest
             lambda?
             letrec-expression
             letrec-gensyms
+            letrec-in-order?
             letrec-names
             letrec-values
             letrec?
@@ -39,7 +37,6 @@
             make-call
             make-constant
             make-lambda
-            make-lambda-case
             make-lexical-ref
             make-lexical-set
             make-sequence
@@ -47,14 +44,15 @@
             sequence-tail
             sequence?)
           (only (csc ir2)
+            make-apply
             make-atom
             make-branch
             make-call-closure
             make-closure
+            make-fix
             make-kargs
             make-klabel
             make-ktail
-            make-lambda-args
             make-update)
           (only (csc loop)
             loop
@@ -67,116 +65,132 @@
       (make-lexical-ref 'generated-symbol (gensym)))
 
 
-    (define (to-cps expr continuation add-continuation)
+    (define (collect-functions-and-variables expr)
+      (loop for name in (letrec-names expr)
+            for gensym in (letrec-gensyms expr)
+            for value in (letrec-values expr)
+            if (lambda? value)
+              collect (let ((continuation (new-ref)))
+                        (make-closure
+                          (make-lexical-ref name gensym)
+                          (cons continuation (lambda-arguments value))
+                          (lambda-rest value)
+                          (ir1->ir2
+                            (lambda-body value)
+                            (lambda (z)
+                              (make-apply continuation (list z))))))
+                      into functions
+            else
+              collect (make-lexical-ref name gensym) into variable-names
+              and collect value into variable-values
+            finally (return (values functions variable-names variable-values))))
+
+
+    (define (ir1->ir2 expr continuation)
       (cond
         ((or (constant? expr)
              (lexical-ref? expr)
              (library-ref? expr))
-          (make-atom expr continuation))
+          (continuation expr))
         ((lexical-set? expr)
-          (let ((ref (new-ref)))
-            (to-cps
-              (lexical-set-expression expr)
-              (add-continuation
-                (make-kargs (list ref)
-                  (make-update (lexical-set-ref expr) ref continuation)))
-              add-continuation)))
+          (ir1->ir2
+            (lexical-set-expression expr)
+            (lambda (val)
+              (make-update (lexical-set-ref expr) val (continuation (make-constant #f))))))
         ((library-define? expr)
-          (let ((ref (new-ref)))
-            (to-cps
-              (library-define-expression expr)
-              (add-continuation
-                (make-kargs (list ref)
-                  (make-update (library-define-ref expr) ref continuation)))
-              add-continuation)))
+          (ir1->ir2
+            (library-define-expression expr)
+            (lambda (val)
+              (make-update (library-define-ref expr) val (continuation (make-constant #f))))))
         ((define-syntax? expr)
           ; no-op
-          (make-atom (make-constant #f) continuation))
+          (continuation (make-constant #f)))
         ((if? expr)
-          (let* ((true-id (add-continuation
-                            (make-klabel
-                              (to-cps (if-consequent expr) continuation add-continuation))))
-                 (false-id (add-continuation
-                             (make-klabel
-                               (to-cps (if-alternate expr) continuation add-continuation))))
-                 (test-ref (new-ref))
-                 (branch-id (add-continuation
-                              (make-kargs (list test-ref)
-                                (make-branch test-ref true-id false-id)))))
-            (to-cps (if-test expr) branch-id add-continuation)))
+          (ir1->ir2
+            (if-test expr)
+            (lambda (val)
+              (define continuation-ref (new-ref))
+              (define result-ref (new-ref))
+              (make-fix
+                (list (make-closure continuation-ref (list result-ref) #f
+                        (continuation result-ref)))
+                (make-branch val
+                  (ir1->ir2
+                    (if-consequent expr)
+                    (lambda (result)
+                      (make-apply continuation-ref (list result))))
+                  (ir1->ir2
+                    (if-alternate expr)
+                    (lambda (result)
+                      (make-apply continuation-ref (list result)))))))))
         ((call? expr)
-          ; Technically the order of evaluation is unspecified. We evaluate
-          ; expressions left to right.
-          (loop with terms = (cons (call-procedure expr) (call-arguments expr))
-                with temps = (map (lambda (x) (new-ref)) terms)
-                with expr = (make-call-closure (car temps) (cdr temps) continuation)
-                for term in (reverse terms)
-                for temp in (reverse temps)
-                do (set! expr (to-cps term
-                                (add-continuation
-                                  (make-kargs (list temp) expr))
-                                add-continuation))
-                finally (return expr)))
+          (let ((return-address (new-ref))
+                (result (new-ref)))
+            (make-fix
+              (list (make-closure return-address (list result) #f (continuation result)))
+              (ir1->ir2
+                (call-procedure expr)
+                (lambda (f)
+                  ; Technically the order of evaluation is unspecified.
+                  ; We evaluate expressions left to right.
+                  ;
+                  ; I would use the loop macro, but it mutates the loop
+                  ; variables which plays badly with building a lambda.
+                  (let loop ((args (reverse (call-arguments expr)))
+                             (exprs (lambda (vals)
+                                      (make-apply f (cons return-address (reverse vals))))))
+                    (if (null? args)
+                      (exprs '())
+                      (loop (cdr args)
+                            (lambda (vals)
+                              (ir1->ir2
+                                (car args)
+                                (lambda (val)
+                                  (exprs (cons val vals)))))))))))))
         ((sequence? expr)
-          (to-cps
+          (ir1->ir2
             (sequence-head expr)
-            (add-continuation
-              (make-klabel
-                (to-cps (sequence-tail expr) continuation add-continuation)))
-            add-continuation))
+            (lambda (x)
+              (ir1->ir2
+                (sequence-tail expr)
+                continuation))))
         ((lambda? expr)
-          (make-closure
-            (loop with tail = (add-continuation (make-ktail))
-                  for lambda-case = (lambda-body expr) then (lambda-case-alternate lambda-case)
-                  while lambda-case
-                  collect (let ((args (lambda-case-arguments lambda-case))
-                                (rest (lambda-case-rest lambda-case)))
-                            (make-lambda-args
-                              (length args)
-                              (not (not rest))
-                              (add-continuation
-                                (make-kargs (map make-lexical-ref
-                                              (append args (list rest))
-                                              (lambda-case-gensyms lambda-case))
-                                  (to-cps (lambda-case-body lambda-case) tail add-continuation))))))
-            continuation))
+          (let ((f (new-ref))
+                (k (new-ref)))
+            (make-fix
+              (list
+                (make-closure f (cons k (lambda-arguments expr)) (lambda-rest expr)
+                  (ir1->ir2
+                    (lambda-body expr)
+                    (lambda (ret)
+                      (make-apply k (list ret))))))
+              (continuation f))))
         ((letrec? expr)
-          ; We re-write a letrec into a corresponding lambda form.
-          (let ((names (letrec-names expr))
-                (gensyms (letrec-gensyms expr)))
-            (to-cps
-              (make-call
-                (make-lambda
-                  (make-lambda-case
-                    names
-                    #f
-                    gensyms
-                    (make-sequence
-                      (loop for name in names
-                            for gensym in gensyms
-                            for value in (letrec-values expr)
-                            for set = (make-lexical-set (make-lexical-ref name gensym) value)
-                            for body = set then (make-sequence body set)
-                            finally (return body))
+          (let-values (((functions variable-names variable-values) (collect-functions-and-variables expr)))
+            (make-fix functions
+              (ir1->ir2
+                ; We re-write a letrec into a corresponding lambda form.
+                (if (letrec-in-order? expr)
+                  (loop for name in (reverse variable-names)
+                        for value in (reverse variable-values)
+                        for expr = (make-call
+                                     (make-lambda
+                                       (list name)
+                                       #f
+                                       (letrec-expression expr))
+                                     (list value))
+                                 then (make-call
+                                        (make-lambda
+                                          (list name)
+                                          #f
+                                          expr)
+                                        (list value))
+                        finally (return expr))
+                  (make-call
+                    (make-lambda
+                      variable-names
+                      #f
                       (letrec-expression expr))
-                    #f))
-                (map (lambda (x) (make-constant #f)) names))
-              continuation
-              add-continuation)))
-        (else (error "unexpected type in to-cps" expr))))
-
-
-    ; Returns a map from integers to CPS continuations.
-    ; By convention the continuation at key 0 is the entrypoint.
-    (define (ir1->ir2 program)
-      (define current-continuation-id 0)
-      (define soup (make-map (lambda (x) x) <))
-      (define (add-continuation continuation)
-        (set! current-continuation-id (+ 1 current-continuation-id))
-        (set! soup (insert soup
-                     current-continuation-id
-                     continuation))
-        current-continuation-id)
-      (define ktail (add-continuation (make-ktail)))
-      (define entrypoint (to-cps program ktail add-continuation))
-      (insert soup 0 (make-klabel entrypoint)))))
+                    variable-values))
+                continuation))))
+        (else (error "unexpected type in ir1->ir2" expr))))))