better restoring continuation

src/continuation.c

@@ -5,88 +5,166 @@
 // Marks a continuation chain as immutable.
 // Unwinding through an immutable continuation chain requires
 // copying each continuation as needed.
-void continuation_set_immutable(obj cc) {
-    while (Mcontinuationp(cc)) {
-        Mcontinuation_immutablep(cc) = 1;
-        cc = Mcontinuation_prev(cc);
+void continuation_set_immutable(obj k) {
+    while (Mcontinuationp(k) && !Mcontinuation_immutablep(k)) {
+        Mcontinuation_immutablep(k) = 1;
+        k = Mcontinuation_prev(k);
     }
 }
 
 // Safely returns a mutable version of the current continuation.
 // If the continuation chain is immutable, a copy is made.
 // Otherwise, the argument is returned.
-obj continuation_mutable(obj cc) {
-    if (Mcontinuation_immutablep(cc)) {
+obj continuation_mutable(obj k) {
+    if (Mcontinuation_immutablep(k)) {
         // immutable => need to make a copy
-        obj k;
+        obj k2;
 
-        switch (Mcontinuation_type(cc)) {
+        switch (Mcontinuation_type(k)) {
         // bottom of continuation chain (immutable, by default)
         case NULL_CONT_TYPE:
-            k = cc;
+            k2 = k;
             break;
 
         // application
         case APP_CONT_TYPE:
-            k = Mapp_continuation(Mcontinuation_prev(cc), Mcontinuation_env(cc), NULL);
-            Mcontinuation_app_hd(k) = Mcontinuation_app_hd(cc);
-            Mcontinuation_app_tl(k) = Mcontinuation_app_tl(cc);
+            k2 = Mapp_continuation(Mcontinuation_prev(k), Mcontinuation_env(k), NULL);
+            Mcontinuation_app_hd(k2) = Mcontinuation_app_hd(k);
+            Mcontinuation_app_tl(k2) = Mcontinuation_app_tl(k);
             break;
 
         // if expressions
         case COND_CONT_TYPE:
-            k = Mcond_continuation(
-                Mcontinuation_prev(cc),
-                Mcontinuation_env(cc),
-                Mcontinuation_cond_ift(cc),
-                Mcontinuation_cond_iff(cc)
+            k2 = Mcond_continuation(
+                Mcontinuation_prev(k),
+                Mcontinuation_env(k),
+                Mcontinuation_cond_ift(k),
+                Mcontinuation_cond_iff(k)
             );
             break;
 
         // begin expressions
         case SEQ_CONT_TYPE:
-            k = Mseq_continuation(
-                Mcontinuation_prev(cc),
-                Mcontinuation_env(cc),
-                Mcontinuation_seq_value(cc)
+            k2 = Mseq_continuation(
+                Mcontinuation_prev(k),
+                Mcontinuation_env(k),
+                Mcontinuation_seq_value(k)
             );
             break;
 
         // let expressions
         case LET_CONT_TYPE:
-            k = Mlet_continuation(
-                Mcontinuation_prev(cc),
-                Mcontinuation_env(cc),
-                Mcontinuation_let_bindings(cc),
-                Mcontinuation_let_body(cc)
+            k2 = Mlet_continuation(
+                Mcontinuation_prev(k),
+                Mcontinuation_env(k),
+                Mcontinuation_let_bindings(k),
+                Mcontinuation_let_body(k)
             );
-            Mcontinuation_let_env(k) = Mcontinuation_let_env(cc);
+            Mcontinuation_let_env(k2) = Mcontinuation_let_env(k);
             break;
 
         // set! expressions
         case SETB_CONT_TYPE:
-            k = Msetb_continuation(
-                Mcontinuation_prev(cc),
-                Mcontinuation_env(cc),
-                Mcontinuation_setb_name(cc)
+            k2 = Msetb_continuation(
+                Mcontinuation_prev(k),
+                Mcontinuation_env(k),
+                Mcontinuation_setb_name(k)
             );
             break;
 
         // call/cc expressions
         case CALLCC_CONT_TYPE:
-            return Mcallcc_continuation(Mcontinuation_prev(cc), Mcontinuation_env(cc));
+            k2 = Mcallcc_continuation(Mcontinuation_prev(k), Mcontinuation_env(k));
+            break;
 
         default:
-            minim_error1("continuation_pop", "unimplemented", Mfixnum(Mcontinuation_type(cc)));
+            minim_error1("continuation_pop", "unimplemented", Mfixnum(Mcontinuation_type(k)));
         }
 
-        return k;
+        return k2;
     } else {
         // mutable => return it
-        return cc;
+        return k;
+    }
+}
+
+// Length of a continuation chain.
+static uptr continuation_length(obj k) {
+    uptr l = 0;
+    for (; Mcontinuationp(k); k = Mcontinuation_prev(k), ++l);
+    return l;
+}
+
+// Extracts the tail of a continuation chain.
+static obj continuation_tail(obj k, iptr l) {
+    for (uptr i = 0; i < l; ++i, k = Mcontinuation_prev(k));
+    return k;
+}
+
+// Extracts the common tail of two continuation chains.
+static obj common_tail(obj k1, obj k2) {
+    uptr l1, l2;
+
+    // eliminate excess frames
+    l1 = continuation_length(k1);
+    l2 = continuation_length(k2);
+    if (l1 > l2) {
+        k1 = continuation_tail(k1, l1 - l2);
+    } else  if (l2 > l1) {
+        k2 = continuation_tail(k2, l2 - l1);
+    }
+
+    // unwind both until a common ancestor is found
+    while (Mcontinuationp(k1)) {
+        if (k1 == k2)
+            return k1;
+
+        k1 = Mcontinuation_prev(k1);
+        k2 = Mcontinuation_prev(k2);
     }
+
+    return k1;
 }
 
+// Restores a continuation.
+// The result is a new continuation chain formed by merging
+// the common ancestors of the continuation and current continuation.
+obj continuation_restore(obj cc, obj k) {
+    obj tl, re;
+
+    tl = common_tail(cc, k);
+    if (!Mcontinuationp(tl)) {
+        // fully unwound, just reinstate `k` fully
+        return k;
+    } else if (tl == k) {
+        // `tl` is `k` so just return it
+        return tl;
+    } else {
+        // found a common ancestor, reinstate `k` up to that point
+        // we can assume that `tl` is immutable since `k`
+        // must be captured by `call/cc` (and variants)
+        if (!Mcontinuation_immutablep(tl))
+            minim_error1("continuation_restore", "must be immutable", tl);
+
+        // walk back to the ancestor and track the frames
+        // must be at least one
+        re = Mcons(continuation_mutable(k), Mnull);
+        for (k = Mcontinuation_prev(k); k != tl; k = Mcontinuation_prev(k)) {
+            re = Mcons(continuation_mutable(k), re);
+        }
+
+        // fprintf(stderr, "reinstating %lu frames (%lu saved)\n", list_length(re), continuation_length(tl));
+
+        // reinstate the frames and link them
+        k = tl;
+        for (; !Mnullp(re); re = Mcdr(re)) {
+            Mcontinuation_prev(Mcar(re)) = k;
+            k = Mcar(re);
+        }
+
+        return k;
+    }
+}
 
 // For debugging
 void print_continuation(obj cc) {

src/eval.c

@@ -176,7 +176,7 @@ static obj eval_k(obj e, obj env, obj k) {
         k = Mcontinuation_prev(k);
         goto loop;
     } else if (Mcontinuationp(f)) {
-        k = f;
+        k = continuation_restore(k, f);
         goto do_k;
     } else {
         minim_error1("eval_expr", "application: not a procedure", x);
@@ -268,46 +268,42 @@ static obj eval_k(obj e, obj env, obj k) {
 
     // set! expressions
     case SETB_CONT_TYPE:
-        k = continuation_mutable(k);
         env = Mcontinuation_env(k);
+
+        // update binding
         e = env_find(env, Mcontinuation_setb_name(k));
-        if (Mfalsep(x)) {
+        if (Mfalsep(e)) {
             minim_error1("set!", "unbound variable", Mcontinuation_setb_name(k));
-        } else if (Munboundp(x)) {
-            minim_error1("set!", "uninitialized variable", Mcontinuation_setb_name(k));   
+        } else {
+            Mcdr(e) = x;
         }
 
-        // update binding, result is void
-        Mcdr(e) = x;
+        // result is void
         x = Mvoid;
         k = Mcontinuation_prev(k);
         goto do_k;
 
     // call/cc expressions
     case CALLCC_CONT_TYPE:
-        // actually do call/cc
-        k = Mcontinuation_prev(k);
+        env = Mcontinuation_env(k);
         if (Mprimp(x)) {
             iptr arity = Mprim_arity(x);
             if (arity < 0 ? arity == -1 : arity != 1)
                 minim_error1("call/cc", "expected a procedure of at least 1 argument", x);
-
-            x = do_prim(x, Mlist1(k));
-            goto do_k;
         } else if (Mclosurep(x)) {
             iptr arity = Mclosure_arity(x);
             if (arity < 0 ? arity == -1 : arity != 1)
                 minim_error1("call/cc", "expected a procedure of at least 1 argument", x);
-
-            e = Mclosure_body(x);
-            env = do_closure(x, Mlist1(k));
-            goto loop;
         } else if (Mcontinuationp(x)) {
-            goto do_k;
+            // do nothing
         } else {
             minim_error1("call/cc", "expected a procedure", x);
         }
 
+        f = x;
+        args = Mlist1(Mcontinuation_prev(k));
+        goto do_app;
+
     // unknown
     default:
         minim_error1("eval_expr", "unimplemented", Mfixnum(Mcontinuation_type(k)));

src/minim.h

@@ -347,17 +347,17 @@ obj Mlength(obj x);
 // Marks a continuation chain as immutable.
 // Unwinding through an immutable continuation chain requires
 // copying each continuation as needed.
-void continuation_set_immutable(obj cc);
+void continuation_set_immutable(obj k);
 
 // Safely returns a mutable version of the current continuation.
 // If the continuation chain is immutable, a copy is made.
 // Otherwise, the argument is returned.
-obj continuation_mutable(obj x);
+obj continuation_mutable(obj k);
 
 // Restores a continuation.
 // The result is a new continuation chain formed by merging
 // the common ancestors of the continuation and current continuation.
-
+obj continuation_restore(obj cc, obj k);
 
 // For debugging
 void print_continuation(obj cc);

tests/syntax.c

@@ -189,6 +189,13 @@ int test_callcc(void) {
     check_equal("(let ([x (call/cc (lambda (k) k))]) "
                   "(x (lambda (ignore) \"hi\")))",
                 "\"hi\"");
+
+    check_equal("(letrec ([k* #f] "
+                         "[y (fx1+ (call/cc (lambda (k) (set! k* k) 0)))]) "
+                  "(if (fx2< y 5) "
+                      "(k* y) "
+                      "y))",
+                "5");
 
     return passed;
 }