Java21 Continuations Demo

Continuation.java

package demo.cont;

import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodType;
import java.util.function.Consumer;

public final class Continuation {

    private final Object delegate;

    private static final Class<?> IMPL_CLASS;
    private static final MethodHandle NEW;
    private static final MethodHandle YIELD;
    private static final MethodHandle RUN;
    private static final MethodHandle IS_DONE;

    static {
        try {
            IMPL_CLASS = Class.forName("jdk.internal.vm.Continuation");

            var lookup = MethodHandles.privateLookupIn(IMPL_CLASS, MethodHandles.lookup());
            NEW = lookup.findConstructor(IMPL_CLASS, MethodType.methodType(void.class, Scope.IMPL_CLASS, Runnable.class));
            RUN = lookup.findVirtual(IMPL_CLASS, "run", MethodType.methodType(void.class));
            YIELD = lookup.findStatic(IMPL_CLASS, "yield", MethodType.methodType(boolean.class, Scope.IMPL_CLASS));
            IS_DONE = lookup.findVirtual(IMPL_CLASS, "isDone", MethodType.methodType(boolean.class));

        } catch (Throwable t) {
            throw new ExceptionInInitializerError(t);
        }
    }

    public <T> Continuation(String scopeName, Consumer<Scope<T>> code) {
        this(scope(scopeName), code);
    }

    public <T> Continuation(Scope<T> scope, Consumer<Scope<T>> code) {
        try {
            this.delegate = NEW.invoke(scope.delegate, (Runnable) () -> code.accept(scope));
        } catch (Throwable t) {
            throw new RuntimeException(t);
        }
    }

    public void run() {
        try {
            RUN.invoke(delegate);
        } catch (Throwable t) {
            throw new RuntimeException(t);
        }
    }

    public <T> T next(Scope<T> scope) {
        run();
        return getState(scope);
    }

    public static <T> void yield(Scope<T> scope) {
        try {
            YIELD.invoke(scope.delegate);
        } catch (Throwable t) {
            throw new RuntimeException(t);
        }
    }

    public static <T> void yield(Scope<T> scope, T state) {
        ScopedValue.where(scope.state, state).run(() -> {
            try {
                YIELD.invoke(scope.delegate);
            } catch (Throwable t) {
                throw new RuntimeException(t);
            }
        });
    }

    public static <T> T getState(Scope<T> scope) {
        return scope.state.get();
    }

    public boolean isDone() {
        try {
            return (boolean) IS_DONE.invoke(delegate);
        } catch (Throwable t) {
            throw new RuntimeException(t);
        }
    }

    public static <T> Scope<T> scope(String name) {
        return new Scope<>(name);
    }

    public static final class Scope<T> {

        private final Object delegate;

        private final ScopedValue<T> state;

        private static final Class<?> IMPL_CLASS;
        private static final MethodHandle NEW;

        static {
            try {
                IMPL_CLASS = Class.forName("jdk.internal.vm.ContinuationScope");

                var lookup = MethodHandles.privateLookupIn(IMPL_CLASS, MethodHandles.lookup());
                NEW = lookup.findConstructor(IMPL_CLASS, MethodType.methodType(void.class, String.class));

            } catch (Throwable t) {
                throw new ExceptionInInitializerError(t);
            }
        }

        private Scope(String name) {
            try {
                this.delegate = NEW.invoke(name);
                this.state = ScopedValue.newInstance();
            } catch (Throwable e) {
                throw new RuntimeException(e);
            }
        }

        @Override
        public String toString() {
            return delegate.toString();
        }
    }
}

ContinuationDemo.java

package demo.cont;

import java.util.Scanner;

/**
 * Run on JDK 21+ with <code>--add-opens java.base/jdk.internal.vm=ALL-UNNAMED --enable-preview</code>
 */
public class ContinuationDemo {

    static void countUp(Continuation.Scope<?> scope) {
        for (int i = 0; i < 10; i++) {
            System.out.println(i);
            Continuation.yield(scope); // gives up control until continuation is resumed
        }
        System.out.println("Done");
    }

    public static void main(String[] args) {

        var continuation = new Continuation("myscope", ContinuationDemo::countUp);

        try (var scanner = new Scanner(System.in)) {
            while (!continuation.isDone()) { // continuation code block is not finished
                System.out.print("Press enter to run one more step: ");
                scanner.nextLine();
                continuation.run(); // resume continuation
            }
            System.out.println("No more input.");
        }
    }
}

ContinuationDemo2.java

package demo.cont;

import java.util.Scanner;

/**
 * Run on JDK 21+ with <code>--add-opens java.base/jdk.internal.vm=ALL-UNNAMED --enable-preview</code>
 */
public class ContinuationDemo2 {

    static void countUp(Continuation.Scope<?> scope) {
        for (int i = 0; i < 10; i++) {
            System.out.printf("%s: %s%n", scope, i);
            Continuation.yield(scope); // gives up control until continuation is resumed
        }
        System.out.printf("%s: Done%n", scope);
    }

    public static void main(String[] args) {

        var counter1 = new Continuation("counter1", ContinuationDemo2::countUp);
        var counter2 = new Continuation("counter2", ContinuationDemo2::countUp);

        try (var scanner = new Scanner(System.in)) {
            loop:
            while (!counter1.isDone() || !counter2.isDone()) { // continuation code blocks are not finished
                System.out.print("Enter 1 or 2 to select counter to increment, or 0 to exit: ");
                int counter = scanner.nextInt();
                switch (counter) {
                    case 1:
                        if (!counter1.isDone()) {
                            counter1.run(); // throws "IllegalStateException: Continuation terminated" if already completed
                        }
                        break;
                    case 2:
                        if (!counter2.isDone()) {
                            counter2.run();
                        }
                        break;
                    case 0:
                        break loop;
                }
            }
            System.out.println("No more input.");
        }
    }
}

ContinuationDemo3.java

package demo.cont;

import java.util.Scanner;

/**
 * Run on JDK 21+ with <code>--add-opens java.base/jdk.internal.vm=ALL-UNNAMED --enable-preview</code>
 */
public class ContinuationDemo3 {

    static void generator(Continuation.Scope<Integer> scope) {
        for (int i = 0; ; i++) {
            Continuation.yield(scope, i);
        }
    }

    public static void main(String[] args) {

        var scope = Continuation.<Integer>scope("counter");
        var counter = new Continuation(scope, ContinuationDemo3::generator);

        try (var scanner = new Scanner(System.in)) {
            while (!counter.isDone()) {
                scanner.nextLine();
                int value = counter.next(scope);
                System.out.printf("State: %s%n", value);
            }
        }
    }
}

I added implements Iterable<T> to Continuation if anyone else wants something similar:

Usage is like:

var scope = Continuation.<Integer>scope("tree");
var continuation = new Continuation<Integer>(scope, s -> traverseTree(s, tree));

try (var scanner = new Scanner(System.in)) {
    for (var value : continuation) {
        System.out.println(value);
        System.out.println("Press enter to continue...");
        scanner.nextLine();
    }
}

Implementation:

final class Continuation<T> implements Iterable<T> {

    private final Object delegate;
    private final Scope<T> scope;

    private static final Class<?> IMPL_CLASS;
    private static final MethodHandle NEW;
    private static final MethodHandle YIELD;
    private static final MethodHandle RUN;
    private static final MethodHandle IS_DONE;

    static {
        try {
            IMPL_CLASS = Class.forName("jdk.internal.vm.Continuation");

            var lookup = MethodHandles.privateLookupIn(IMPL_CLASS, MethodHandles.lookup());
            NEW = lookup.findConstructor(IMPL_CLASS,
                    MethodType.methodType(void.class, Scope.IMPL_CLASS, Runnable.class));
            RUN = lookup.findVirtual(IMPL_CLASS, "run", MethodType.methodType(void.class));
            YIELD = lookup.findStatic(IMPL_CLASS, "yield", MethodType.methodType(boolean.class, Scope.IMPL_CLASS));
            IS_DONE = lookup.findVirtual(IMPL_CLASS, "isDone", MethodType.methodType(boolean.class));

        } catch (Throwable t) {
            throw new ExceptionInInitializerError(t);
        }
    }

    public Continuation(String scopeName, Consumer<Scope<T>> code) {
        this(scope(scopeName), code);
    }

    public Continuation(Scope<T> scope, Consumer<Scope<T>> code) {
        try {
            this.scope = scope;
            this.delegate = NEW.invoke(scope.delegate, (Runnable) () -> code.accept(scope));
        } catch (Throwable t) {
            throw new RuntimeException(t);
        }
    }

    public void run() {
        try {
            RUN.invoke(delegate);
        } catch (Throwable t) {
            throw new RuntimeException(t);
        }
    }

    public T next(Scope<T> scope) {
        run();
        return getState(scope);
    }

    public static <T> void yield(Scope<T> scope) {
        try {
            YIELD.invoke(scope.delegate);
        } catch (Throwable t) {
            throw new RuntimeException(t);
        }
    }

    public static <T> void yield(Scope<T> scope, T state) {
        ScopedValue.where(scope.state, state).run(() -> {
            try {
                YIELD.invoke(scope.delegate);
            } catch (Throwable t) {
                throw new RuntimeException(t);
            }
        });
    }

    public static <T> T getState(Scope<T> scope) {
        return scope.state.get();
    }

    public boolean isDone() {
        try {
            return (boolean) IS_DONE.invoke(delegate);
        } catch (Throwable t) {
            throw new RuntimeException(t);
        }
    }

    public static <T> Scope<T> scope(String name) {
        return new Scope<>(name);
    }

    public static final class Scope<T> {

        private final Object delegate;

        private final ScopedValue<T> state;

        private static final Class<?> IMPL_CLASS;
        private static final MethodHandle NEW;

        static {
            try {
                IMPL_CLASS = Class.forName("jdk.internal.vm.ContinuationScope");

                var lookup = MethodHandles.privateLookupIn(IMPL_CLASS, MethodHandles.lookup());
                NEW = lookup.findConstructor(IMPL_CLASS, MethodType.methodType(void.class, String.class));

            } catch (Throwable t) {
                throw new ExceptionInInitializerError(t);
            }
        }

        private Scope(String name) {
            try {
                this.delegate = NEW.invoke(name);
                this.state = ScopedValue.newInstance();
            } catch (Throwable e) {
                throw new RuntimeException(e);
            }
        }

        @Override
        public String toString() {
            return delegate.toString();
        }
    }

    @Override
    public Iterator<T> iterator() {
        return new Iterator<T>() {
            Continuation<T> current = Continuation.this;

            @Override
            public boolean hasNext() {
                return !current.isDone();
            }

            @Override
            public T next() {
                if (!hasNext()) {
                    throw new NoSuchElementException();
                }
                T result = current.next(scope);
                return result;
            }
        };
    }
}