Implement conversion of shared variables to boxes.

Also, narrow bindings to their minimal necessary scope. This reduces
the number of variables which must be considered 'shared'.

compiler.ss

@@ -6,27 +6,103 @@
     (pretty-print x)
     x))
 
+(define (subst old new lst)
+  (foldr (lambda (x rst)
+           (cons (if (eq? x old)
+                     new
+                     x)
+                 rst))
+         '()
+         lst))
+
 (define (simplify-form form)
   (if (pair? form)
-    (case (car form)
+      (case (car form)
-      [(let) (simplify-let form)]
+        [(let) (simplify-let form)]
-      [(let*) (simplify-let* form)]
+        [(let*) (simplify-let* form)]
-      [(letrec) (simplify-letrec form)]
+        [(letrec) (simplify-letrec form)]
-      [(if) (simplify-if form)]
+        [(if) (simplify-if form)]
-      [(lambda) (simplify-lambda form)]
+        [(lambda) (simplify-lambda form)]
-      [(begin) (simplify-form `(let () ,@(cdr form)))]
+        [(begin) (simplify-form `(let () ,@(cdr form)))]
-      [(set!) (simplify-set! form)]
+        [(set!) (simplify-set! form)]
-      [(let/cc) (simplify-form `(call/cc (lambda (,(second form)) ,@(cddr form))))]
+        [(let/cc) (simplify-form `(call/cc (lambda (,(second form)) ,@(cddr form))))]
-      [(car cdr cons call/cc)
+        [(car cdr cons call/cc)
-       (simplify-primitive (case (first form)
+         (simplify-primitive (case (first form)
-                             [(car)     '%car]
+                               [(make-box) '%make-box]
-                             [(cdr)     '%cdr]
+                               [(set-box!) '%set-box!]
-                             [(cons)    '%cons]
+                               [(unbox)    '%unbox]
-                             [(call/cc) '%call/cc])
+                               [(cons)     '%cons]
-                           (cdr form))]
+                               [(set-car!) '%set-car!]
-      [(%bind %if %tail-call %apply %lambda %set! %cons %car %cdr %call/cc quote) form]
+                               [(car)      '%car]
-      [else (simplify-funcall form)])
+                               [(set-cdr!) '%set-cdr!]
-    (if (eq? form '()) '%nil form)))
+                               [(cdr)      '%cdr]
+                               [(call/cc)  '%call/cc]
+                               )
+                             (cdr form))]
+        [(%bind %if %set! %lambda quote
+          %tail-call %apply %call/cc
+          %cons %set-car! %car %set-cdr! %cdr
+          %make-box %set-box! %unbox)
+         form]
+        [else (simplify-funcall form)])
+      (if (eq? form '()) '%nil form)))
+
+(define (form-sets? form var [call-may-set? #t])
+  (define (recurse simple-form)
+    (if (pair? simple-form)
+        (case (car simple-form)
+          [(%bind) (and (not (memq var (second simple-form)))
+                        (ormap recurse (cddr simple-form)))]
+          [(%set!) (eq? (second simple-form) var)]
+          [(%tail-call %apply %call/cc) call-may-set?]
+          [(quote %lambda %if
+            %make-box %set-box! %unbox
+            %cons %set-car! %car %set-cdr! %cdr)
+           #f]
+          [else (error "Invalid simple form:" simple-form)])
+        #f))
+  (recurse (simplify-form form)))
+
+(define (form-uses? form var [call-may-use? #t] [descend? #t])
+  (define (recurse simple-form)
+    (if (pair? simple-form)
+        (case (car simple-form)
+          [(%bind) (and (not (memq var (second simple-form)))
+                        (ormap recurse (cddr simple-form)))]
+          [(%set!) (recurse (third simple-form))]
+          [(%tail-call %apply %call/cc) (and call-may-use? #t)]
+          [(%if %make-box %set-box! %unbox
+            %cons %set-car! %car %set-cdr! %cdr)
+           (ormap recurse (cdr simple-form))]
+          [(%lambda) (and descend?
+                          (memq var (free-variables simple-form))
+                          #t)]
+          [(quote) #f]
+          [else (error "Invalid simple form:" simple-form)])
+        (eq? simple-form var)))
+  (recurse (simplify-form form)))
+
+(define (form-captures? form var [input? #t] [output? #t])
+  (define (recurse simple-form)
+    (if (pair? simple-form)
+        (case (car simple-form)
+          [(%bind) (and (not (memq var (second simple-form)))
+                        (ormap recurse (cddr simple-form)))]
+          [(%set!) (recurse (third simple-form))]
+          [(%tail-call %apply %call/cc %if quote
+            %make-box %set-box! %unbox
+            %cons %set-car! %car %set-cdr! %cdr)
+           #f]
+          [(%lambda) (and (memq var (free-variables simple-form input? output?)) #t)]
+          [else (error "Invalid simple form:" simple-form)])
+        #f))
+  (recurse (simplify-form form)))
+
+(define (form-captures-input? form var)
+  (form-captures? form var #t #f))
+
+(define (form-captures-output? form var)
+  (form-captures? form var #f #t))
 
 (define (simplify-set! form)
   (let ([value-form (simplify-form (third form))])
@@ -37,7 +113,7 @@
                         (cond
                           [(pair? after) (cons subform after)]
                           [(or (not (pair? subform))
-                               (memq (first subform) '(%apply %call/cc %car %cdr %cons %bind %if quote)))
+                               (memq (first subform) '(%apply %call/cc %car %cdr %cons %bind %if %unbox quote)))
                            `((set! ,(second form) ,subform))]
                           [(and (pair? subform) (eq? (first subform) '%tail-call))
                            `(,subform)] ; The %set! wouldn't be executed anyway.
@@ -254,18 +330,135 @@
    (define rest+bodyexprs (if rest `(let ([,rest ,argv-temp]) ,@bodyexprs)
                                    `(begin ,@bodyexprs)))
    `(%lambda 
-      ,(transform-to-cps
+      ,((compose transform-to-cps
-         (add-tail-call k rval
+                 (lambda (x) (add-tail-call k rval x))
-           (flatten-binds
+                 flatten-binds)
-             `(%bind (,rval ,k)
+        `(%bind (,rval ,k)
-                (%set! ,k %k)
+           (%set! ,k %k)
-                ,(simplify-form
+           ,(simplify-form
-                   `(set! ,rval (let ([,argv-temp %argv])
+              `(set! ,rval (let ([,argv-temp %argv])
-                                  ,(foldr add-req
+                             ,(foldr add-req
-                                          (foldr add-opt
+                                     (foldr add-opt
-                                                 rest+bodyexprs
+                                            rest+bodyexprs
-                                                 optionals)
+                                            optionals)
-                                          requireds))))))))))
+                                     requireds))))))))
+
+(define (narrow-binds simple-lambda-form)
+  (define bind (second simple-lambda-form))
+  (define (at-top-level? var)
+    (or (ormap (lambda (x) (form-sets? x var #f)) (cddr bind))
+        (ormap (lambda (x) (form-uses? x var #f #f)) (cddr bind))))
+  (define (captured-twice? var)
+    (let/cc return
+      (foldl (lambda (subform once?)
+               (if (form-captures? subform var)
+                   (if once? (return #t) #t)
+                   once?))
+             (at-top-level? var)
+             (cddr bind))
+      #f))
+
+  (define extra-binds
+    (filter-not captured-twice?
+                (filter-not at-top-level?
+                            (second bind))))
+   `(%lambda
+     (%bind ,(remove* extra-binds (second bind))
+       ,@(map (lambda (subform)
+                (foldl (lambda (var subform)
+                         (if (and (pair? subform)
+                                  (eq? (first subform) '%set!)
+                                  (pair? (third subform))
+                                  (eq? (first (third subform)) '%lambda))
+                             (let* ([dest (second subform)]
+                                    [value (third subform)]
+                                    [bind (second value)])
+                               `(%set! ,dest ,(narrow-binds
+                                                `(%lambda
+                                                   ,(if (form-captures? value var)
+                                                        `(%bind (,@(second bind) ,var)
+                                                           ,@(cddr bind))
+                                                        bind)))))
+                             subform))
+                         subform
+                         extra-binds))
+              (cddr bind)))))
+
+(define (promote-to-box var form)
+  (define (recurse subform) (promote-to-box var subform))
+  (if (pair? form)
+    (case (car form)
+      [(%bind)
+       (flatten-binds
+        `(%bind ,(second form)
+           ,@(if (memq var (second form))
+                 `((%set! ,var (%make-box %undef)))
+                 '())
+           ,@(map recurse (cddr form))))]
+      [(%lambda)
+       `(%lambda ,(recurse (second form)))]
+      [(%set!)
+       (let ([value (recurse (third form))]
+             [kind (if (eq? (second form) var) '%set-box! '%set!)])
+         ; If value is (%bind), could only come from clause below.
+         (if (and (pair? value) (eq? (first value) '%bind))
+             (if (and (pair? (fourth value)) (not (eq? kind '%set!)))
+                 (let ([tmp (gensym)])
+                   `(%bind (,@(second value) ,tmp)
+                      ,(third value)
+                      (%set! ,tmp ,(fourth value))
+                      (,kind ,(second form) ,tmp)))
+                 `(%bind ,(second value)
+                    ,(third value)
+                    (,kind ,(second form) ,(fourth value))))
+             (if (and (pair? value) (not (eq? kind '%set!)))
+                 (let ([tmp (gensym)])
+                   `(%bind (,tmp)
+                      (%set! ,tmp ,value)
+                      (,kind ,(second form) ,tmp)))
+                 `(,kind ,(second form) ,value))))]
+      [(%if %tail-call %cons %car %cdr %call/cc %make-box %set-box! %unbox)
+       (if (memq var (cdr form))
+           (let ([tmp (gensym)])
+             `(%bind (,tmp)
+                (%set! ,tmp (%unbox ,var))
+                (,(first form) ,@(subst var tmp (cdr form)))))
+           form)]
+      [(quote) form]
+      [else (error "Unsimplified form:" form)])
+    (if (eq? form var)
+        `(%unbox ,form)
+        form)))
+
+; form needs to be flattened (%bind ...)
+(define (is-shared-var? var form)
+  (define captured-input?  (ormap (lambda (f) (form-captures-input?  f var)) (cddr form)))
+  (define captured-output? (ormap (lambda (f) (form-captures-output? f var)) (cddr form)))
+  (define (set-after-first-use?)
+    (let/cc return
+      (foldr (lambda (subform set-after?)
+               (if (or set-after? (form-sets? subform var captured-output?))
+                   (if (form-uses? subform var captured-input?)
+                       (return #t)
+                       #t)
+                   #f))
+             #f
+             (cddr form))
+      #f))
+  (and (or captured-input?
+           captured-output?)
+       (set-after-first-use?)))
+
+(define (promote-shared-vars simple-lambda-form)
+  (define bind (second simple-lambda-form))
+  `(%lambda
+     ,(foldl (lambda (var frm)
+               (if (is-shared-var? var frm)
+                   (promote-to-box var frm)
+                   frm))
+             bind
+             (second bind))))
+
 ; <= (%bind (var...)
 ;      @before
 ;      (%apply x y)
@@ -393,7 +586,7 @@
            form
            `(%bind ,(second form) ,@(map recurse (cddr form))))]
       [(quote) form]
-      [(%if %tail-call %apply %lambda %set! %cons %car %cdr %call/cc)
+      [(%if %tail-call %apply %lambda %set! %cons %car %cdr %call/cc %make-box %set-box! %unbox)
        `(,(first form) ,@(map recurse (cdr form)))]
       [else (error "Unsimplified form:" form)])
     (if (eq? form old-var)
@@ -414,12 +607,11 @@
                             frm
                             (map cons (second form) new-vars)))
                    (cddr form))))]
-      [(%if %tail-call %apply %lambda %set! %cons %car %cdr %call/cc quote) form]
+      [(%if %tail-call %apply %lambda %set! %cons %car %cdr %call/cc %make-box %set-box! %unbox quote) form]
       [else (error "Unsimplified form:" form)])
     form))
 
 (define (flatten-binds form)
-  ;(pretty-print form) (write-char #\Newline)
   (if (pair? form)
     (case (car form)
       [(%bind)
@@ -432,7 +624,7 @@
                                         (list new-form)))
                                     (map flatten-binds (cddr form)))])
          `(%bind ,bound-vars ,@subforms))]
-      [(%if %tail-call %apply %set! %lambda %cons %car %cdr %call/cc quote) form]
+      [(%if %tail-call %apply %set! %lambda %cons %car %cdr %call/cc %make-box %set-box! %unbox quote) form]
       [else (error "Unsimplified form:" form)])
     form))
 
@@ -447,13 +639,13 @@
                    (cons (second form) (recurse (third form)))
                    (recurse (third form)))]
       [(quote) '()]
-      [(%if %tail-call %apply %lambda %cons %car %cdr %call/cc)
+      [(%if %tail-call %apply %lambda %cons %car %cdr %call/cc %make-box %set-box! %unbox)
        (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))))
+             (not (memq form '(%nil %undef %self %argv %ctx %k))))
         (list form)
         '())))
 
@@ -472,7 +664,7 @@
          `(%bind ,(filter (lambda (x) (memq x ref-vars)) (second form))
             ,@(map reduce-variables (cddr form))))]
       [(quote) form]
-      [(%if %tail-call %apply %set! %lambda %cons %car %cdr %call/cc)
+      [(%if %tail-call %set! %lambda %cons %car %cdr %call/cc %make-box %set-box! %unbox)
          `(,(first form) ,@(map reduce-variables (cdr form)))]
       [else (error "Unsimplified form:" form)])
     form))
@@ -482,30 +674,34 @@
   (if (pair? form)
     (case (car form)
       [(%bind)
-       (let ([free-vars (free-variables form)])
+       (let ([bound-vars (second form)])
-         `(%bind ,(second form)
+         `(%bind ,bound-vars
             ,@(foldr (lambda (subform after)
                        (if (and (pair? subform)
                                 (eq? (first subform) '%set!)
-                                (not (memq (second subform) free-vars))
+                                (memq (second subform) bound-vars)
                                 (not (memq (second subform)
                                            (append-map free-input-variables after))))
                            after
                            (cons subform after)))
                      '()
                      (map reduce-set! (cddr form)))))]
-      [(%if %tail-call %apply %cons %car %cdr %call/cc quote) form]
+      [(%if %tail-call %cons %car %cdr %call/cc %make-box %set-box! %unbox quote) form]
       [(%set! %lambda) `(,(first form) ,@(map reduce-set! (cdr form)))]
       [else (error "Unsimplified form:" form)])
     form))
 
-(define (optimize form)
+(define (simplify-toplevel-lambda form)
+  (promote-shared-vars
+    (narrow-binds
+      (simplify-lambda form))))
+
+(define (optimize-simplified-lambda form)
   (reduce-variables
     (reduce-set!
       form)))
 
-(pretty-print (optimize (simplify-form `(lambda () ,(read)))))
+(pretty-print (optimize-simplified-lambda (simplify-toplevel-lambda `(lambda () ,(read)))))
-;(pretty-print (simplify-form (read)))
+;(pretty-print (simplify-toplevel-lambda `(lambda () ,(read))))
-;(pretty-print (optimize (trace simplify-form (read))))
 
 ; vim:set sw=2 expandtab: