Bit off a bigger batch than I could chew

- Implement crumbling for batching engine.
- It's ugly! (the code)
- Need to support translucency sorting in batching.
  - Plumb RenderType#sortOnUpload through to DrawBuffer.
  - Allocate room in DrawBuffer for BufferBuilder to write indices.
  - Set nextElementByte on injectForRender.
- Remove RenderStage from DrawBuffer, isn't really needed.
- Make some fields public, so they can be accessed by BatchedCrumbling.
- Track TransformCalls in BatchedInstancer.
-

src/main/java/com/jozufozu/flywheel/api/backend/Engine.java

@@ -30,6 +30,8 @@ public interface Engine extends InstancerProvider {
 
 	/**
 	 * Render the given instances as a crumbling overlay.
+	 * <br>
+	 * This is guaranteed to be called between the first and last calls to {@link #renderStage}.
 	 *
 	 * @param executor        The task executor running the frame plan.
 	 * @param context         The render context for this frame.

src/main/java/com/jozufozu/flywheel/api/vertex/VertexList.java

@@ -1,5 +1,8 @@
 package com.jozufozu.flywheel.api.vertex;
 
+import org.joml.Vector3f;
+import org.joml.Vector4f;
+
 /**
  * A read only view of a vertex buffer.
  *
@@ -73,4 +76,12 @@ public interface VertexList {
 	default boolean isEmpty() {
 		return vertexCount() == 0;
 	}
+
+    default Vector3f getNormal(int i, Vector3f dest) {
+		return dest.set(normalX(i), normalY(i), normalZ(i));
+	}
+
+	default Vector4f getPos(int i, Vector4f dest) {
+		return dest.set(x(i), y(i), z(i));
+	}
 }

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

@@ -0,0 +1,109 @@
+package com.jozufozu.flywheel.backend.engine.batching;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.jetbrains.annotations.NotNull;
+
+import com.jozufozu.flywheel.api.backend.Engine;
+import com.jozufozu.flywheel.api.event.RenderStage;
+import com.jozufozu.flywheel.api.instance.Instance;
+import com.jozufozu.flywheel.api.model.Mesh;
+import com.jozufozu.flywheel.api.vertex.ReusableVertexList;
+import com.jozufozu.flywheel.backend.engine.InstanceHandleImpl;
+import com.jozufozu.flywheel.extension.RenderTypeExtension;
+
+import it.unimi.dsi.fastutil.Pair;
+import net.minecraft.client.renderer.RenderType;
+import net.minecraft.client.resources.model.ModelBakery;
+
+public class BatchedCrumbling {
+	static void renderCrumbling(List<Engine.CrumblingBlock> crumblingBlocks, BatchContext batchContext, BatchedDrawTracker drawTracker) {
+		var instancesPerType = doCrumblingSort(crumblingBlocks);
+
+		for (var entry : instancesPerType.entrySet()) {
+			// TODO: separate concept of renderstage from drawbuffer
+			var drawBuffer = RenderTypeExtension.getDrawBufferSet(entry.getKey())
+					.getBuffer(RenderStage.AFTER_BLOCK_ENTITIES);
+
+			var bucket = entry.getValue();
+
+			drawBuffer.prepare(bucket.vertexCount);
+
+			ReusableVertexList vertexList = drawBuffer.slice(0, 0);
+			long basePtr = vertexList.ptr();
+
+			int totalVertices = 0;
+
+			for (var pair : bucket.instances) {
+				var instance = pair.first();
+				var instancer = pair.second();
+
+				totalVertices += bufferOne(instancer, totalVertices, vertexList, drawBuffer, instance);
+			}
+
+			vertexList.ptr(basePtr);
+			vertexList.vertexCount(totalVertices);
+
+			// apply these in bulk
+			BatchingTransforms.applyDecalUVs(vertexList);
+			BatchingTransforms.applyMatrices(vertexList, batchContext.matrices());
+
+			drawTracker._draw(drawBuffer);
+		}
+	}
+
+	@NotNull
+	private static Map<RenderType, CrumblingBucket> doCrumblingSort(List<Engine.CrumblingBlock> crumblingBlocks) {
+		Map<RenderType, CrumblingBucket> out = new HashMap<>();
+
+		for (Engine.CrumblingBlock crumblingBlock : crumblingBlocks) {
+			var crumblingType = ModelBakery.DESTROY_TYPES.get(crumblingBlock.progress());
+
+			for (Instance instance : crumblingBlock.instances()) {
+				if (!(instance.handle() instanceof InstanceHandleImpl impl)) {
+					continue;
+				}
+                if (!(impl.instancer instanceof BatchedInstancer<?> instancer)) {
+					continue;
+                }
+
+				var bucket = out.computeIfAbsent(crumblingType, $ -> new CrumblingBucket());
+				bucket.instances.add(Pair.of(instance, instancer));
+
+				for (TransformCall<?> transformCall : instancer.getTransformCalls()) {
+					var mesh = transformCall.mesh;
+					bucket.vertexCount += mesh.getVertexCount();
+				}
+            }
+		}
+		return out;
+	}
+
+	public static <I extends Instance> int bufferOne(BatchedInstancer<I> batchedInstancer, int baseVertex, ReusableVertexList vertexList, DrawBuffer drawBuffer, Instance instance) {
+		int totalVertices = 0;
+
+		for (TransformCall<I> transformCall : batchedInstancer.getTransformCalls()) {
+            Mesh mesh = transformCall.mesh.mesh;
+
+			vertexList.ptr(drawBuffer.ptrForVertex(baseVertex + totalVertices));
+			vertexList.vertexCount(mesh.vertexCount());
+
+			mesh.write(vertexList);
+			batchedInstancer.type.getVertexTransformer()
+					.transform(vertexList, (I) instance);
+
+			totalVertices += mesh.vertexCount();
+		}
+
+		return totalVertices;
+	}
+
+	static final class CrumblingBucket {
+		private int vertexCount;
+		private final List<Pair<Instance, BatchedInstancer<?>>> instances = new ArrayList<>();
+
+	}
+}

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

@@ -18,7 +18,7 @@ import com.mojang.blaze3d.vertex.VertexFormat;
 import net.minecraft.client.renderer.RenderType;
 
 class BatchedDrawManager extends InstancerStorage<BatchedInstancer<?>> {
-	private final BatchedDrawTracker drawTracker = new BatchedDrawTracker();
+	public final BatchedDrawTracker drawTracker = new BatchedDrawTracker();
 	private final Map<RenderStage, BatchedStagePlan> stagePlans = new EnumMap<>(RenderStage.class);
 	private final Map<VertexFormat, BatchedMeshPool> meshPools = new HashMap<>();
 

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

@@ -7,6 +7,7 @@ import java.util.Set;
 
 import com.jozufozu.flywheel.api.event.RenderStage;
 import com.jozufozu.flywheel.extension.BufferBuilderExtension;
+import com.mojang.blaze3d.systems.RenderSystem;
 import com.mojang.blaze3d.vertex.BufferBuilder;
 
 public class BatchedDrawTracker {
@@ -25,16 +26,17 @@ public class BatchedDrawTracker {
 		((BufferBuilderExtension) scratch).flywheel$freeBuffer();
 	}
 
-	public void markActive(DrawBuffer buffer) {
+	public void markActive(RenderStage stage, DrawBuffer buffer) {
 		synchronized (activeBuffers) {
-			activeBuffers.get(buffer.getRenderStage())
+			activeBuffers.get(stage)
 					.add(buffer);
 		}
 	}
 
-	public void markInactive(DrawBuffer buffer) {
+	// TODO: remove?
+	public void markInactive(RenderStage stage, DrawBuffer buffer) {
 		synchronized (activeBuffers) {
-			activeBuffers.get(buffer.getRenderStage())
+			activeBuffers.get(stage)
 					.remove(buffer);
 		}
 	}
@@ -67,12 +69,12 @@ public class BatchedDrawTracker {
 		buffers.clear();
 	}
 
-	private void _draw(DrawBuffer buffer) {
+	public void _draw(DrawBuffer buffer) {
 		if (buffer.hasVertices()) {
 			BufferBuilderExtension scratch = (BufferBuilderExtension) this.scratch;
 			buffer.inject(scratch);
 			buffer.getRenderType()
-					.end(this.scratch, null);
+					.end(this.scratch, RenderSystem.getVertexSorting());
 		}
 		buffer.reset();
 	}

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

@@ -1,5 +1,6 @@
 package com.jozufozu.flywheel.backend.engine.batching;
 
+import java.util.ArrayList;
 import java.util.List;
 
 import com.jozufozu.flywheel.api.instance.Instance;
@@ -7,6 +8,8 @@ import com.jozufozu.flywheel.api.instance.InstanceType;
 import com.jozufozu.flywheel.backend.engine.AbstractInstancer;
 
 public class BatchedInstancer<I extends Instance> extends AbstractInstancer<I> {
+	private final List<TransformCall<I>> transformCalls = new ArrayList<>();
+
 	public BatchedInstancer(InstanceType<I> type) {
 		super(type);
 	}
@@ -26,4 +29,12 @@ public class BatchedInstancer<I extends Instance> extends AbstractInstancer<I> {
 	public void update() {
 		removeDeletedInstances();
 	}
+
+	public void addTransformCall(TransformCall<I> transformCall) {
+		transformCalls.add(transformCall);
+	}
+
+	public List<TransformCall<I>> getTransformCalls() {
+		return transformCalls;
+	}
 }

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

@@ -150,7 +150,7 @@ public class BatchedMeshPool {
 	}
 
 	public class BufferedMesh {
-		private final Mesh mesh;
+		public final Mesh mesh;
 		private final int byteSize;
 		private final int vertexCount;
 		private final Vector4fc boundingSphere;

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

@@ -84,7 +84,7 @@ public class BatchedStagePlan implements SimplyComposedPlan<BatchContext> {
 				return;
 			}
 
-			tracker.markActive(buffer);
+			tracker.markActive(stage, buffer);
 			buffer.prepare(vertexCount);
 
 			var vertexCounter = new AtomicInteger(0);

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

@@ -45,7 +45,10 @@ public class BatchingEngine extends AbstractEngine {
 
 	@Override
 	public void renderCrumblingInstances(TaskExecutor executor, RenderContext context, List<CrumblingBlock> crumblingBlocks) {
-		// TODO: implement
+		executor.syncUntil(flushFlag::isRaised);
+
+		var batchContext = BatchContext.create(context, this.renderOrigin);
+		BatchedCrumbling.renderCrumbling(crumblingBlocks, batchContext, this.drawManager.drawTracker);
 	}
 
 	@Override

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

@@ -0,0 +1,53 @@
+package com.jozufozu.flywheel.backend.engine.batching;
+
+import org.joml.Matrix3f;
+import org.joml.Matrix4f;
+import org.joml.Vector3f;
+import org.joml.Vector4f;
+
+import com.jozufozu.flywheel.api.vertex.MutableVertexList;
+import com.jozufozu.flywheel.api.vertex.ReusableVertexList;
+import com.jozufozu.flywheel.lib.vertex.VertexTransformations;
+import com.mojang.blaze3d.vertex.PoseStack;
+import com.mojang.blaze3d.vertex.SheetedDecalTextureGenerator;
+
+import net.minecraft.core.Direction;
+
+public class BatchingTransforms {
+	public static void applyMatrices(MutableVertexList vertexList, PoseStack.Pose matrices) {
+		Matrix4f modelMatrix = matrices.pose();
+		Matrix3f normalMatrix = matrices.normal();
+
+		for (int i = 0; i < vertexList.vertexCount(); i++) {
+			VertexTransformations.transformPos(vertexList, i, modelMatrix);
+			VertexTransformations.transformNormal(vertexList, i, normalMatrix);
+		}
+	}
+
+	/**
+	 * Performs the same operation as {@link SheetedDecalTextureGenerator} in-place.
+	 * <br>
+	 * Call this in world space.
+	 *
+	 * @param vertexList The vertex list to apply the transformations to.
+	 */
+	public static void applyDecalUVs(ReusableVertexList vertexList) {
+		Vector3f normal = new Vector3f();
+		Vector4f pos = new Vector4f();
+
+		for (int i = 0; i < vertexList.vertexCount(); i++) {
+			vertexList.getPos(i, pos);
+			vertexList.getNormal(i, normal);
+
+			Direction direction = Direction.getNearest(normal.x(), normal.y(), normal.z());
+			pos.rotateY((float)Math.PI);
+			pos.rotateX((-(float)Math.PI / 2F));
+			pos.rotate(direction.getRotation());
+			float u = -pos.x();
+			float v = -pos.y();
+
+			vertexList.u(i, u);
+			vertexList.v(i, v);
+		}
+	}
+}

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

@@ -14,19 +14,20 @@ import com.jozufozu.flywheel.lib.memory.MemoryBlock;
 import com.mojang.blaze3d.vertex.VertexFormat;
 
 import net.minecraft.client.renderer.RenderType;
+import net.minecraft.util.Mth;
 
 /**
  * A byte buffer that can be used to draw vertices through multiple {@link ReusableVertexList}s.
- *
- * The number of vertices needs to be known ahead of time.
+ * <br>
+ * Note: The number of vertices needs to be known ahead of time.
  */
 public class DrawBuffer {
 	private static final List<DrawBuffer> ALL = new ArrayList<>();
 
 	private final RenderType renderType;
-	private final RenderStage renderStage;
 	private final VertexFormat format;
 	private final int stride;
+	private final boolean sortOnUpload;
 	private final VertexListProvider provider;
 
 	private MemoryBlock data;
@@ -36,11 +37,11 @@ public class DrawBuffer {
 	private int vertexCount;
 	private int verticesToDraw;
 
-	public DrawBuffer(RenderType renderType, RenderStage renderStage, VertexFormat format, int stride, VertexListProvider provider) {
+	public DrawBuffer(RenderType renderType, VertexFormat format, int stride, boolean sortOnUpload, VertexListProvider provider) {
 		this.renderType = renderType;
-		this.renderStage = renderStage;
 		this.format = format;
 		this.stride = stride;
+		this.sortOnUpload = sortOnUpload;
 		this.provider = provider;
 
 		ALL.add(this);
@@ -64,6 +65,15 @@ public class DrawBuffer {
 		// is called and reallocates the buffer if there is not space for one more vertex.
 		int byteSize = stride * (vertexCount + 1);
 
+		// We'll need to allocate space for the index buffer if this render type needs sorting.
+		if (sortOnUpload) {
+			int i = renderType.mode()
+					.indexCount(vertexCount);
+			VertexFormat.IndexType indexType = VertexFormat.IndexType.least(i);
+			int extraBytes = Mth.roundToward(i * indexType.bytes, 4);
+			byteSize += extraBytes;
+		}
+
 		if (data == null) {
 			data = MemoryBlock.malloc(byteSize);
 			buffer = data.asBuffer();
@@ -116,10 +126,6 @@ public class DrawBuffer {
 		return renderType;
 	}
 
-	public RenderStage getRenderStage() {
-		return renderStage;
-	}
-
 	public VertexFormat getVertexFormat() {
 		return format;
 	}

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

@@ -13,18 +13,20 @@ import net.minecraft.client.renderer.RenderType;
 public class DrawBufferSet {
 	private final RenderType renderType;
 	private final VertexFormat format;
+	private final boolean sortOnUpload;
 	private final int stride;
 	private final VertexListProvider provider;
 	private final Map<RenderStage, DrawBuffer> buffers = new EnumMap<>(RenderStage.class);
 
-	public DrawBufferSet(RenderType renderType) {
+	public DrawBufferSet(RenderType renderType, boolean sortOnUpload) {
 		this.renderType = renderType;
+		this.sortOnUpload = sortOnUpload;
 		format = renderType.format();
 		stride = format.getVertexSize();
 		provider = VertexListProviderRegistry.getProvider(format);
 	}
 
 	public DrawBuffer getBuffer(RenderStage stage) {
-		return buffers.computeIfAbsent(stage, renderStage -> new DrawBuffer(renderType, renderStage, format, stride, provider));
+		return buffers.computeIfAbsent(stage, renderStage -> new DrawBuffer(renderType, format, stride, sortOnUpload, provider));
 	}
 }

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

@@ -3,8 +3,6 @@ package com.jozufozu.flywheel.backend.engine.batching;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import org.joml.FrustumIntersection;
-import org.joml.Matrix3f;
-import org.joml.Matrix4f;
 import org.joml.Vector4f;
 import org.joml.Vector4fc;
 
@@ -14,22 +12,25 @@ import com.jozufozu.flywheel.api.instance.InstanceVertexTransformer;
 import com.jozufozu.flywheel.api.material.Material;
 import com.jozufozu.flywheel.api.material.MaterialVertexTransformer;
 import com.jozufozu.flywheel.api.task.Plan;
-import com.jozufozu.flywheel.api.vertex.MutableVertexList;
 import com.jozufozu.flywheel.api.vertex.ReusableVertexList;
 import com.jozufozu.flywheel.lib.task.ForEachSlicePlan;
-import com.jozufozu.flywheel.lib.vertex.VertexTransformations;
 import com.mojang.blaze3d.vertex.PoseStack;
 
 import net.minecraft.client.multiplayer.ClientLevel;
 
 public class TransformCall<I extends Instance> {
-	private final BatchedInstancer<I> instancer;
-	private final int meshVertexCount;
+	public final BatchedInstancer<I> instancer;
+	public final Material material;
+	public final BatchedMeshPool.BufferedMesh mesh;
 
+	private final int meshVertexCount;
 	private final Plan<PlanContext> drawPlan;
 
 	public TransformCall(BatchedInstancer<I> instancer, Material material, BatchedMeshPool.BufferedMesh mesh) {
 		this.instancer = instancer;
+		this.material = material;
+		this.mesh = mesh;
+		instancer.addTransformCall(this);
 
 		InstanceVertexTransformer<I> instanceVertexTransformer = instancer.type.getVertexTransformer();
 		InstanceBoundingSphereTransformer<I> boundingSphereTransformer = instancer.type.getBoundingSphereTransformer();
@@ -56,7 +57,7 @@ public class TransformCall<I extends Instance> {
 				mesh.copyTo(vertexList.ptr());
 				instanceVertexTransformer.transform(vertexList, instance);
 				materialVertexTransformer.transform(vertexList, ctx.level);
-				applyMatrices(vertexList, ctx.matrices);
+				BatchingTransforms.applyMatrices(vertexList, ctx.matrices);
 			}
 		});
 	}
@@ -73,16 +74,6 @@ public class TransformCall<I extends Instance> {
 		return drawPlan;
 	}
 
-	private static void applyMatrices(MutableVertexList vertexList, PoseStack.Pose matrices) {
-		Matrix4f modelMatrix = matrices.pose();
-		Matrix3f normalMatrix = matrices.normal();
-
-		for (int i = 0; i < vertexList.vertexCount(); i++) {
-			VertexTransformations.transformPos(vertexList, i, modelMatrix);
-			VertexTransformations.transformNormal(vertexList, i, normalMatrix);
-		}
-	}
-
 	public record PlanContext(FrustumIntersection frustum, AtomicInteger vertexCounter, DrawBuffer buffer, ClientLevel level, PoseStack.Pose matrices) {
 	}
 }

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

@@ -37,6 +37,9 @@ public abstract class BufferBuilderMixin implements BufferBuilderExtension {
 	@Shadow
 	private boolean building;
 
+	@Shadow
+	private int nextElementByte;
+
 	@Override
 	public void flywheel$freeBuffer() {
 		if (buffer != null) {
@@ -58,5 +61,6 @@ public abstract class BufferBuilderMixin implements BufferBuilderExtension {
 
 		currentElement = format.getElements().get(0);
 		elementIndex = 0;
+		nextElementByte = vertexCount * format.getVertexSize();
 	}
 }

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

@@ -1,7 +1,9 @@
 package com.jozufozu.flywheel.mixin;
 
 import org.jetbrains.annotations.NotNull;
+import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
 import org.spongepowered.asm.mixin.Unique;
 
 import com.jozufozu.flywheel.backend.engine.batching.DrawBufferSet;
@@ -11,6 +13,9 @@ import net.minecraft.client.renderer.RenderType;
 
 @Mixin(RenderType.class)
 public class RenderTypeMixin implements RenderTypeExtension {
+	@Shadow
+	@Final
+	private boolean sortOnUpload;
 	@Unique
 	private DrawBufferSet flywheel$drawBufferSet;
 
@@ -18,7 +23,7 @@ public class RenderTypeMixin implements RenderTypeExtension {
 	@NotNull
 	public DrawBufferSet flywheel$getDrawBufferSet() {
 		if (flywheel$drawBufferSet == null) {
-			flywheel$drawBufferSet = new DrawBufferSet((RenderType) (Object) this);
+			flywheel$drawBufferSet = new DrawBufferSet((RenderType) (Object) this, sortOnUpload);
 		}
 		return flywheel$drawBufferSet;
 	}