import java.util.concurrent.Future;
import jdk.incubator.concurrent.ExtentLocal;
import jdk.incubator.concurrent.StructuredTaskScope;

public class ExtentLocalBreakingDemo {

    private static final int ITERATIONS = 10_000;
    private static final int THRESHOLD = 64;
    private static final int N = 50_000;
    private static final ExtentLocal<int[]> DATA = ExtentLocal.newInstance();

    public static void main(String[] args) throws Exception {

        for (int i = 0; i < ITERATIONS; ++i) {
            System.out.println("Iteration " + i);
            int[] data = new int[N];

            fill(data, 42);

            for (int j = 0; j < N; ++j) {
                if (data[j] != 42) {
                    System.out.println("data[" + j + "] != 42");
                    System.exit(1);
                }
            }

            System.out.println("OK");
        }
    }

    private static void fill(int[] arr, int value) {
        ExtentLocal
                .where(DATA, arr)
                .run(() -> doFill(value, 0, arr.length));
    }

    private static Void doFill(int value, int begin, int end) {
        if (end <= begin) {
            return null;
        }

        if (end - begin <= THRESHOLD) {
            for (int i = begin; i != end; ++i) {
                DATA.get()[i] = value;
            }
            return null;
        }

        try (var scope = new StructuredTaskScope.ShutdownOnFailure()) {
            int middle = begin + (end - begin) / 2;

            Future<Void> left = scope.fork(() -> doFill(value, begin, middle));
            Future<Void> right = scope.fork(() -> doFill(value, middle, end));

            // (1) BEGIN
            scope.join();
            scope.throwIfFailed();
            // (1) END

            // (2) BEGIN
            try {
                left.get();
                right.get();
                DATA.get();
            } catch (NullPointerException e) {
                synchronized (ExtentLocalBreakingDemo.class) {
                    System.err.println(Thread.currentThread());
                    e.printStackTrace();
                }
                throw e;
            }
            // (2) END
        } catch (Exception e) {
            throw new RuntimeException(e);
        }

        return null;
    }
}