Consistent batches

- Reintroduce BatchedDrawManager.
- BatchingEngine no longer implements plan itself, instead uses
  composition of other plans.
- Add DynamicNestedPlan, runs many plans provided at execution time.
- Add ContextFunction and ContextSupplier to match *Consumer and
  *Runnable.
- Add unit tests for DynamicNestedPlan and IfElsePlan.

src/main/java/com/jozufozu/flywheel/backend/engine/batching/BatchedDrawManager.java

@@ -0,0 +1,81 @@
+package com.jozufozu.flywheel.backend.engine.batching;
+
+import java.util.Collection;
+import java.util.EnumMap;
+import java.util.HashMap;
+import java.util.Map;
+
+import com.jozufozu.flywheel.api.event.RenderStage;
+import com.jozufozu.flywheel.api.instance.Instance;
+import com.jozufozu.flywheel.api.instance.InstanceType;
+import com.jozufozu.flywheel.api.model.Mesh;
+import com.jozufozu.flywheel.api.model.Model;
+import com.jozufozu.flywheel.api.task.TaskExecutor;
+import com.jozufozu.flywheel.backend.engine.InstancerKey;
+import com.jozufozu.flywheel.backend.engine.InstancerStorage;
+import com.mojang.blaze3d.vertex.VertexFormat;
+
+import net.minecraft.client.renderer.RenderType;
+
+class BatchedDrawManager extends InstancerStorage<BatchedInstancer<?>> {
+	private final BatchedDrawTracker drawTracker = new BatchedDrawTracker();
+	private final Map<RenderStage, BatchedStagePlan> stagePlans = new EnumMap<>(RenderStage.class);
+	private final Map<VertexFormat, BatchedMeshPool> meshPools = new HashMap<>();
+
+	public Collection<BatchedStagePlan> getStagePlans() {
+		return stagePlans.values();
+	}
+
+	public void renderStage(TaskExecutor executor, RenderStage stage) {
+		var stagePlan = stagePlans.get(stage);
+
+		if (stagePlan == null) {
+			return;
+		}
+
+		executor.syncUntil(stagePlan.flag::isRaised);
+		stagePlan.flag.lower();
+
+		drawTracker.draw(stage);
+	}
+
+	@Override
+	protected <I extends Instance> BatchedInstancer<?> create(InstanceType<I> type) {
+		return new BatchedInstancer<>(type);
+	}
+
+	@Override
+	protected <I extends Instance> void add(InstancerKey<I> key, BatchedInstancer<?> instancer, Model model, RenderStage stage) {
+		var stagePlan = stagePlans.computeIfAbsent(stage, renderStage -> new BatchedStagePlan(renderStage, drawTracker));
+		var meshes = model.getMeshes();
+		for (var entry : meshes.entrySet()) {
+			var material = entry.getKey();
+			RenderType renderType = material.getFallbackRenderType();
+			var transformCall = new TransformCall<>(instancer, material, alloc(entry.getValue(), renderType.format()));
+			stagePlan.put(renderType, transformCall);
+		}
+	}
+
+	@Override
+	public void flush() {
+		super.flush();
+
+		for (var pool : meshPools.values()) {
+			pool.flush();
+		}
+	}
+
+	@Override
+	public void invalidate() {
+		super.invalidate();
+
+		meshPools.values()
+				.forEach(BatchedMeshPool::delete);
+		meshPools.clear();
+	}
+
+	private BatchedMeshPool.BufferedMesh alloc(Mesh mesh, VertexFormat format) {
+		return meshPools.computeIfAbsent(format, BatchedMeshPool::new)
+				.alloc(mesh);
+	}
+}

src/main/java/com/jozufozu/flywheel/backend/engine/batching/BatchingEngine.java

@@ -1,55 +1,23 @@
 package com.jozufozu.flywheel.backend.engine.batching;
 
-import java.util.EnumMap;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 
 import com.jozufozu.flywheel.api.event.RenderContext;
 import com.jozufozu.flywheel.api.event.RenderStage;
 import com.jozufozu.flywheel.api.instance.Instance;
-import com.jozufozu.flywheel.api.instance.InstanceType;
-import com.jozufozu.flywheel.api.model.Mesh;
-import com.jozufozu.flywheel.api.model.Model;
 import com.jozufozu.flywheel.api.task.Plan;
 import com.jozufozu.flywheel.api.task.TaskExecutor;
 import com.jozufozu.flywheel.backend.engine.AbstractEngine;
 import com.jozufozu.flywheel.backend.engine.AbstractInstancer;
-import com.jozufozu.flywheel.backend.engine.InstancerKey;
 import com.jozufozu.flywheel.backend.engine.InstancerStorage;
+import com.jozufozu.flywheel.lib.task.DynamicNestedPlan;
 import com.jozufozu.flywheel.lib.task.Flag;
+import com.jozufozu.flywheel.lib.task.MapContextPlan;
 import com.jozufozu.flywheel.lib.task.NamedFlag;
-import com.jozufozu.flywheel.lib.task.SimplyComposedPlan;
-import com.jozufozu.flywheel.lib.task.Synchronizer;
-import com.mojang.blaze3d.vertex.VertexFormat;
+import com.jozufozu.flywheel.lib.task.SimplePlan;
 
-import net.minecraft.client.renderer.RenderType;
-
-public class BatchingEngine extends AbstractEngine implements SimplyComposedPlan<RenderContext> {
-	private final BatchedDrawTracker drawTracker = new BatchedDrawTracker();
-
-	// TODO: reintroduce BatchedDrawManager
-	private final InstancerStorage<BatchedInstancer<?>> storage = new InstancerStorage<>() {
-        @Override
-        protected <I extends Instance> BatchedInstancer<?> create(InstanceType<I> type) {
-            return new BatchedInstancer<>(type);
-        }
-
-        @Override
-		protected <I extends Instance> void add(InstancerKey<I> key, BatchedInstancer<?> instancer, Model model, RenderStage stage) {
-            var stagePlan = stagePlans.computeIfAbsent(stage, renderStage -> new BatchedStagePlan(renderStage, drawTracker));
-            var meshes = model.getMeshes();
-            for (var entry : meshes.entrySet()) {
-                var material = entry.getKey();
-                RenderType renderType = material.getFallbackRenderType();
-                var transformCall = new TransformCall<>(instancer, material, alloc(entry.getValue(), renderType.format()));
-                stagePlan.put(renderType, transformCall);
-            }
-        }
-    };
-
-	private final Map<RenderStage, BatchedStagePlan> stagePlans = new EnumMap<>(RenderStage.class);
-	private final Map<VertexFormat, BatchedMeshPool> meshPools = new HashMap<>();
+public class BatchingEngine extends AbstractEngine {
+	private final BatchedDrawManager drawManager = new BatchedDrawManager();
 
 	private final Flag flushFlag = new NamedFlag("flushed");
 
@@ -57,26 +25,13 @@ public class BatchingEngine extends AbstractEngine implements SimplyComposedPlan
 		super(maxOriginDistance);
 	}
 
-	@Override
-	public void execute(TaskExecutor taskExecutor, RenderContext context, Runnable onCompletion) {
-		flush();
-
-		// Now it's safe to read stage plans in #renderStage.
-		flushFlag.raise();
-
-		BatchContext ctx = BatchContext.create(context, renderOrigin);
-
-		var sync = new Synchronizer(stagePlans.values()
-				.size(), onCompletion);
-
-		for (var stagePlan : stagePlans.values()) {
-			stagePlan.execute(taskExecutor, ctx, sync);
-		}
-	}
-
 	@Override
 	public Plan<RenderContext> createFramePlan() {
-		return this;
+		return SimplePlan.<RenderContext>of(() -> {
+			drawManager.flush();
+			flushFlag.raise();
+		}).then(MapContextPlan.map((RenderContext ctx) -> BatchContext.create(ctx, renderOrigin))
+				.to(DynamicNestedPlan.of(drawManager::getStagePlans)));
 	}
 
 	@Override
@@ -86,16 +41,7 @@ public class BatchingEngine extends AbstractEngine implements SimplyComposedPlan
 			flushFlag.lower();
 		}
 
-		var stagePlan = stagePlans.get(stage);
-
-		if (stagePlan == null) {
-			drawTracker.draw(stage);
-			return;
-		}
-
-		executor.syncUntil(stagePlan.flag::isRaised);
-		stagePlan.flag.lower();
-		drawTracker.draw(stage);
+		drawManager.renderStage(executor, stage);
 	}
 
 	@Override
@@ -105,28 +51,11 @@ public class BatchingEngine extends AbstractEngine implements SimplyComposedPlan
 
 	@Override
 	protected InstancerStorage<? extends AbstractInstancer<?>> getStorage() {
-		return storage;
+		return drawManager;
 	}
 
 	@Override
 	public void delete() {
-		storage.invalidate();
-
-		meshPools.values()
-				.forEach(BatchedMeshPool::delete);
-		meshPools.clear();
-	}
-
-	private void flush() {
-		storage.flush();
-
-		for (var pool : meshPools.values()) {
-			pool.flush();
-		}
-	}
-
-	private BatchedMeshPool.BufferedMesh alloc(Mesh mesh, VertexFormat format) {
-		return meshPools.computeIfAbsent(format, BatchedMeshPool::new)
-				.alloc(mesh);
+		drawManager.invalidate();
 	}
 }

src/main/java/com/jozufozu/flywheel/lib/task/ContextFunction.java

@@ -0,0 +1,13 @@
+package com.jozufozu.flywheel.lib.task;
+
+import com.jozufozu.flywheel.api.task.Plan;
+
+/**
+ * A function like interface for use with {@link Plan}s.
+ * @param <C> The context type.
+ * @param <R> The return type.
+ */
+@FunctionalInterface
+public interface ContextFunction<C, R> {
+	R apply(C context);
+}

src/main/java/com/jozufozu/flywheel/lib/task/ContextSupplier.java

@@ -0,0 +1,17 @@
+package com.jozufozu.flywheel.lib.task;
+
+/**
+ * A {@link ContextFunction} that ignores the context object.
+ *
+ * @param <C> The context type.
+ * @param <R> The return type.
+ */
+@FunctionalInterface
+public interface ContextSupplier<C, R> extends ContextFunction<C, R> {
+	R get();
+
+	@Override
+	default R apply(C ignored) {
+		return get();
+	}
+}

src/main/java/com/jozufozu/flywheel/lib/task/DynamicNestedPlan.java

@@ -0,0 +1,38 @@
+package com.jozufozu.flywheel.lib.task;
+
+import java.util.Collection;
+
+import com.jozufozu.flywheel.api.task.Plan;
+import com.jozufozu.flywheel.api.task.TaskExecutor;
+
+/**
+ * A plan that executes many other plans provided dynamically.
+ *
+ * @param plans A function to get a collection of plans based on the context.
+ * @param <C>   The type of the context object.
+ */
+public record DynamicNestedPlan<C>(ContextFunction<C, Collection<? extends Plan<C>>> plans) implements SimplyComposedPlan<C> {
+	public static <C> Plan<C> of(ContextSupplier<C, Collection<? extends Plan<C>>> supplier) {
+		return new DynamicNestedPlan<>(supplier);
+	}
+
+	public static <C> Plan<C> of(ContextFunction<C, Collection<? extends Plan<C>>> function) {
+		return new DynamicNestedPlan<>(function);
+	}
+
+	@Override
+	public void execute(TaskExecutor taskExecutor, C context, Runnable onCompletion) {
+		var plans = this.plans.apply(context);
+
+		if (plans.isEmpty()) {
+			onCompletion.run();
+			return;
+		}
+
+		var sync = new Synchronizer(plans.size(), onCompletion);
+
+		for (var plan : plans) {
+			plan.execute(taskExecutor, context, sync);
+		}
+	}
+}

src/test/java/com/jozufozu/flywheel/lib/task/PlanExecutionTest.java

@@ -1,5 +1,7 @@
 package com.jozufozu.flywheel.lib.task;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -114,29 +116,74 @@ class PlanExecutionTest {
 	}
 
 	@Test
-	void unitPlan() {
+	void emptyPlansDontCallTheExecutor() {
 		var done = new AtomicBoolean(false);
 
 		UnitPlan.of()
 				.execute(null, Unit.INSTANCE, () -> done.set(true));
 
 		Assertions.assertTrue(done.get());
-	}
-
-	@Test
-	void emptyPlan() {
-		var done = new AtomicBoolean(false);
+		done.set(false);
 
 		SimplePlan.of()
 				.execute(null, Unit.INSTANCE, () -> done.set(true));
-		Assertions.assertTrue(done.get());
 
+		Assertions.assertTrue(done.get());
 		done.set(false);
+
 		NestedPlan.of()
 				.execute(null, Unit.INSTANCE, () -> done.set(true));
+
 		Assertions.assertTrue(done.get());
 	}
 
+	@Test
+	void ifElsePlan() {
+		var branch = new AtomicInteger(0);
+
+		var plan = IfElsePlan.<Boolean>on(b -> b)
+				.ifTrue(SimplePlan.of(() -> branch.set(1)))
+				.ifFalse(SimplePlan.of(() -> branch.set(2)))
+				.plan();
+
+		runAndWait(plan, true);
+
+		Assertions.assertEquals(1, branch.get());
+
+		runAndWait(plan, false);
+
+		Assertions.assertEquals(2, branch.get());
+	}
+
+	@Test
+	void dynamicNestedPlan() {
+		var counter = new AtomicInteger(0);
+
+		List<Plan<Unit>> plans = new ArrayList<>();
+
+		// We'll re-use this same plan but append to the list of plans it executes.
+		var plan = DynamicNestedPlan.of(() -> plans);
+
+		runAndWait(plan);
+
+		Assertions.assertEquals(0, counter.get());
+
+		plans.add(SimplePlan.of(counter::incrementAndGet));
+
+		runAndWait(plan);
+
+		Assertions.assertEquals(1, counter.get());
+
+		counter.set(0);
+
+		plans.add(SimplePlan.of(counter::incrementAndGet));
+		plans.add(SimplePlan.of(counter::incrementAndGet));
+
+		runAndWait(plan);
+
+		Assertions.assertEquals(3, counter.get());
+	}
+
 	@Test
 	void mainThreadPlan() {
 		var done = new AtomicBoolean(false);
@@ -213,19 +260,25 @@ class PlanExecutionTest {
 	}
 
 	public static void runAndWait(Plan<Unit> plan) {
-		new TestBarrier(plan).runAndWait();
+		new TestBarrier<Unit>(plan, Unit.INSTANCE).runAndWait();
 	}
 
-	private static final class TestBarrier {
-		private final Plan<Unit> plan;
+	public static <C> void runAndWait(Plan<C> plan, C ctx) {
+		new TestBarrier<>(plan, ctx).runAndWait();
+	}
+
+	private static final class TestBarrier<C> {
+		private final Plan<C> plan;
+		private final C ctx;
 		private boolean done = false;
 
-		private TestBarrier(Plan<Unit> plan) {
+		private TestBarrier(Plan<C> plan, C ctx) {
 			this.plan = plan;
+			this.ctx = ctx;
 		}
 
 		public void runAndWait() {
-			plan.execute(EXECUTOR, Unit.INSTANCE, this::doneWithPlan);
+			plan.execute(EXECUTOR, ctx, this::doneWithPlan);
 
 			synchronized (this) {
 				// early exit in case the plan is already done for e.g. UnitPlan