Multiprompt Delimited Continuations with Project Loom

DelimCC.java

import java.lang.Continuation;
import java.lang.ContinuationScope;
import java.util.function.Supplier;
import java.util.function.Function;
import java.util.Stack;

// multiprompt delimited continuations in terms of the current API
// this implementation has Felleisen classification -F+
class DelimCC {

  private static Object result;
  private static Object arg;
  private static CPS body;

  public static <A> A reset(Prompt<A> prompt, Supplier<A> f) {
    var k = new Continuation(prompt, () -> {
      result = f.get();
    });
    return runCont(k);
  }

  public static <A, R> A shift(Prompt<R> prompt, CPS<A, R> body) {
    DelimCC.body = body;
    Continuation.yield(prompt);
    return (A) arg;
  }

  private static <A> A runCont(Continuation k) {

    k.run();

    final Stack<Continuation> frames = new Stack<>();

    while (!k.isDone()) {
      
      // IDEA:
      //   1) Push a separate (one-time) prompt on `shift`. 
      //   2) On resume, capture the continuation on a heap allocated 
      //      stack of `Continuation`s
      //   3) Trampoline those continuations at the position of the original
      //      `reset`.
      //
      // This only works since continuations are one-shot. Otherwise the
      // captured frames would contain references to the continuation and
      // would be evaluated out of scope.
      final var bodyPrompt = new ContinuationScope() {};
      final var bodyCont = new Continuation(bodyPrompt, () -> {
        result = ((CPS<?, A>) body).apply(value -> {
          // yield and wait until the subcontinuation has been
          // evaluated.
          arg = value;
          // yielding here returns control to the outer continuation
          Continuation.yield(bodyPrompt);
          return (A) result;
        });
        body = null;
      });

      bodyCont.run(); // start it

      // continuation was called within body
      if (!bodyCont.isDone()) {
        frames.push(bodyCont);
        k.run();

      // continuation was discarded or escaped the dynamic scope of
      // bodyCont.
      } else {
        break;
      }
    }

    while (!frames.isEmpty()) {
      var frame = frames.pop();

      if (!frame.isDone()) {
        frame.run();
      }
    }

    return (A) result;
  }

  public static void main(String[] args) {
    test();
    test2();
  }

  static void test() {
    var p1 = new Prompt<Integer>();
    var p2 = new Prompt<Integer>();

    var res1 = reset(p1, () -> 5 - shift(p1, (Cont<Integer, Integer> k) -> k.resume(2) * 7));
    var res2 = reset(p1, () -> 1 + shift(p1, (Cont<Integer, Integer> k) -> k.resume(2))
                                 + shift(p1, (Cont<Integer, Integer> k) -> k.resume(3)));
    var res3 = reset(p1, () -> 1 + shift(p1, (Cont<Integer, Integer> k) -> k.resume(2))
                 + reset(p2, () -> 2 + shift(p2, (Cont<Integer, Integer> k) -> k.resume(3) * 3)
                                     + shift(p1, (Cont<Integer, Integer> k) -> k.resume(4) * 2)));

    var res4 = reset(p1, () -> 1 + shift(p1, (Cont<Integer, Integer> k) -> k.resume(2))
                 + reset(p2, () -> 2 + shift(p2, (Cont<Integer, Integer> k) -> k.resume(3) * 3)
                                     + shift(p1, (Cont<Integer, Integer> k) -> 42)));

    System.out.println(res1); // 21
    System.out.println(res2); // 6
    System.out.println(res3); // 60
    System.out.println(res4); // 42

    try {
      reset(p1, () -> 1 + shift(p1, (Cont<Integer, Integer> k) -> {
        return k.resume(2) + shift(p1, (Cont<Integer, Integer> k2) -> k2.resume(3));
      }));
    } catch (java.lang.IllegalStateException e) {
      System.out.println("Prompt not found");
    }
  }

  static void test2() {
    var p1 = new Prompt<Integer>();
    var res = reset(p1, () -> {
      var n = 10000;
      var r = 0;
      while (n > 0) {
        r += shift(p1, (Cont<Integer, Integer> k) -> k.resume(1));
        n--;
      }
      return r;
    });
    System.out.println(res);
  }
}

class Prompt<A> extends ContinuationScope {}

interface Cont<A, R> {
  R resume(A value);
}

interface CPS<A, R> {
  R apply(Cont<A, R> continuation);
}