RenderLayers directly store DrawBuffers

- Inspired by pepper's BlockEntityTypeExtensions
 - Document the batching engine internals.

src/main/java/com/jozufozu/flywheel/backend/Backend.java

@@ -1,6 +1,10 @@
 package com.jozufozu.flywheel.backend;
 
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 
 import javax.annotation.Nullable;
 

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

@@ -0,0 +1,83 @@
+package com.jozufozu.flywheel.backend.instancing;
+
+import java.util.HashSet;
+import java.util.Set;
+
+import com.jozufozu.flywheel.backend.model.BufferBuilderExtension;
+import com.jozufozu.flywheel.backend.model.DirectVertexConsumer;
+import com.mojang.blaze3d.vertex.BufferBuilder;
+
+import net.minecraft.client.renderer.RenderType;
+
+public class BatchDrawingTracker {
+
+	protected final Set<RenderType> activeTypes = new HashSet<>();
+	private final BufferBuilder scratch;
+
+	public BatchDrawingTracker() {
+		scratch = new BufferBuilder(8);
+
+		((BufferBuilderExtension) scratch).flywheel$freeBuffer();
+	}
+
+	/**
+	 * Get a direct vertex consumer for drawing the given number of vertices to the given RenderType.
+	 * @param renderType The RenderType to draw to.
+	 * @param vertexCount The number of vertices that will be drawn.
+	 * @return A direct vertex consumer.
+	 */
+	public DirectVertexConsumer getDirectConsumer(RenderType renderType, int vertexCount) {
+		activeTypes.add(renderType);
+		return RenderTypeExtension.getDrawBuffer(renderType)
+				.begin(vertexCount);
+	}
+
+	/**
+	 * Draws all active DrawBuffers and reset them.
+	 */
+	public void endBatch() {
+		// TODO: when/if this causes trouble with shaders, try to inject our BufferBuilders
+		//  into the RenderBuffers from context.
+
+		for (RenderType renderType : activeTypes) {
+			_draw(renderType);
+		}
+
+		activeTypes.clear();
+	}
+
+	/**
+	 * Draw and reset the DrawBuffer for the given RenderType.
+	 * @param renderType The RenderType to draw.
+	 */
+	public void endBatch(RenderType renderType) {
+		_draw(renderType);
+
+		activeTypes.remove(renderType);
+	}
+
+	/**
+	 * Resets all DrawBuffers to 0 vertices.
+	 */
+	public void clear() {
+		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/DrawBuffer.java

@@ -0,0 +1,77 @@
+package com.jozufozu.flywheel.backend.instancing;
+
+import java.nio.ByteBuffer;
+
+import com.jozufozu.flywheel.backend.model.BufferBuilderExtension;
+import com.jozufozu.flywheel.backend.model.DirectVertexConsumer;
+import com.mojang.blaze3d.platform.MemoryTracker;
+import com.mojang.blaze3d.vertex.VertexFormat;
+
+import net.minecraft.client.renderer.RenderType;
+
+/**
+ * A byte buffer that can be used to draw vertices through a {@link DirectVertexConsumer}.
+ *
+ * The number of vertices needs to be known ahead of time.
+ */
+public class DrawBuffer {
+
+	private final RenderType parent;
+	private ByteBuffer backingBuffer;
+	private int expectedVertices;
+
+	public DrawBuffer(RenderType parent) {
+		this.parent = parent;
+	}
+
+	/**
+	 * Creates a direct vertex consumer that can be used to write vertices into this buffer.
+	 * @param vertexCount The number of vertices to reserve memory for.
+	 * @return A direct vertex consumer.
+	 * @throws IllegalStateException If the buffer is already in use.
+	 */
+	public DirectVertexConsumer begin(int vertexCount) {
+		if (expectedVertices != 0) {
+			throw new IllegalStateException("Already drawing");
+		}
+
+		this.expectedVertices = vertexCount;
+
+		VertexFormat format = parent.format();
+
+		int byteSize = format.getVertexSize() * vertexCount;
+
+		if (backingBuffer == null) {
+			backingBuffer = MemoryTracker.create(byteSize);
+		}
+		if (byteSize > backingBuffer.capacity()) {
+			backingBuffer = MemoryTracker.resize(backingBuffer, byteSize);
+		}
+
+		return new DirectVertexConsumer(backingBuffer, format, vertexCount);
+	}
+
+	/**
+	 * Injects the backing buffer into the given builder and prepares it for rendering.
+	 * @param bufferBuilder The buffer builder to inject into.
+	 */
+	public void inject(BufferBuilderExtension bufferBuilder) {
+		bufferBuilder.flywheel$injectForRender(backingBuffer, parent.format(), expectedVertices);
+	}
+
+	/**
+	 * @return {@code true} if the buffer has any vertices.
+	 */
+	public boolean hasVertices() {
+		return expectedVertices > 0;
+	}
+
+	/**
+	 * Reset the draw buffer to have no vertices.
+	 *
+	 * Does not clear the backing buffer.
+	 */
+	public void reset() {
+		this.expectedVertices = 0;
+	}
+}

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

@@ -0,0 +1,25 @@
+package com.jozufozu.flywheel.backend.instancing;
+
+import net.minecraft.client.renderer.RenderType;
+
+/**
+ * Duck interface to make RenderType store a DrawBuffer.
+ *
+ * @see RenderType
+ */
+public interface RenderTypeExtension {
+
+	/**
+	 * @return The DrawBuffer associated with this RenderType.
+	 */
+	DrawBuffer flywheel$getDrawBuffer();
+
+	/**
+	 * Helper function to cast a RenderType to a RenderTypeExtension and get its DrawBuffer.
+	 * @param type The RenderType to get the DrawBuffer from.
+	 * @return The DrawBuffer associated with the given RenderType.
+	 */
+	static DrawBuffer getDrawBuffer(RenderType type) {
+		return ((RenderTypeExtension) type).flywheel$getDrawBuffer();
+	}
+}

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

@@ -1,79 +0,0 @@
-package com.jozufozu.flywheel.backend.instancing;
-
-import java.nio.ByteBuffer;
-import java.util.HashMap;
-import java.util.Map;
-
-import com.jozufozu.flywheel.backend.model.BufferBuilderHack;
-import com.jozufozu.flywheel.backend.model.DirectVertexConsumer;
-import com.mojang.blaze3d.platform.MemoryTracker;
-import com.mojang.blaze3d.vertex.BufferBuilder;
-import com.mojang.blaze3d.vertex.VertexFormat;
-
-import net.minecraft.client.renderer.RenderType;
-
-public class SuperBufferSource {
-
-	protected final Map<RenderType, DrawBuffer> buffers = new HashMap<>();
-	private final BufferBuilder scratch;
-
-	public SuperBufferSource() {
-		scratch = new BufferBuilder(8);
-
-		((BufferBuilderHack) scratch).flywheel$freeBuffer();
-	}
-
-	public DirectVertexConsumer getBuffer(RenderType renderType, int vertexCount) {
-		return buffers.computeIfAbsent(renderType, DrawBuffer::new)
-				.begin(vertexCount);
-	}
-
-	public void endBatch() {
-		// TODO: when/if this causes trouble with shaders, try to inject our BufferBuilders
-		//  into the RenderBuffers from context.
-
-		BufferBuilderHack hack = (BufferBuilderHack) scratch;
-
-		for (Map.Entry<RenderType, DrawBuffer> entry : buffers.entrySet()) {
-			DrawBuffer builder = entry.getValue();
-
-			if (builder.expectedVertices > 0) {
-				RenderType type = entry.getKey();
-
-				hack.flywheel$hackBegin(builder.backingBuffer, type.format(), builder.expectedVertices);
-
-				type.end(scratch, 0, 0, 0);
-
-				builder.expectedVertices = 0;
-			}
-		}
-	}
-
-	private static class DrawBuffer {
-
-		private final RenderType type;
-		private ByteBuffer backingBuffer;
-		private int expectedVertices;
-
-		public DrawBuffer(RenderType type) {
-			this.type = type;
-		}
-
-		public DirectVertexConsumer begin(int vertexCount) {
-			this.expectedVertices = vertexCount;
-
-			VertexFormat format = type.format();
-
-			int byteSize = format.getVertexSize() * vertexCount;
-
-			if (backingBuffer == null) {
-				backingBuffer = MemoryTracker.create(byteSize);
-			} if (byteSize > backingBuffer.capacity()) {
-				backingBuffer = MemoryTracker.resize(backingBuffer, byteSize);
-			}
-
-			return new DirectVertexConsumer(backingBuffer, format, vertexCount);
-		}
-
-	}
-}

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

@@ -7,7 +7,8 @@ import com.jozufozu.flywheel.api.InstanceData;
 import com.jozufozu.flywheel.api.MaterialGroup;
 import com.jozufozu.flywheel.api.struct.Batched;
 import com.jozufozu.flywheel.api.struct.StructType;
-import com.jozufozu.flywheel.backend.instancing.SuperBufferSource;
+import com.jozufozu.flywheel.backend.OptifineHandler;
+import com.jozufozu.flywheel.backend.instancing.BatchDrawingTracker;
 import com.jozufozu.flywheel.backend.instancing.TaskEngine;
 import com.jozufozu.flywheel.backend.model.DirectVertexConsumer;
 import com.mojang.blaze3d.vertex.PoseStack;
@@ -34,7 +35,7 @@ public class BatchedMaterialGroup implements MaterialGroup {
 		}
 	}
 
-	public void render(PoseStack stack, SuperBufferSource source, TaskEngine pool) {
+	public void render(PoseStack stack, BatchDrawingTracker source, TaskEngine pool) {
 
 		int vertexCount = 0;
 		for (BatchedMaterial<?> material : materials.values()) {
@@ -44,14 +45,14 @@ public class BatchedMaterialGroup implements MaterialGroup {
 			}
 		}
 
-		DirectVertexConsumer consumer = source.getBuffer(state, vertexCount);
+		DirectVertexConsumer consumer = source.getDirectConsumer(state, vertexCount);
 
 		// avoids rendering garbage, but doesn't fix the issue of some instances not being buffered
 		consumer.memSetZero();
 
 		for (BatchedMaterial<?> material : materials.values()) {
 			for (CPUInstancer<?> instancer : material.models.values()) {
-				instancer.sbb.context.outputColorDiffuse = !consumer.hasOverlay();
+				instancer.sbb.context.outputColorDiffuse = !consumer.hasOverlay() && !OptifineHandler.usingShaders();
 				instancer.submitTasks(stack, pool, consumer);
 			}
 		}

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

@@ -6,8 +6,8 @@ import java.util.Map;
 
 import com.jozufozu.flywheel.api.MaterialGroup;
 import com.jozufozu.flywheel.backend.RenderLayer;
+import com.jozufozu.flywheel.backend.instancing.BatchDrawingTracker;
 import com.jozufozu.flywheel.backend.instancing.Engine;
-import com.jozufozu.flywheel.backend.instancing.SuperBufferSource;
 import com.jozufozu.flywheel.backend.instancing.TaskEngine;
 import com.jozufozu.flywheel.event.RenderLayerEvent;
 import com.mojang.blaze3d.platform.Lighting;
@@ -21,7 +21,7 @@ import net.minecraft.core.Vec3i;
 public class BatchingEngine implements Engine {
 
 	private final Map<RenderLayer, Map<RenderType, BatchedMaterialGroup>> layers;
-	private final SuperBufferSource superBufferSource = new SuperBufferSource();
+	private final BatchDrawingTracker batchTracker = new BatchDrawingTracker();
 
 	public BatchingEngine() {
 		this.layers = new EnumMap<>(RenderLayer.class);
@@ -43,14 +43,11 @@ public class BatchingEngine implements Engine {
 
 	@Override
 	public void render(TaskEngine taskEngine, RenderLayerEvent event) {
-
 		Map<RenderType, BatchedMaterialGroup> groups = layers.get(event.getLayer());
 		for (BatchedMaterialGroup group : groups.values()) {
-			group.render(event.stack, superBufferSource, taskEngine);
+			group.render(event.stack, batchTracker, taskEngine);
 		}
 
-		taskEngine.syncPoint();
-
 		// FIXME: this probably breaks some vanilla stuff but it works much better for flywheel
 		Matrix4f mat = new Matrix4f();
 		mat.setIdentity();
@@ -60,7 +57,8 @@ public class BatchingEngine implements Engine {
 			Lighting.setupLevel(mat);
 		}
 
-		superBufferSource.endBatch();
+		taskEngine.syncPoint();
+		batchTracker.endBatch();
 	}
 
 	@Override

src/main/java/com/jozufozu/flywheel/backend/model/BufferBuilderExtension.java

@@ -0,0 +1,27 @@
+package com.jozufozu.flywheel.backend.model;
+
+import java.nio.ByteBuffer;
+
+import com.mojang.blaze3d.vertex.BufferBuilder;
+import com.mojang.blaze3d.vertex.VertexFormat;
+
+/**
+ * Duck interface used on {@link BufferBuilder} to provide lower level access to the backing memory.
+ *
+ * @see com.jozufozu.flywheel.mixin.BufferBuilderMixin
+ */
+public interface BufferBuilderExtension {
+
+	/**
+	 * Frees the internal ByteBuffer, if it exists.
+	 */
+	void flywheel$freeBuffer();
+
+	/**
+	 * Prepares the BufferBuilder for drawing the contents of the given buffer.
+	 * @param buffer The buffer to draw.
+	 * @param format The format of the buffer.
+	 * @param vertexCount The number of vertices in the buffer.
+	 */
+	void flywheel$injectForRender(ByteBuffer buffer, VertexFormat format, int vertexCount);
+}

src/main/java/com/jozufozu/flywheel/backend/model/BufferBuilderHack.java

@@ -1,16 +0,0 @@
-package com.jozufozu.flywheel.backend.model;
-
-import java.nio.ByteBuffer;
-
-import com.mojang.blaze3d.vertex.BufferBuilder;
-import com.mojang.blaze3d.vertex.VertexFormat;
-
-/**
- * Duck interface used on {@link BufferBuilder} to provide lower level access to the backing memory.
- */
-public interface BufferBuilderHack {
-
-	void flywheel$freeBuffer();
-
-	void flywheel$hackBegin(ByteBuffer buffer, VertexFormat format, int vertexCount);
-}

src/main/java/com/jozufozu/flywheel/backend/model/DirectVertexConsumer.java

@@ -13,7 +13,7 @@ import com.mojang.blaze3d.vertex.VertexFormatElement;
 /**
  * An unsafe vertex consumer allowing for unchecked writes into a ByteBuffer.
  *
- * @see BufferBuilderHack
+ * @see BufferBuilderExtension
  */
 public class DirectVertexConsumer implements VertexConsumer {
 	public final VertexFormat format;

src/main/java/com/jozufozu/flywheel/core/LastActiveCamera.java

@@ -2,6 +2,11 @@ package com.jozufozu.flywheel.core;
 
 import net.minecraft.client.Camera;
 
+/**
+ * A class tracking which object last had {@link Camera#setup} called on it.
+ *
+ * @see com.jozufozu.flywheel.mixin.CameraMixin
+ */
 public class LastActiveCamera {
 
 	private static Camera camera;

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

@@ -9,13 +9,13 @@ import org.lwjgl.system.MemoryUtil;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
 
-import com.jozufozu.flywheel.backend.model.BufferBuilderHack;
+import com.jozufozu.flywheel.backend.model.BufferBuilderExtension;
 import com.mojang.blaze3d.vertex.BufferBuilder;
 import com.mojang.blaze3d.vertex.VertexFormat;
 import com.mojang.blaze3d.vertex.VertexFormatElement;
 
 @Mixin(BufferBuilder.class)
-public abstract class BufferBuilderMixin implements BufferBuilderHack {
+public abstract class BufferBuilderMixin implements BufferBuilderExtension {
 	@Shadow
 	private ByteBuffer buffer;
 
@@ -50,7 +50,7 @@ public abstract class BufferBuilderMixin implements BufferBuilderHack {
 	}
 
 	@Override
-	public void flywheel$hackBegin(@Nonnull ByteBuffer buffer, @Nonnull VertexFormat format, int vertexCount) {
+	public void flywheel$injectForRender(@Nonnull ByteBuffer buffer, @Nonnull VertexFormat format, int vertexCount) {
 		this.building = true;
 		this.mode = VertexFormat.Mode.QUADS;
 

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

@@ -0,0 +1,24 @@
+package com.jozufozu.flywheel.mixin;
+
+import javax.annotation.Nonnull;
+
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Unique;
+
+import com.jozufozu.flywheel.backend.instancing.DrawBuffer;
+import com.jozufozu.flywheel.backend.instancing.RenderTypeExtension;
+
+import net.minecraft.client.renderer.RenderType;
+
+@Mixin(RenderType.class)
+public class RenderTypeMixin implements RenderTypeExtension {
+
+	@Unique
+	private final DrawBuffer flywheel$drawBuffer = new DrawBuffer((RenderType) (Object) this);
+
+	@Override
+	@Nonnull
+	public DrawBuffer flywheel$getDrawBuffer() {
+		return flywheel$drawBuffer;
+	}
+}

src/main/resources/flywheel.mixins.json

@@ -1,36 +1,37 @@
 {
-  "required": true,
+    "required": true,
-  "minVersion": "0.8",
+    "minVersion": "0.8",
-  "package": "com.jozufozu.flywheel.mixin",
+    "package": "com.jozufozu.flywheel.mixin",
-  "compatibilityLevel": "JAVA_17",
+    "compatibilityLevel": "JAVA_17",
-  "refmap": "flywheel.refmap.json",
+    "refmap": "flywheel.refmap.json",
-  "client": [
+    "client": [
-    "BlockEntityTypeMixin",
+        "BlockEntityTypeMixin",
-    "BufferBuilderMixin",
+        "BufferBuilderMixin",
-    "BufferUploaderAccessor",
+        "BufferUploaderAccessor",
-    "CameraMixin",
+        "CameraMixin",
-    "CancelEntityRenderMixin",
+        "CancelEntityRenderMixin",
-    "ChunkRebuildHooksMixin",
+        "ChunkRebuildHooksMixin",
-    "EntityTypeMixin",
+        "EntityTypeMixin",
-    "FixFabulousDepthMixin",
+        "FixFabulousDepthMixin",
-    "FrustumMixin",
+        "FrustumMixin",
-    "InstanceAddMixin",
+        "InstanceAddMixin",
-    "InstanceRemoveMixin",
+        "InstanceRemoveMixin",
-    "LevelRendererAccessor",
+        "LevelRendererAccessor",
-    "LevelRendererMixin",
+        "LevelRendererMixin",
-    "PausedPartialTickAccessor",
+        "PausedPartialTickAccessor",
-    "RenderTexturesMixin",
+        "RenderTexturesMixin",
-    "ShaderCloseMixin",
+        "RenderTypeMixin",
-    "ShaderInstanceAccessor",
+        "ShaderCloseMixin",
-    "atlas.AtlasDataMixin",
+        "ShaderInstanceAccessor",
-    "atlas.SheetDataAccessor",
+        "atlas.AtlasDataMixin",
-    "light.LightUpdateMixin",
+        "atlas.SheetDataAccessor",
-    "light.NetworkLightUpdateMixin",
+        "light.LightUpdateMixin",
-    "matrix.Matrix3fMixin",
+        "light.NetworkLightUpdateMixin",
-    "matrix.Matrix4fMixin",
+        "matrix.Matrix3fMixin",
-    "matrix.PoseStackMixin"
+        "matrix.Matrix4fMixin",
-  ],
+        "matrix.PoseStackMixin"
-  "injectors": {
+    ],
-    "defaultRequire": 1
+    "injectors": {
+        "defaultRequire": 1
   }
 }