BatchingTransformManager

- Rename BatchingDrawManager to BatchingTransformManager
- Rename BatchDrawingTracker to BatchingDrawTracker
- Move BatchDrawingTracker from BatchingTransformManager to
BatchingEngine
- Do not prepare DrawBuffer if number of vertices is 0
- Add DrawBuffer#isPrepared
- Add checks to DrawBuffer methods
- Remove commented code from CPUInstancerFactory

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

@@ -9,12 +9,12 @@ import com.mojang.blaze3d.vertex.BufferBuilder;
 
 import net.minecraft.client.renderer.RenderType;
 
-public class BatchDrawingTracker {
+public class BatchingDrawTracker {
 
-	protected final Set<RenderType> activeTypes = new HashSet<>();
+	private final Set<RenderType> activeTypes = new HashSet<>();
 	private final BufferBuilder scratch;
 
-	public BatchDrawingTracker() {
+	public BatchingDrawTracker() {
 		scratch = new BufferBuilder(8);
 
 		((BufferBuilderExtension) scratch).flywheel$freeBuffer();
@@ -25,13 +25,20 @@ public class BatchDrawingTracker {
 		return RenderTypeExtension.getDrawBuffer(renderType);
 	}
 
+	/**
+	 * Draw and reset the DrawBuffer for the given RenderType.
+	 * @param renderType The RenderType to draw.
+	 */
+	public void draw(RenderType renderType) {
+		_draw(renderType);
+
+		activeTypes.remove(renderType);
+	}
+
 	/**
 	 * Draws all active DrawBuffers and reset them.
 	 */
-	public void endBatch() {
+	public void drawAll() {
-		// TODO: when/if this causes trouble with shaders, try to inject our BufferBuilders
-		//  into the RenderBuffers from context.
-
 		for (RenderType renderType : activeTypes) {
 			_draw(renderType);
 		}
@@ -39,38 +46,27 @@ public class BatchDrawingTracker {
 		activeTypes.clear();
 	}
 
-	/**
+	private void _draw(RenderType renderType) {
-	 * Draw and reset the DrawBuffer for the given RenderType.
+		DrawBuffer buffer = RenderTypeExtension.getDrawBuffer(renderType);
-	 * @param renderType The RenderType to draw.
-	 */
-	public void endBatch(RenderType renderType) {
-		_draw(renderType);
 
-		activeTypes.remove(renderType);
+		if (buffer.hasVertices()) {
+			BufferBuilderExtension scratch = (BufferBuilderExtension) this.scratch;
+			buffer.inject(scratch);
+			renderType.end(this.scratch, 0, 0, 0);
+			buffer.reset();
+		}
 	}
 
 	/**
 	 * Resets all DrawBuffers to 0 vertices.
 	 */
-	public void clear() {
+	public void reset() {
 		for (RenderType type : activeTypes) {
 			RenderTypeExtension.getDrawBuffer(type)
 					.reset();
 		}
+
 		activeTypes.clear();
 	}
 
-	private void _draw(RenderType renderType) {
-		DrawBuffer drawBuffer = RenderTypeExtension.getDrawBuffer(renderType);
-
-		BufferBuilderExtension scratch = (BufferBuilderExtension) this.scratch;
-		if (drawBuffer.hasVertices()) {
-			drawBuffer.inject(scratch);
-
-			renderType.end(this.scratch, 0, 0, 0);
-
-			drawBuffer.reset();
-		}
-	}
-
 }

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

@@ -24,7 +24,8 @@ import net.minecraft.world.phys.Vec3;
 
 public class BatchingEngine implements Engine {
 
-	protected final BatchingDrawManager drawManager = new BatchingDrawManager();
+	protected final BatchingTransformManager transformManager = new BatchingTransformManager();
+	protected final BatchingDrawTracker drawTracker = new BatchingDrawTracker();
 	protected final Map<StructType<?>, CPUInstancerFactory<?>> factories = new HashMap<>();
 
 	@SuppressWarnings("unchecked")
@@ -36,12 +37,12 @@ public class BatchingEngine implements Engine {
 
 	@NotNull
 	private <D extends InstancedPart> CPUInstancerFactory<D> createFactory(StructType<D> type) {
-		return new CPUInstancerFactory<>(type, drawManager::create);
+		return new CPUInstancerFactory<>(type, transformManager::create);
 	}
 
 	@Override
 	public void beginFrame(TaskEngine taskEngine, RenderContext context) {
-		drawManager.flush();
+		transformManager.flush();
 
 		Vec3 cameraPos = context.camera().getPosition();
 		var stack = FlwUtil.copyPoseStack(context.stack());
@@ -51,25 +52,29 @@ public class BatchingEngine implements Engine {
 	}
 
 	public void submitTasks(TaskEngine taskEngine, PoseStack stack, ClientLevel level) {
-		BatchingDrawManager.TransformSet drawSet = drawManager.get(RenderStage.AFTER_FINAL_END_BATCH);
+		BatchingTransformManager.TransformSet transformSet = transformManager.get(RenderStage.AFTER_FINAL_END_BATCH);
-		for (var entry : drawSet) {
+		for (var entry : transformSet) {
 			var renderType = entry.getKey();
-			var renderList = entry.getValue();
+			var transformCalls = entry.getValue();
 
 			int vertices = 0;
-			for (var transformSet : renderList) {
+			for (var transformCall : transformCalls) {
-				vertices += transformSet.getTotalVertexCount();
+				vertices += transformCall.getTotalVertexCount();
 			}
 
-			DrawBuffer buffer = drawManager.batchTracker.getBuffer(renderType);
+			if (vertices == 0) {
+				continue;
+			}
+
+			DrawBuffer buffer = drawTracker.getBuffer(renderType);
 			buffer.prepare(vertices);
 
 			int startVertex = 0;
-			for (var transformSet : renderList) {
+			for (var transformCall : transformCalls) {
-				transformSet.submitTasks(taskEngine, buffer, startVertex, stack, level);
+				transformCall.submitTasks(taskEngine, buffer, startVertex, stack, level);
-				startVertex += transformSet.getTotalVertexCount();
+				startVertex += transformCall.getTotalVertexCount();
+			}
 		}
-		};
 	}
 
 	@Override
@@ -81,7 +86,7 @@ public class BatchingEngine implements Engine {
 			return;
 		}
 
-		drawManager.batchTracker.endBatch();
+		drawTracker.drawAll();
 	}
 
 	@Override
@@ -103,7 +108,7 @@ public class BatchingEngine implements Engine {
 	@Override
 	public void delete() {
 		factories.clear();
-		drawManager.delete();
+		transformManager.delete();
 	}
 
 	@Override

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

@@ -2,6 +2,7 @@ package com.jozufozu.flywheel.backend.instancing.batching;
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.EnumMap;
 import java.util.Iterator;
 import java.util.List;
@@ -17,15 +18,19 @@ import com.jozufozu.flywheel.core.model.Model;
 
 import net.minecraft.client.renderer.RenderType;
 
-public class BatchingDrawManager {
+public class BatchingTransformManager {
 
 	private final List<UninitializedModel> uninitializedModels = new ArrayList<>();
 	private final List<CPUInstancer<?>> allInstancers = new ArrayList<>();
-	public final Map<RenderStage, TransformSet> renderLists = new EnumMap<>(RenderStage.class);
+	private final Map<RenderStage, TransformSet> transformSets = new EnumMap<>(RenderStage.class);
-	public final BatchDrawingTracker batchTracker = new BatchDrawingTracker();
+	private final Map<RenderStage, TransformSet> transformSetsView = Collections.unmodifiableMap(transformSets);
 
 	public TransformSet get(RenderStage stage) {
-		return renderLists.getOrDefault(stage, TransformSet.EMPTY);
+		return transformSets.getOrDefault(stage, TransformSet.EMPTY);
+	}
+
+	public Map<RenderStage, TransformSet> getTransformSetsView() {
+		return transformSetsView;
 	}
 
 	public void create(CPUInstancer<?> instancer, Model model) {
@@ -55,8 +60,8 @@ public class BatchingDrawManager {
 			var material = transformCall.getMaterial();
 			var renderType = material.getBatchingRenderType();
 
-//			renderLists.computeIfAbsent(material.getRenderStage(), TransformSet::new)
+//			transformSets.computeIfAbsent(material.getRenderStage(), TransformSet::new)
-			renderLists.computeIfAbsent(RenderStage.AFTER_FINAL_END_BATCH, TransformSet::new)
+			transformSets.computeIfAbsent(RenderStage.AFTER_FINAL_END_BATCH, TransformSet::new)
 					.put(renderType, transformCall);
 		}
 		allInstancers.add(instancer);

src/main/java/com/jozufozu/flywheel/backend/instancing/batching/CPUInstancerFactory.java

@@ -26,14 +26,6 @@ public class CPUInstancerFactory<D extends InstancedPart> implements InstancerFa
 		return models.computeIfAbsent(modelKey, this::createInstancer);
 	}
 
-//	/**
-//	 * Clear all instance data without freeing resources.
-//	 */
-//	public void clear() {
-//		models.values()
-//				.forEach(BatchedModel::clear);
-//	}
-
 	private CPUInstancer<D> createInstancer(Model model) {
 		var instancer = new CPUInstancer<>(type);
 		creationListener.accept(instancer, model);

src/main/java/com/jozufozu/flywheel/backend/instancing/batching/DrawBuffer.java

@@ -28,7 +28,8 @@ public class DrawBuffer {
 	private MemoryBlock memory;
 	private ByteBuffer buffer;
 
-	private int expectedVertices;
+	private boolean prepared;
+	private int vertexCount;
 
 	@ApiStatus.Internal
 	public DrawBuffer(VertexFormat format) {
@@ -45,11 +46,11 @@ public class DrawBuffer {
 	 * @throws IllegalStateException If the buffer is already in use.
 	 */
 	public void prepare(int vertexCount) {
-		if (expectedVertices != 0) {
+		if (prepared) {
-			throw new IllegalStateException("Already drawing!");
+			throw new IllegalStateException("Cannot prepare DrawBuffer twice!");
 		}
 
-		this.expectedVertices = vertexCount;
+		this.vertexCount = vertexCount;
 
 		// Add one extra vertex to uphold the vanilla assumption that BufferBuilders have at least
 		// enough buffer space for one more vertex. Rubidium checks for this extra space when popNextBuffer
@@ -65,9 +66,14 @@ public class DrawBuffer {
 		}
 
 		memory.clear();
+		prepared = true;
 	}
 
 	public ReusableVertexList slice(int startVertex, int vertexCount) {
+		if (!prepared) {
+			throw new IllegalStateException("Cannot slice DrawBuffer that is not prepared!");
+		}
+
 		ReusableVertexList vertexList = provider.createVertexList();
 		vertexList.ptr(memory.ptr());
 		vertexList.shiftPtr(startVertex);
@@ -80,19 +86,27 @@ public class DrawBuffer {
 	 * @param bufferBuilder The buffer builder to inject into.
 	 */
 	public void inject(BufferBuilderExtension bufferBuilder) {
+		if (!prepared) {
+			throw new IllegalStateException("Cannot inject DrawBuffer that is not prepared!");
+		}
+
 		buffer.clear();
-		bufferBuilder.flywheel$injectForRender(buffer, format, expectedVertices);
+		bufferBuilder.flywheel$injectForRender(buffer, format, vertexCount);
+	}
+
+	public boolean isPrepared() {
+		return prepared;
 	}
 
 	public int getVertexCount() {
-		return expectedVertices;
+		return vertexCount;
 	}
 
 	/**
 	 * @return {@code true} if the buffer has any vertices.
 	 */
 	public boolean hasVertices() {
-		return expectedVertices > 0;
+		return vertexCount > 0;
 	}
 
 	/**
@@ -101,7 +115,8 @@ public class DrawBuffer {
 	 * Does not clear the backing buffer.
 	 */
 	public void reset() {
-		this.expectedVertices = 0;
+		prepared = false;
+		vertexCount = 0;
 	}
 
 	public void free() {

src/main/java/com/jozufozu/flywheel/backend/instancing/instancing/GPUInstancer.java

@@ -25,7 +25,6 @@ public class GPUInstancer<D extends InstancedPart> extends AbstractInstancer<D>
 
 	boolean anyToUpdate;
 
-
 	public GPUInstancer(StructType<D> type) {
 		super(type);
 		this.instanceFormat = type.getLayout();