Add batching stage support

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

@@ -1,8 +1,10 @@
 package com.jozufozu.flywheel.backend.instancing.batching;
 
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.Set;
 
+import com.jozufozu.flywheel.api.RenderStage;
 import com.jozufozu.flywheel.extension.BufferBuilderExtension;
 import com.jozufozu.flywheel.extension.RenderTypeExtension;
 import com.mojang.blaze3d.vertex.BufferBuilder;
@@ -10,7 +12,6 @@ import com.mojang.blaze3d.vertex.BufferBuilder;
 import net.minecraft.client.renderer.RenderType;
 
 public class BatchingDrawTracker {
-
 	private final Set<RenderType> activeTypes = new HashSet<>();
 	private final BufferBuilder scratch;
 
@@ -20,13 +21,17 @@ public class BatchingDrawTracker {
 		((BufferBuilderExtension) scratch).flywheel$freeBuffer();
 	}
 
-	public DrawBuffer getBuffer(RenderType renderType) {
+	public DrawBuffer getBuffer(RenderType renderType, RenderStage stage) {
+		return getBufferSet(renderType).getBuffer(stage);
+	}
+
+	public DrawBufferSet getBufferSet(RenderType renderType) {
 		activeTypes.add(renderType);
-		return RenderTypeExtension.getDrawBuffer(renderType);
+		return RenderTypeExtension.getDrawBufferSet(renderType);
 	}
 
 	/**
-	 * Draw and reset the DrawBuffer for the given RenderType.
+	 * Draw and reset all DrawBuffers for the given RenderType.
 	 * @param renderType The RenderType to draw.
 	 */
 	public void draw(RenderType renderType) {
@@ -36,7 +41,27 @@ public class BatchingDrawTracker {
 	}
 
 	/**
-	 * Draws all active DrawBuffers and reset them.
+	 * Draw and reset all DrawBuffers for the given RenderStage.
+	 * @param stage The RenderStage to draw.
+	 */
+	public void draw(RenderStage stage) {
+		Iterator<RenderType> iterator = activeTypes.iterator();
+		while (iterator.hasNext()) {
+			RenderType renderType = iterator.next();
+			DrawBufferSet bufferSet = RenderTypeExtension.getDrawBufferSet(renderType);
+			DrawBuffer buffer = bufferSet.deactivateBuffer(stage);
+			if (buffer == null) {
+				continue;
+			}
+			if (bufferSet.getActiveStagesView().isEmpty()) {
+				iterator.remove();
+			}
+			_draw(buffer, renderType);
+		}
+	}
+
+	/**
+	 * Draw and reset all active DrawBuffers.
 	 */
 	public void drawAll() {
 		for (RenderType renderType : activeTypes) {
@@ -47,8 +72,14 @@ public class BatchingDrawTracker {
 	}
 
 	private void _draw(RenderType renderType) {
-		DrawBuffer buffer = RenderTypeExtension.getDrawBuffer(renderType);
+		DrawBufferSet bufferSet = RenderTypeExtension.getDrawBufferSet(renderType);
+		for (RenderStage stage : bufferSet.getActiveStagesView()) {
+			DrawBuffer buffer = bufferSet.deactivateBuffer(stage);
+			_draw(buffer, renderType);
+		}
+	}
 
+	private void _draw(DrawBuffer buffer, RenderType renderType) {
 		if (buffer.hasVertices()) {
 			BufferBuilderExtension scratch = (BufferBuilderExtension) this.scratch;
 			buffer.inject(scratch);
@@ -59,15 +90,14 @@ public class BatchingDrawTracker {
 	}
 
 	/**
-	 * Resets all DrawBuffers to 0 vertices.
+	 * Reset all active DrawBuffers.
 	 */
 	public void reset() {
 		for (RenderType type : activeTypes) {
-			RenderTypeExtension.getDrawBuffer(type)
+			RenderTypeExtension.getDrawBufferSet(type)
 					.reset();
 		}
 
 		activeTypes.clear();
 	}
-
 }

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

@@ -23,7 +23,6 @@ import net.minecraft.core.Vec3i;
 import net.minecraft.world.phys.Vec3;
 
 public class BatchingEngine implements Engine {
-
 	protected final BatchingTransformManager transformManager = new BatchingTransformManager();
 	protected final BatchingDrawTracker drawTracker = new BatchingDrawTracker();
 	protected final Map<StructType<?>, CPUInstancerFactory<?>> factories = new HashMap<>();
@@ -54,41 +53,38 @@ public class BatchingEngine implements Engine {
 	}
 
 	public void submitTasks(TaskEngine taskEngine, PoseStack stack, ClientLevel level) {
-		BatchingTransformManager.TransformSet transformSet = transformManager.get(RenderStage.AFTER_FINAL_END_BATCH);
-		for (var entry : transformSet) {
-			var renderType = entry.getKey();
-			var transformCalls = entry.getValue();
+		for (var transformSetEntry : transformManager.getTransformSetsView().entrySet()) {
+			var stage = transformSetEntry.getKey();
+			var transformSet = transformSetEntry.getValue();
 
-			int vertices = 0;
-			for (var transformCall : transformCalls) {
-				vertices += transformCall.getTotalVertexCount();
-			}
+			for (var entry : transformSet) {
+				var renderType = entry.getKey();
+				var transformCalls = entry.getValue();
 
-			if (vertices == 0) {
-				continue;
-			}
+				int vertices = 0;
+				for (var transformCall : transformCalls) {
+					vertices += transformCall.getTotalVertexCount();
+				}
 
-			DrawBuffer buffer = drawTracker.getBuffer(renderType);
-			buffer.prepare(vertices);
+				if (vertices == 0) {
+					continue;
+				}
 
-			int startVertex = 0;
-			for (var transformCall : transformCalls) {
-				transformCall.submitTasks(taskEngine, buffer, startVertex, stack, level);
-				startVertex += transformCall.getTotalVertexCount();
+				DrawBuffer buffer = drawTracker.getBuffer(renderType, stage);
+				buffer.prepare(vertices);
+
+				int startVertex = 0;
+				for (var transformCall : transformCalls) {
+					transformCall.submitTasks(taskEngine, buffer, startVertex, stack, level);
+					startVertex += transformCall.getTotalVertexCount();
+				}
 			}
 		}
 	}
 
 	@Override
 	public void renderStage(TaskEngine taskEngine, RenderContext context, RenderStage stage) {
-		// FIXME: properly support material stages
-		// This also breaks block outlines on batched block entities
-		// and makes translucent blocks occlude everything Flywheel renders
-		if (stage != RenderStage.AFTER_FINAL_END_BATCH) {
-			return;
-		}
-
-		drawTracker.drawAll();
+		drawTracker.draw(stage);
 	}
 
 	@Override

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

@@ -19,7 +19,6 @@ import com.jozufozu.flywheel.core.model.Model;
 import net.minecraft.client.renderer.RenderType;
 
 public class BatchingTransformManager {
-
 	private final List<UninitializedModel> uninitializedModels = new ArrayList<>();
 	private final List<CPUInstancer<?>> allInstancers = new ArrayList<>();
 	private final Map<RenderStage, TransformSet> transformSets = new EnumMap<>(RenderStage.class);
@@ -60,15 +59,13 @@ public class BatchingTransformManager {
 			var material = transformCall.getMaterial();
 			var renderType = material.getBatchingRenderType();
 
-//			transformSets.computeIfAbsent(material.getRenderStage(), TransformSet::new)
-			transformSets.computeIfAbsent(RenderStage.AFTER_FINAL_END_BATCH, TransformSet::new)
+			transformSets.computeIfAbsent(material.getRenderStage(), TransformSet::new)
 					.put(renderType, transformCall);
 		}
 		allInstancers.add(instancer);
 	}
 
 	public static class TransformSet implements Iterable<Map.Entry<RenderType, Collection<TransformCall<?>>>> {
-
 		public static final TransformSet EMPTY = new TransformSet(ImmutableListMultimap.of());
 
 		final ListMultimap<RenderType, TransformCall<?>> transformCalls;

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

@@ -4,8 +4,6 @@ import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.List;
 
-import org.jetbrains.annotations.ApiStatus;
-
 import com.jozufozu.flywheel.api.vertex.ReusableVertexList;
 import com.jozufozu.flywheel.api.vertex.VertexListProvider;
 import com.jozufozu.flywheel.backend.memory.MemoryBlock;
@@ -31,11 +29,10 @@ public class DrawBuffer {
 	private boolean prepared;
 	private int vertexCount;
 
-	@ApiStatus.Internal
-	public DrawBuffer(VertexFormat format) {
+	DrawBuffer(VertexFormat format, int stride, VertexListProvider provider) {
 		this.format = format;
-		stride = format.getVertexSize();
-		provider = VertexListProvider.get(format);
+		this.stride = stride;
+		this.provider = provider;
 
 		ALL.add(this);
 	}
@@ -120,6 +117,9 @@ public class DrawBuffer {
 	}
 
 	public void free() {
+		if (memory == null) {
+			return;
+		}
 		memory.free();
 		memory = null;
 		buffer = null;

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

@@ -0,0 +1,66 @@
+package com.jozufozu.flywheel.backend.instancing.batching;
+
+import java.util.Collections;
+import java.util.EnumMap;
+import java.util.EnumSet;
+import java.util.Map;
+import java.util.Set;
+
+import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.Nullable;
+
+import com.jozufozu.flywheel.api.RenderStage;
+import com.jozufozu.flywheel.api.vertex.VertexListProvider;
+import com.mojang.blaze3d.vertex.VertexFormat;
+
+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);
+	private final Set<RenderStage> activeStages = EnumSet.noneOf(RenderStage.class);
+	private final Set<RenderStage> activeStagesView = Collections.unmodifiableSet(activeStages);
+
+	@ApiStatus.Internal
+	public DrawBufferSet(VertexFormat format) {
+		this.format = format;
+		stride = format.getVertexSize();
+		provider = VertexListProvider.get(format);
+	}
+
+	public Set<RenderStage> getActiveStagesView() {
+		return activeStagesView;
+	}
+
+	public DrawBuffer getBuffer(RenderStage stage) {
+		activeStages.add(stage);
+		return buffers.computeIfAbsent(stage, $ -> createBuffer());
+	}
+
+	@Nullable
+	public DrawBuffer deactivateBuffer(RenderStage stage) {
+		if (activeStages.remove(stage)) {
+			return buffers.get(stage);
+		}
+		return null;
+	}
+
+	public void reset(RenderStage stage) {
+		if (activeStages.remove(stage)) {
+			buffers.get(stage).reset();
+		}
+	}
+
+	public void reset() {
+		for (RenderStage stage : activeStages) {
+			buffers.get(stage).reset();
+		}
+
+		activeStages.clear();
+	}
+
+	private DrawBuffer createBuffer() {
+		return new DrawBuffer(format, stride, provider);
+	}
+}

src/main/java/com/jozufozu/flywheel/extension/RenderTypeExtension.java

@@ -1,27 +1,26 @@
 package com.jozufozu.flywheel.extension;
 
-import com.jozufozu.flywheel.backend.instancing.batching.DrawBuffer;
+import com.jozufozu.flywheel.backend.instancing.batching.DrawBufferSet;
 
 import net.minecraft.client.renderer.RenderType;
 
 /**
- * Duck interface to make RenderType store a DrawBuffer.
+ * Duck interface to make RenderType store a DrawBufferSet.
  *
  * @see RenderType
  */
 public interface RenderTypeExtension {
+	/**
+	 * @return The DrawBufferSet associated with this RenderType.
+	 */
+	DrawBufferSet flywheel$getDrawBufferSet();
 
 	/**
-	 * @return The DrawBuffer associated with this RenderType.
+	 * Helper function to cast a RenderType to a RenderTypeExtension and get its DrawBufferSet.
+	 * @param type The RenderType to get the DrawBufferSet from.
+	 * @return The DrawBufferSet associated with the given 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();
+	static DrawBufferSet getDrawBufferSet(RenderType type) {
+		return ((RenderTypeExtension) type).flywheel$getDrawBufferSet();
 	}
 }

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

@@ -4,7 +4,7 @@ import org.jetbrains.annotations.NotNull;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Unique;
 
-import com.jozufozu.flywheel.backend.instancing.batching.DrawBuffer;
+import com.jozufozu.flywheel.backend.instancing.batching.DrawBufferSet;
 import com.jozufozu.flywheel.extension.RenderTypeExtension;
 
 import net.minecraft.client.renderer.RenderType;
@@ -12,14 +12,14 @@ import net.minecraft.client.renderer.RenderType;
 @Mixin(RenderType.class)
 public class RenderTypeMixin implements RenderTypeExtension {
 	@Unique
-	private DrawBuffer flywheel$drawBuffer;
+	private DrawBufferSet flywheel$drawBufferSet;
 
 	@Override
 	@NotNull
-	public DrawBuffer flywheel$getDrawBuffer() {
-		if (flywheel$drawBuffer == null) {
-			flywheel$drawBuffer = new DrawBuffer(((RenderType) (Object) this).format());
+	public DrawBufferSet flywheel$getDrawBufferSet() {
+		if (flywheel$drawBufferSet == null) {
+			flywheel$drawBufferSet = new DrawBufferSet(((RenderType) (Object) this).format());
 		}
-		return flywheel$drawBuffer;
+		return flywheel$drawBufferSet;
 	}
 }