Thread context register (%ctx) through %apply and %call/cc forms.

Also remove top-level forms without side-effects during CPS transformation.

compiler.ss

@@ -15,6 +15,13 @@
          '()
          lst))
 
+(define (find x lst)
+  (let/cc return
+    (for ([i (in-naturals 0)]
+          [y (in-list lst)])
+      (when (eq? y x) (return i)))
+    #f))
+
 (define (simplify-form form)
   (if (pair? form)
       (case (car form)
@@ -302,7 +309,7 @@
             (values (cons (car arglist) req) opt rst)))
        (values '() '() #f)))
 
-(define (add-tail-call k rval form)
+(define (add-return ctx k rval form)
   (define k-argv (gensym))
   `(%bind (,@(second form) ,k-argv)
      ,@(foldr (lambda (subform after)
@@ -316,7 +323,7 @@
                         (eq? (first (third subform)) '%apply))
                    (let ([fn (second (third subform))]
                          [argv (third (third subform))])
-                     `((%tail-call ,fn ,argv ,k)))]
+                     `((%tail-call ,fn ,argv ,ctx ,k)))]
                   [(and (pair? subform)
                         (eq? (first subform) '%set!)
                         (eq? (second subform) rval)
@@ -324,7 +331,7 @@
                         (eq? (first (third subform)) '%call/cc))
                    (let ([fn (second (third subform))])
                      `((%set! ,k-argv (%cons %k %nil))
-                       (%tail-call ,fn ,k-argv %k)))]
+                       (%tail-call ,fn ,k-argv ,ctx %k)))]
                   [(and (pair? subform)
                         (eq? (first subform) '%tail-call))
                    `(,subform)]
@@ -332,21 +339,22 @@
                         (eq? (first subform) '%apply))
                    `((%tail-call ,(second subform)
                                  ,(third subform)
+                                 ,ctx
                                  ,k))]
                   [(and (pair? subform)
                         (eq? (first subform) '%set!)
                         (eq? (second subform) rval)
                         (eq? (third subform) '%void))
-                   `((%tail-call ,k %nil #f))]
+                   `((%tail-call ,k %nil #f #f))]
                   [(and (pair? subform)
                         (eq? (first subform) '%set!)
                         (eq? (second subform) rval))
                    `(,subform
                      (%set! ,rval (%cons ,rval %nil))
-                     (%tail-call ,k ,rval #f))]
+                     (%tail-call ,k ,rval #f #f))]
                   [else
                    `(,subform
-                     (%tail-call ,k %nil #f))]))
+                     (%tail-call ,k %nil #f #f))]))
                '()
                (cddr form))))
 
@@ -356,6 +364,7 @@
 
    (define-values (requireds optionals rest) (split-arglist arglist))
    (define argv-temp (gensym))
+   (define ctx (gensym))
    (define k (gensym))
    (define rval (gensym))
 
@@ -372,10 +381,11 @@
    (define rest+bodyexprs (if rest `(let ([,rest ,argv-temp]) ,@bodyexprs)
                                    `(begin ,@bodyexprs)))
    `(%lambda () ()
-      ,((compose transform-to-cps
+      ,((compose (lambda (x) (transform-to-cps ctx x))
-                 (lambda (x) (add-tail-call k rval x))
+                 (lambda (x) (add-return ctx k rval x))
                  flatten-binds)
-        `(%bind (,rval ,k)
+        `(%bind (,rval ,ctx ,k)
+           (%set! ,ctx %ctx)
            (%set! ,k %k)
            ,(simplify-form
               `(set! ,rval (let ([,argv-temp %argv])
@@ -387,9 +397,11 @@
 
 (define (narrow-binds simple-lambda-form)
   (define bind (fourth 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?)
@@ -404,6 +416,7 @@
     (filter-not captured-twice?
                 (filter-not at-top-level?
                             (second bind))))
+
   `(%lambda ,(second simple-lambda-form) ,(third simple-lambda-form)
      (%bind ,(remove* extra-binds (second bind))
        ,@(map (lambda (subform)
@@ -505,7 +518,7 @@
 ; => (%bind (var... k)
 ;      @before
 ;      (%set! k (lambda _ @after))
-;      (%tail-call x y k)))
+;      (%tail-call x y ctx k)))
 
 ; <= (%bind (var...)
 ;      @before
@@ -516,7 +529,7 @@
 ;      (%set! k (lambda (x)
 ;                 (%set! v x)
 ;                 @after))
-;      (%tail-call x y k)))
+;      (%tail-call x y ctx k)))
 
 ; <= (%bind (var...)
 ;      @before
@@ -526,9 +539,9 @@
 ;      @before
 ;      (%set! k (lambda _ @after))
 ;      (%set! k-argv (%cons k %nil))
-;      (%tail-call l k-argv k))
+;      (%tail-call l k-argv ctx k))
 
-(define (transform-to-cps bind)
+(define (transform-to-cps ctx bind)
   (define (cps-prepend subform after)
     (cond
       ; (%set! v (%apply x y))
@@ -545,6 +558,7 @@
                                ,@after)))
                (%tail-call ,(second (third subform))
                            ,(third (third subform))
+                           ,ctx
                            ,k))))]
       ; (%apply x y)
       [(and (pair? subform)
@@ -556,6 +570,7 @@
                                ,@after)))
                (%tail-call ,(second subform)
                            ,(third subform)
+                           ,ctx
                            ,k))))]
       ; (%set! v (%call/cc x))
       [(and (pair? subform)
@@ -573,6 +588,7 @@
                (%set! ,k-argv (%cons ,k %nil))
                (%tail-call ,(second (third subform))
                            ,k-argv
+                           ,ctx
                            ,k))))]
       ; (%call/cc x)
       [(and (pair? subform)
@@ -586,8 +602,14 @@
                (%set! ,k-argv (%cons ,k %nil))
                (%tail-call ,(second subform)
                            ,k-argv
+                           ,ctx
                            ,k))))]
-      [else (cons subform after)]))
+      ; keep all other forms with side-effects as-is
+      [(and (pair? subform)
+            (memq (first subform) '(%set! %set-box! %set-car! %set-cdr! %tail-call)))
+         (cons subform after)]
+      ; discard any form without side-effects
+      [else after]))
   (flatten-binds
     `(%bind ,(second bind)
        ,@(foldr cps-prepend '() (cddr bind)))))