Add update conversion.

Now variables can't be updated after they're created, so they can be
freely copied into closures.

1 files changed, 195 insertions, 38 deletions

diff --git a/csc/cps.csc b/csc/cps.csc
index b6decbb..cf46e05 100644
--- a/csc/cps.csc
+++ b/csc/cps.csc
@@ -2,9 +2,13 @@
   (export
     ir1->ir2)
   (import (scheme base)
-          (only (csc gensym) gensym)
+          (only (csc gensym)
+            gensym
+            gensym->int)
           (only (csc hash-map)
             insert
+            key-not-found-error?
+            lookup
             make-map
             merge)
           (only (csc ir1)
@@ -20,7 +24,11 @@
             constant?
             if?
             lambda?
+            letrec-gensyms
+            letrec-names
+            letrec-values
             letrec?
+            lexical-ref-gensym
             lexical-ref?
             lexical-set?
             library-define?
@@ -33,6 +41,14 @@
             make-sequence
             sequence?)
           (only (csc ir2)
+            %apply
+            %branch
+            %fix
+            %primitive
+            closure-arguments
+            closure-body
+            closure-name
+            closure-rest
             make-apply
             make-atom
             make-branch
@@ -42,56 +58,71 @@
             make-kargs
             make-klabel
             make-ktail
-            make-update)
+            make-primitive)
           (only (csc loop)
             loop
             return)
-          (only (csc match) match))
+          (only (csc match)
+            define-match-record-type
+            match))
   (begin
 
 
+    ; Update is a CPS expression that is used internally as part of
+    ; CPS conversion.
+    ; Update expressions are then removed by box-conversion.
+    (define-match-record-type <update>
+      (make-update ref atom continuation)
+      update?
+      %update
+      (ref update-ref)
+      (atom update-atom)
+      (continuation update-continuation))
+
+
     (define (new-ref)
       (make-lexical-ref 'generated-symbol (gensym)))
 
 
     (define (collect-functions-and-variables expr)
-      (match expr
-        ((% %letrec _ names gensyms vals _)
-          (loop for name in names
-                for gensym in gensyms
-                for value in vals
-                if (lambda? value)
-                  collect (match value
-                            ((% %lambda args rest body)
-                              (define continuation (new-ref))
-                              (make-closure
-                                (make-lexical-ref name gensym)
-                                (cons continuation args)
-                                rest
-                                (ir1->ir2
-                                  body
-                                  (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))))))
+      (let ((names (letrec-names expr))
+            (gensyms (letrec-gensyms expr))
+            (vals (letrec-values expr)))
+        (loop for name in names
+              for gensym in gensyms
+              for value in vals
+              if (lambda? value)
+                collect (match value
+                          ((% %lambda args rest body)
+                            (define continuation (new-ref))
+                            (make-closure
+                              (make-lexical-ref name gensym)
+                              (cons continuation args)
+                              rest
+                              (to-cps
+                                body
+                                (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)
+    (define (to-cps expr continuation)
       (match expr
         (_ (when (or (constant? expr)
                      (lexical-ref? expr)
                      (library-ref? expr)))
           (continuation expr))
         ((% %lexical-set ref arg)
-          (ir1->ir2
+          (to-cps
             arg
             (lambda (val)
               (make-update ref val (continuation (make-constant #f))))))
         ((% %library-define ref arg)
-          (ir1->ir2
+          (to-cps
             arg
             (lambda (val)
               (make-update ref val (continuation (make-constant #f))))))
@@ -99,7 +130,7 @@
           ; no-op
           (continuation (make-constant #f)))
         ((% %if test consequent alternate)
-          (ir1->ir2
+          (to-cps
             test
             (lambda (val)
               (define continuation-ref (new-ref))
@@ -108,11 +139,11 @@
                 (list (make-closure continuation-ref (list result-ref) #f
                         (continuation result-ref)))
                 (make-branch val
-                  (ir1->ir2
+                  (to-cps
                     consequent
                     (lambda (result)
                       (make-apply continuation-ref (list result))))
-                  (ir1->ir2
+                  (to-cps
                     alternate
                     (lambda (result)
                       (make-apply continuation-ref (list result)))))))))
@@ -121,7 +152,7 @@
           (define result (new-ref))
           (make-fix
             (list (make-closure return-address (list result) #f (continuation result)))
-            (ir1->ir2
+            (to-cps
               proc
               (lambda (f)
                 ; Technically the order of evaluation is unspecified.
@@ -136,15 +167,15 @@
                     (exprs '())
                     (loop (cdr args*)
                           (lambda (vals)
-                            (ir1->ir2
+                            (to-cps
                               (car args*)
                               (lambda (val)
                                 (exprs (cons val vals))))))))))))
         ((% %sequence head tail)
-          (ir1->ir2
+          (to-cps
             head
             (lambda (x)
-              (ir1->ir2
+              (to-cps
                 tail
                 continuation))))
         ((% %lambda args rest body)
@@ -153,7 +184,7 @@
           (make-fix
             (list
               (make-closure f (cons k args) rest
-                (ir1->ir2
+                (to-cps
                   body
                   (lambda (ret)
                     (make-apply k (list ret))))))
@@ -161,7 +192,7 @@
         ((% %letrec in-order? _ _ _ body)
           (define-values (functions variable-names variable-values) (collect-functions-and-variables expr))
           (make-fix functions
-            (ir1->ir2
+            (to-cps
               ; We re-write a letrec into a corresponding lambda form.
               (if in-order?
                 (loop for name in (reverse variable-names)
@@ -186,4 +217,130 @@
                     body)
                   variable-values))
               continuation)))
-        (_ (error "unexpected type in ir1->ir2" expr))))))
+        (_ (error "unexpected type in to-cps" expr))))
+
+
+    (define (make-ref-map)
+      (make-map
+        (lambda (ref)
+          (gensym->int (lexical-ref-gensym ref)))
+        (lambda (x y) (< (gensym->int (lexical-ref-gensym x))
+                         (gensym->int (lexical-ref-gensym y))))))
+
+
+    (define (get-boxed expr)
+      (match expr
+        ((% %update ref _ continuation)
+          (define m (get-boxed continuation))
+          (when (lexical-ref? ref)
+            (set! m (insert m ref #t)))
+          m)
+        ((% %primitive _ _ _ continuation)
+          (get-boxed continuation))
+        ((% %branch _ true false)
+          (merge
+            (get-boxed true)
+            (get-boxed false)))
+        ((% %apply proc args)
+          (make-ref-map))
+        ((% %fix funs body)
+          (loop with m = (get-boxed body)
+                for fun in funs
+                do (set! m (merge m (get-boxed (closure-body fun))))
+                finally (return m)))
+        (_ (error "Unexpected form in get-boxed"))))
+
+
+    (define (all-closure-args fun)
+      (define args (closure-arguments fun))
+      (define rest (closure-rest fun))
+      (when rest
+        (set! args (cons rest args)))
+      args)
+
+
+    ; Rewrites the given expression to have no more <update> forms.
+    (define (box-conversion expr)
+      (define boxed-refs (get-boxed expr))
+      (define (boxed? ref)
+        (or (library-ref? ref) ; globals are always boxed
+            (and (lexical-ref? ref)
+                 (guard (e ((key-not-found-error? e) #f))
+                   (lookup boxed-refs ref)))))
+      (define (convert-arg-list args)
+        (define boxed-args (loop for arg in args
+                                 if (boxed? arg)
+                                   collect arg))
+        (define vars (loop for x in boxed-args
+                           collect (new-ref)))
+        (define new-args (loop with v* = vars
+                               for arg in args
+                               collect (if (boxed? arg)
+                                         (car v*)
+                                         arg)
+                               if (boxed? arg)
+                                 do (set! v* (cdr v*))))
+        (values new-args boxed-args vars))
+      (let convert ((expr expr))
+        (match expr
+          ((% %update ref atom continuation)
+            (make-primitive 'poke (list atom ref (make-constant 0)) '() (convert continuation)))
+          ((% %primitive op args res continuation)
+            ; Note that no reference in res can be boxed.
+            (define-values (new-args boxed-args vars) (convert-arg-list args))
+            (define new-expr (make-primitive op new-args res (convert continuation)))
+            (loop for arg in boxed-args
+                  for var in vars
+                  do (set! new-expr (make-primitive 'peek (list arg (make-constant 0)) (list var)
+                                      new-expr))
+                  finally (return new-expr)))
+          ((% %branch atom true false)
+            (if (boxed? atom)
+              (let ((var (new-ref)))
+                (make-primitive 'peek (list atom (make-constant 0)) (list var)
+                                (make-branch var (convert true) (convert false))))
+              (make-branch atom (convert true) (convert false))))
+          ((% %apply proc args)
+            (define-values (new-params boxed-params vars) (convert-arg-list (cons proc args)))
+            (define new-expr (make-apply (car new-params) (cdr new-params)))
+            (loop for p in boxed-params
+                  for var in vars
+                  do (set! new-expr (make-primitive 'peek (list p (make-constant 0)) (list var)
+                                      new-expr))
+                  finally (return new-expr)))
+          ((% %fix funs body)
+            (define-values (new-names boxed-names temp-names) (convert-arg-list (loop for fun in funs
+                                                                                      collect (closure-name fun))))
+            (define new-funs (loop for fun in funs
+                                   for new-name in new-names
+                                   for rest = (closure-rest fun)
+                                   collect (let-values (((new-args boxed-args temp-args) (convert-arg-list (all-closure-args fun))))
+                                             (make-closure
+                                               new-name
+                                               (if rest
+                                                 (cdr new-args)
+                                                 new-args)
+                                               (if rest
+                                                 (car new-args)
+                                                 #f)
+                                               (let ((new-expr (convert (closure-body fun))))
+                                                 (loop for arg in boxed-args
+                                                       for var in temp-args
+                                                       do (set! new-expr (make-primitive 'alloc (list (make-constant 1)) (list arg)
+                                                                           (make-primitive 'poke (list var arg (make-constant 0)) '()
+                                                                             new-expr)))
+                                                       finally (return new-expr)))))))
+            (define new-body (convert body))
+            (loop for name in boxed-names
+                  for var in temp-names
+                  do (set! new-body (make-primitive 'poke (list var name (make-constant 0)) '()
+                                      new-body)))
+            (define new-expr (make-fix new-funs new-body))
+            (loop for name in boxed-names
+                  do (set! new-expr (make-primitive 'alloc (list (make-constant 1)) (list name)
+                                      new-expr))
+                  finally (return expr))))))
+
+
+    (define (ir1->ir2 expr continuation)
+      (box-conversion (to-cps expr continuation)))))