Batch in action

- Rename TransformSet -> BatchingStage
- Inline BatchingTransformManager into BatchingEngine
- Reuse one Plan object for each DrawBuffer used by a stage.
- Separate DrawBuffer acquisition from marking as active.
- Remove some unused methods in BatchingDrawTracker
- Rename variables in AnimationTickHolder
- Make flw.loadRenderDoc=false behave as expected.

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

@@ -13,11 +13,9 @@ import com.mojang.blaze3d.vertex.BufferBuilder;
 import net.minecraft.client.renderer.RenderType;
 
 public class BatchingDrawTracker {
-	private static final RenderStage[] RENDER_STAGES = RenderStage.values();
-
 	private final Map<RenderStage, Set<DrawBuffer>> activeBuffers = new EnumMap<>(RenderStage.class);
 	{
-		for (RenderStage stage : RENDER_STAGES) {
+		for (RenderStage stage : RenderStage.values()) {
 			activeBuffers.put(stage, new HashSet<>());
 		}
 	}
@@ -30,44 +28,38 @@ public class BatchingDrawTracker {
 		((BufferBuilderExtension) scratch).flywheel$freeBuffer();
 	}
 
-	public DrawBuffer getBuffer(RenderType renderType, RenderStage stage) {
-		DrawBuffer buffer = RenderTypeExtension.getDrawBufferSet(renderType).getBuffer(stage);
-		activeBuffers.get(stage).add(buffer);
-		return buffer;
+	public static DrawBuffer getBuffer(RenderType renderType, RenderStage stage) {
+		return RenderTypeExtension.getDrawBufferSet(renderType)
+				.getBuffer(stage);
+	}
+
+	public void markActive(RenderStage stage, DrawBuffer buffer) {
+		synchronized (activeBuffers) {
+			activeBuffers.get(stage)
+					.add(buffer);
+		}
 	}
 
 	/**
 	 * Draw and reset all DrawBuffers for the given RenderStage.
+	 *
 	 * @param stage The RenderStage to draw.
 	 */
 	public void draw(RenderStage stage) {
 		Set<DrawBuffer> buffers = activeBuffers.get(stage);
 		for (DrawBuffer buffer : buffers) {
 			_draw(buffer);
-			buffer.reset();
 		}
 		buffers.clear();
 	}
 
-	/**
-	 * Draw and reset all active DrawBuffers.
-	 */
-	public void drawAll() {
-		for (Set<DrawBuffer> buffers : activeBuffers.values()) {
-			for (DrawBuffer buffer : buffers) {
-				_draw(buffer);
-				buffer.reset();
-			}
-			buffers.clear();
-		}
-	}
-
 	private void _draw(DrawBuffer buffer) {
 		if (buffer.hasVertices()) {
 			BufferBuilderExtension scratch = (BufferBuilderExtension) this.scratch;
 			buffer.inject(scratch);
 			buffer.getRenderType().end(this.scratch, 0, 0, 0);
 		}
+		buffer.reset();
 	}
 
 	/**

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

@@ -1,23 +1,36 @@
 package com.jozufozu.flywheel.backend.engine.batching;
 
+import java.util.ArrayList;
+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.instance.Instancer;
+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.InstancerKey;
+import com.jozufozu.flywheel.lib.task.NestedPlan;
 import com.jozufozu.flywheel.util.FlwUtil;
+import com.mojang.blaze3d.vertex.VertexFormat;
 
+import net.minecraft.client.renderer.RenderType;
 import net.minecraft.world.phys.Vec3;
 
 public class BatchingEngine extends AbstractEngine {
-	private final BatchingTransformManager transformManager = new BatchingTransformManager();
 	private final BatchingDrawTracker drawTracker = new BatchingDrawTracker();
+	private final Map<InstancerKey<?>, CPUInstancer<?>> instancers = new HashMap<>();
+	private final List<UninitializedInstancer> uninitializedInstancers = new ArrayList<>();
+	private final List<CPUInstancer<?>> initializedInstancers = new ArrayList<>();
+	private final Map<RenderStage, BatchingStage> stages = new EnumMap<>(RenderStage.class);
+	private final Map<VertexFormat, BatchedMeshPool> meshPools = new HashMap<>();
 
 	public BatchingEngine(int maxOriginDistance) {
 		super(maxOriginDistance);
@@ -25,7 +38,7 @@ public class BatchingEngine extends AbstractEngine {
 
 	@Override
 	public <I extends Instance> Instancer<I> instancer(InstanceType<I> type, Model model, RenderStage stage) {
-		return transformManager.getInstancer(type, model, stage);
+		return this.getInstancer(type, model, stage);
 	}
 
 	@Override
@@ -35,7 +48,15 @@ public class BatchingEngine extends AbstractEngine {
 		var stack = FlwUtil.copyPoseStack(context.stack());
 		stack.translate(renderOrigin.getX() - cameraPos.x, renderOrigin.getY() - cameraPos.y, renderOrigin.getZ() - cameraPos.z);
 
-		return transformManager.plan(stack.last(), context.level(), drawTracker);
+		flush();
+
+		var plans = new ArrayList<Plan>();
+
+		for (var transformSet : stages.values()) {
+			plans.add(transformSet.plan(stack.last(), context.level()));
+		}
+
+		return new NestedPlan(plans);
 	}
 
 	@Override
@@ -49,12 +70,18 @@ public class BatchingEngine extends AbstractEngine {
 
 	@Override
 	protected void onRenderOriginChanged() {
-		transformManager.clearInstancers();
+		initializedInstancers.forEach(CPUInstancer::clear);
 	}
 
 	@Override
 	public void delete() {
-		transformManager.delete();
+		instancers.clear();
+
+		meshPools.values()
+				.forEach(BatchedMeshPool::delete);
+		meshPools.clear();
+
+		initializedInstancers.clear();
 	}
 
 	@Override
@@ -62,4 +89,47 @@ public class BatchingEngine extends AbstractEngine {
 		info.add("Batching");
 		info.add("Origin: " + renderOrigin.getX() + ", " + renderOrigin.getY() + ", " + renderOrigin.getZ());
 	}
+
+	@SuppressWarnings("unchecked")
+	public <I extends Instance> Instancer<I> getInstancer(InstanceType<I> type, Model model, RenderStage stage) {
+		InstancerKey<I> key = new InstancerKey<>(type, model, stage);
+		CPUInstancer<I> instancer = (CPUInstancer<I>) instancers.get(key);
+		if (instancer == null) {
+			instancer = new CPUInstancer<>(type);
+			instancers.put(key, instancer);
+			uninitializedInstancers.add(new UninitializedInstancer(instancer, model, stage));
+		}
+		return instancer;
+	}
+
+	private void flush() {
+		for (var instancer : uninitializedInstancers) {
+			add(instancer.instancer(), instancer.model(), instancer.stage());
+		}
+		uninitializedInstancers.clear();
+
+		for (var pool : meshPools.values()) {
+			pool.flush();
+		}
+	}
+
+	private void add(CPUInstancer<?> instancer, Model model, RenderStage stage) {
+		var batchingStage = stages.computeIfAbsent(stage, renderStage -> new BatchingStage(renderStage, drawTracker));
+		var meshes = model.getMeshes();
+		for (var entry : meshes.entrySet()) {
+			var material = entry.getKey();
+			RenderType renderType = material.getBatchingRenderType();
+			var transformCall = new TransformCall<>(instancer, material, alloc(entry.getValue(), renderType.format()));
+			batchingStage.put(renderType, transformCall);
+		}
+		initializedInstancers.add(instancer);
+	}
+
+	private BatchedMeshPool.BufferedMesh alloc(Mesh mesh, VertexFormat format) {
+		return meshPools.computeIfAbsent(format, BatchedMeshPool::new)
+				.alloc(mesh);
+	}
+
+	private record UninitializedInstancer(CPUInstancer<?> instancer, Model model, RenderStage stage) {
+	}
 }

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

@@ -0,0 +1,105 @@
+package com.jozufozu.flywheel.backend.engine.batching;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import com.jozufozu.flywheel.api.event.RenderStage;
+import com.jozufozu.flywheel.api.task.Plan;
+import com.jozufozu.flywheel.api.task.TaskExecutor;
+import com.jozufozu.flywheel.lib.task.NestedPlan;
+import com.jozufozu.flywheel.lib.task.Synchronizer;
+import com.jozufozu.flywheel.lib.task.UnitPlan;
+import com.mojang.blaze3d.vertex.PoseStack;
+
+import net.minecraft.client.multiplayer.ClientLevel;
+import net.minecraft.client.renderer.RenderType;
+
+/**
+ * All the rendering that happens within a render stage.
+ */
+public class BatchingStage {
+	private final RenderStage stage;
+	private final BatchingDrawTracker tracker;
+	private final Map<RenderType, BufferPlan> buffers = new HashMap<>();
+
+	public BatchingStage(RenderStage renderStage, BatchingDrawTracker tracker) {
+		stage = renderStage;
+		this.tracker = tracker;
+	}
+
+	public Plan plan(PoseStack.Pose matrices, ClientLevel level) {
+		var plans = new ArrayList<Plan>();
+
+		for (var bufferPlan : buffers.values()) {
+			plans.add(bufferPlan.update(matrices, level));
+		}
+
+		return new NestedPlan(plans);
+	}
+
+	public void put(RenderType renderType, TransformCall<?> transformCall) {
+		buffers.computeIfAbsent(renderType, type -> new BufferPlan(BatchingDrawTracker.getBuffer(type, stage)))
+				.add(transformCall);
+	}
+
+	public boolean isEmpty() {
+		return buffers.isEmpty();
+	}
+
+	private class BufferPlan implements Plan {
+		private final DrawBuffer buffer;
+		private final List<TransformCall<?>> transformCalls = new ArrayList<>();
+		private PoseStack.Pose matrices;
+		private ClientLevel level;
+		private int vertexCount;
+
+		public BufferPlan(DrawBuffer drawBuffer) {
+			buffer = drawBuffer;
+		}
+
+		public Plan update(PoseStack.Pose matrices, ClientLevel level) {
+			this.matrices = matrices;
+			this.level = level;
+
+			vertexCount = setupAndCountVertices();
+			if (vertexCount <= 0) {
+				return UnitPlan.INSTANCE;
+			}
+
+			// Moving this into execute leads to a race condition that causes things to flash in and out of existence.
+			// Sometimes the main thread decides there's nothing to render in a stage before the worker threads have
+			// marked a stage as active. Then in the next frame #markActive complains because it's already prepared.
+			tracker.markActive(stage, buffer);
+			return this;
+		}
+
+		public void add(TransformCall<?> transformCall) {
+			transformCalls.add(transformCall);
+		}
+
+		@Override
+		public void execute(TaskExecutor taskExecutor, Runnable onCompletion) {
+			buffer.prepare(vertexCount);
+
+			var synchronizer = new Synchronizer(transformCalls.size(), onCompletion);
+
+			int startVertex = 0;
+			for (var transformCall : transformCalls) {
+				transformCall.plan(buffer, startVertex, matrices, level)
+						.execute(taskExecutor, synchronizer::decrementAndEventuallyRun);
+				startVertex += transformCall.getTotalVertexCount();
+			}
+		}
+
+		private int setupAndCountVertices() {
+			int vertices = 0;
+			for (var transformCall : transformCalls) {
+				transformCall.setup();
+				vertices += transformCall.getTotalVertexCount();
+			}
+			return vertices;
+		}
+	}
+}

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

@@ -1,98 +0,0 @@
-package com.jozufozu.flywheel.backend.engine.batching;
-
-import java.util.ArrayList;
-import java.util.EnumMap;
-import java.util.HashMap;
-import java.util.List;
-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.instance.Instancer;
-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.backend.engine.InstancerKey;
-import com.jozufozu.flywheel.lib.task.NestedPlan;
-import com.mojang.blaze3d.vertex.PoseStack;
-import com.mojang.blaze3d.vertex.VertexFormat;
-
-import net.minecraft.client.multiplayer.ClientLevel;
-import net.minecraft.client.renderer.RenderType;
-
-public class BatchingTransformManager {
-	private final Map<InstancerKey<?>, CPUInstancer<?>> instancers = new HashMap<>();
-	private final List<UninitializedInstancer> uninitializedInstancers = new ArrayList<>();
-	private final List<CPUInstancer<?>> initializedInstancers = new ArrayList<>();
-	private final Map<RenderStage, TransformSet> transformSets = new EnumMap<>(RenderStage.class);
-	private final Map<VertexFormat, BatchedMeshPool> meshPools = new HashMap<>();
-
-	public Plan plan(PoseStack.Pose matrices, ClientLevel level, BatchingDrawTracker tracker) {
-		flush();
-		var plans = new ArrayList<Plan>();
-
-		for (var transformSet : transformSets.values()) {
-			plans.add(transformSet.plan(matrices, level, tracker));
-		}
-
-		return new NestedPlan(plans);
-	}
-
-	@SuppressWarnings("unchecked")
-	public <I extends Instance> Instancer<I> getInstancer(InstanceType<I> type, Model model, RenderStage stage) {
-		InstancerKey<I> key = new InstancerKey<>(type, model, stage);
-		CPUInstancer<I> instancer = (CPUInstancer<I>) instancers.get(key);
-		if (instancer == null) {
-			instancer = new CPUInstancer<>(type);
-			instancers.put(key, instancer);
-			uninitializedInstancers.add(new UninitializedInstancer(instancer, model, stage));
-		}
-		return instancer;
-	}
-
-	public void flush() {
-		for (var instancer : uninitializedInstancers) {
-			add(instancer.instancer(), instancer.model(), instancer.stage());
-		}
-		uninitializedInstancers.clear();
-
-		for (var pool : meshPools.values()) {
-			pool.flush();
-		}
-	}
-
-	public void delete() {
-		instancers.clear();
-
-		meshPools.values()
-				.forEach(BatchedMeshPool::delete);
-		meshPools.clear();
-
-		initializedInstancers.clear();
-	}
-
-	public void clearInstancers() {
-		initializedInstancers.forEach(CPUInstancer::clear);
-	}
-
-	private void add(CPUInstancer<?> instancer, Model model, RenderStage stage) {
-		TransformSet transformSet = transformSets.computeIfAbsent(stage, TransformSet::new);
-		var meshes = model.getMeshes();
-		for (var entry : meshes.entrySet()) {
-			var material = entry.getKey();
-			RenderType renderType = material.getBatchingRenderType();
-			TransformCall<?> transformCall = new TransformCall<>(instancer, material, alloc(entry.getValue(), renderType.format()));
-			transformSet.put(renderType, transformCall);
-		}
-		initializedInstancers.add(instancer);
-	}
-
-	private BatchedMeshPool.BufferedMesh alloc(Mesh mesh, VertexFormat format) {
-		return meshPools.computeIfAbsent(format, BatchedMeshPool::new)
-				.alloc(mesh);
-	}
-
-	private record UninitializedInstancer(CPUInstancer<?> instancer, Model model, RenderStage stage) {
-	}
-}

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

@@ -15,7 +15,6 @@ public class DrawBufferSet {
 	private final VertexFormat format;
 	private final int stride;
 	private final VertexListProvider provider;
-
 	private final Map<RenderStage, DrawBuffer> buffers = new EnumMap<>(RenderStage.class);
 
 	public DrawBufferSet(RenderType renderType) {

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

@@ -1,62 +0,0 @@
-package com.jozufozu.flywheel.backend.engine.batching;
-
-import java.util.ArrayList;
-
-import com.google.common.collect.ArrayListMultimap;
-import com.google.common.collect.ListMultimap;
-import com.jozufozu.flywheel.api.event.RenderStage;
-import com.jozufozu.flywheel.api.task.Plan;
-import com.jozufozu.flywheel.lib.task.NestedPlan;
-import com.mojang.blaze3d.vertex.PoseStack;
-
-import net.minecraft.client.multiplayer.ClientLevel;
-import net.minecraft.client.renderer.RenderType;
-
-public class TransformSet {
-	private final RenderStage stage;
-	private final ListMultimap<RenderType, TransformCall<?>> transformCalls;
-
-	public TransformSet(RenderStage renderStage) {
-		stage = renderStage;
-		transformCalls = ArrayListMultimap.create();
-	}
-
-	public Plan plan(PoseStack.Pose matrices, ClientLevel level, BatchingDrawTracker tracker) {
-		var plans = new ArrayList<Plan>();
-
-		for (var entry : transformCalls.asMap()
-				.entrySet()) {
-			var renderType = entry.getKey();
-			var transformCalls = entry.getValue();
-
-			int vertices = 0;
-			for (var transformCall : transformCalls) {
-				transformCall.setup();
-				vertices += transformCall.getTotalVertexCount();
-			}
-
-			if (vertices == 0) {
-				continue;
-			}
-
-			DrawBuffer buffer = tracker.getBuffer(renderType, this.stage);
-			buffer.prepare(vertices);
-
-			int startVertex = 0;
-			for (var transformCall : transformCalls) {
-				plans.add(transformCall.plan(buffer, startVertex, matrices, level));
-				startVertex += transformCall.getTotalVertexCount();
-			}
-		}
-
-		return new NestedPlan(plans);
-	}
-
-	public void put(RenderType shaderState, TransformCall<?> transformCall) {
-		transformCalls.put(shaderState, transformCall);
-	}
-
-	public boolean isEmpty() {
-		return transformCalls.isEmpty();
-	}
-}

src/main/java/com/jozufozu/flywheel/lib/util/AnimationTickHolder.java

@@ -9,16 +9,16 @@ import net.minecraft.client.Minecraft;
  */
 public final class AnimationTickHolder {
 	// Wrap around every 24 hours to maintain floating point accuracy.
-	private static final int wrappingInterval = 1_728_000;
+	private static final int WRAPPING_INTERVAL = 1_728_000;
 	private static int ticks;
-	private static int paused_ticks;
+	private static int pausedTicks;
 
 	public static void tick() {
 		if (!Minecraft.getInstance()
 				.isPaused()) {
-			ticks = (ticks + 1) % wrappingInterval;
+			ticks = (ticks + 1) % WRAPPING_INTERVAL;
 		} else {
-			paused_ticks = (paused_ticks + 1) % wrappingInterval;
+			pausedTicks = (pausedTicks + 1) % WRAPPING_INTERVAL;
 		}
 	}
 
@@ -27,7 +27,7 @@ public final class AnimationTickHolder {
 	}
 
 	public static int getTicks(boolean includePaused) {
-		return includePaused ? ticks + paused_ticks : ticks;
+		return includePaused ? ticks + pausedTicks : ticks;
 	}
 
 	public static float getRenderTime() {
@@ -42,6 +42,6 @@ public final class AnimationTickHolder {
 	// Unused but might be useful for debugging.
 	public static void _reset() {
 		ticks = 0;
-		paused_ticks = 0;
+		pausedTicks = 0;
 	}
 }

src/main/java/com/jozufozu/flywheel/mixin/ClientMainMixin.java

@@ -11,8 +11,8 @@ import net.minecraft.client.main.Main;
 public class ClientMainMixin {
 	@Inject(method = "main([Ljava/lang/String;)V", at = @At("HEAD"))
 	private static void flywheel$injectRenderDoc(CallbackInfo ci) {
-		// Only try to load RenderDoc if a system property is set.
-		if (System.getProperty("flw.loadRenderDoc") == null) {
+		// Only try to load RenderDoc if a system property is set to true.
+		if (!Boolean.parseBoolean(System.getProperty("flw.loadRenderDoc"))) {
 			return;
 		}