Buckets of crumbs

- When crumbling, group DrawCalls by ShaderState.
- Pass all instances for a given crumbling progress at once.
- Use a map to associate initialized instancers with draw calls.
- Document most of the fields in InstancedDrawManager.
- Fix race condition in InstancedDrawManager#getInstancer.
- Will follow up for BatchingEngine and IndirectDrawManager.

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

@@ -44,6 +44,7 @@ public class BatchingEngine extends AbstractEngine implements SimplyComposedPlan
 	public <I extends Instance> Instancer<I> instancer(InstanceType<I> type, Model model, RenderStage stage) {
 		InstancerKey<I> key = new InstancerKey<>(type, model, stage);
 		BatchedInstancer<I> instancer = (BatchedInstancer<I>) instancers.get(key);
+		// FIXME: This needs to be synchronized like InstancingEngine
 		if (instancer == null) {
 			instancer = new BatchedInstancer<>(type);
 			instancers.put(key, instancer);

src/main/java/com/jozufozu/flywheel/backend/engine/indirect/IndirectDrawManager.java

@@ -23,6 +23,7 @@ public class IndirectDrawManager {
 	@SuppressWarnings("unchecked")
 	public <I extends Instance> Instancer<I> getInstancer(InstanceType<I> type, Model model, RenderStage stage) {
 		InstancerKey<I> key = new InstancerKey<>(type, model, stage);
+		// FIXME: This needs to be synchronized like InstancingEngine
 		IndirectInstancer<I> instancer = (IndirectInstancer<I>) instancers.get(key);
 		if (instancer == null) {
 			instancer = new IndirectInstancer<>(type);

src/main/java/com/jozufozu/flywheel/backend/engine/instancing/DrawCall.java

@@ -1,5 +1,6 @@
 package com.jozufozu.flywheel.backend.engine.instancing;
 
+import com.jozufozu.flywheel.backend.engine.InstanceHandleImpl;
 import com.jozufozu.flywheel.gl.array.GlVertexArray;
 
 public class DrawCall {
@@ -36,7 +37,7 @@ public class DrawCall {
 			return;
 		}
 
-		instancer.bindToVAO(vao, meshAttributes);
+		instancer.bindIfNeeded(vao, meshAttributes);
 		mesh.setup(vao);
 
 		vao.bindForDraw();
@@ -44,7 +45,7 @@ public class DrawCall {
 		mesh.draw(instanceCount);
 	}
 
-	public void renderOne(int index) {
+	public void renderOne(InstanceHandleImpl impl) {
 		if (isInvalid() || mesh.isEmpty()) {
 			return;
 		}
@@ -52,13 +53,13 @@ public class DrawCall {
 		instancer.update();
 
 		int instanceCount = instancer.getInstanceCount();
-		if (instanceCount <= 0 || index >= instanceCount) {
+		if (instanceCount <= 0 || impl.index >= instanceCount) {
 			return;
 		}
 
 		var vao = lazyScratchVao();
 
-		instancer.bindRaw(vao, meshAttributes, index);
+		instancer.bindRaw(vao, meshAttributes, impl.index);
 		mesh.setup(vao);
 
 		vao.bindForDraw();

src/main/java/com/jozufozu/flywheel/backend/engine/instancing/InstancedDrawManager.java

@@ -23,10 +23,40 @@ import com.jozufozu.flywheel.api.vertex.VertexType;
 import com.jozufozu.flywheel.backend.engine.InstancerKey;
 
 public class InstancedDrawManager {
+	/**
+	 * A map of instancer keys to instancers.
+	 * <br>
+	 * This map is populated as instancers are requested and contains both initialized and uninitialized instancers.
+	 * Write access to this map must be synchronized on {@link #creationLock}.
+	 * <br>
+	 * See {@link #getInstancer} for insertion details.
+	 */
 	private final Map<InstancerKey<?>, InstancedInstancer<?>> instancers = new HashMap<>();
+	/**
+	 * A list of instancers that have not yet been initialized.
+	 * <br>
+	 * All new instancers land here before having resources allocated in {@link #flush}.
+	 * Write access to this list must be synchronized on {@link #creationLock}.
+	 */
 	private final List<UninitializedInstancer> uninitializedInstancers = new ArrayList<>();
-	private final List<InitializedInstancer> initializedInstancers = new ArrayList<>();
+	/**
+	 * Mutex for {@link #instancers} and {@link #uninitializedInstancers}.
+	 */
+	private final Object creationLock = new Object();
+
+	/**
+	 * A map of initialized instancers to their draw calls.
+	 * <br>
+	 * This map is populated in {@link #flush} and contains only initialized instancers.
+	 */
+	private final Map<InstancedInstancer<?>, List<DrawCall>> initializedInstancers = new HashMap<>();
+	/**
+	 * The set of draw calls to make in each {@link RenderStage}.
+	 */
 	private final Map<RenderStage, DrawSet> drawSets = new EnumMap<>(RenderStage.class);
+	/**
+	 * A map of vertex types to their mesh pools.
+	 */
 	private final Map<VertexType, InstancedMeshPool> meshPools = new HashMap<>();
 	private final EBOCache eboCache = new EBOCache();
 
@@ -37,13 +67,27 @@ public class InstancedDrawManager {
 	@SuppressWarnings("unchecked")
 	public <I extends Instance> Instancer<I> getInstancer(InstanceType<I> type, Model model, RenderStage stage) {
 		InstancerKey<I> key = new InstancerKey<>(type, model, stage);
+
 		InstancedInstancer<I> instancer = (InstancedInstancer<I>) instancers.get(key);
-		if (instancer == null) {
-			instancer = new InstancedInstancer<>(type);
-			instancers.put(key, instancer);
-			uninitializedInstancers.add(new UninitializedInstancer(instancer, model, stage));
+		// Happy path: instancer is already initialized.
+        if (instancer != null) {
+            return instancer;
+        }
+
+		// Unhappy path: instancer is not initialized, need to sync to make sure we don't create duplicates.
+        synchronized (creationLock) {
+			// Someone else might have initialized it while we were waiting for the lock.
+            instancer = (InstancedInstancer<I>) instancers.get(key);
+            if (instancer != null) {
+                return instancer;
+            }
+
+			// Create a new instancer and add it to the uninitialized list.
+            instancer = new InstancedInstancer<>(type);
+            instancers.put(key, instancer);
+            uninitializedInstancers.add(new UninitializedInstancer(instancer, model, stage));
+			return instancer;
 		}
-		return instancer;
 	}
 
 	public void flush() {
@@ -68,14 +112,16 @@ public class InstancedDrawManager {
 				.forEach(DrawSet::delete);
 		drawSets.clear();
 
-		initializedInstancers.forEach(InitializedInstancer::deleteInstancer);
+		initializedInstancers.keySet()
+				.forEach(InstancedInstancer::delete);
 		initializedInstancers.clear();
 
 		eboCache.invalidate();
 	}
 
 	public void clearInstancers() {
-		initializedInstancers.forEach(InitializedInstancer::clear);
+		initializedInstancers.keySet()
+				.forEach(InstancedInstancer::clear);
 	}
 
 	private void add(InstancedInstancer<?> instancer, Model model, RenderStage stage) {
@@ -94,7 +140,7 @@ public class InstancedDrawManager {
 			drawSet.put(shaderState, drawCall);
 			drawCalls.add(drawCall);
 		}
-		initializedInstancers.add(new InitializedInstancer(instancer, drawCalls));
+		initializedInstancers.put(instancer, drawCalls);
 	}
 
 	private InstancedMeshPool.BufferedMesh alloc(Mesh mesh) {
@@ -103,13 +149,7 @@ public class InstancedDrawManager {
 	}
 
 	public List<DrawCall> drawCallsForInstancer(InstancedInstancer<?> instancer) {
-		for (InitializedInstancer initializedInstancer : initializedInstancers) {
-			if (initializedInstancer.instancer == instancer) {
-				return initializedInstancer.drawCalls;
-			}
-		}
-
-		return List.of();
+		return initializedInstancers.getOrDefault(instancer, List.of());
 	}
 
 	public static class DrawSet implements Iterable<Map.Entry<ShaderState, Collection<DrawCall>>> {
@@ -150,14 +190,4 @@ public class InstancedDrawManager {
 
 	private record UninitializedInstancer(InstancedInstancer<?> instancer, Model model, RenderStage stage) {
 	}
-
-	private record InitializedInstancer(InstancedInstancer<?> instancer, List<DrawCall> drawCalls) {
-		public void deleteInstancer() {
-			instancer.delete();
-		}
-
-		public void clear() {
-			instancer.clear();
-		}
-	}
 }

src/main/java/com/jozufozu/flywheel/backend/engine/instancing/InstancedInstancer.java

@@ -79,7 +79,13 @@ public class InstancedInstancer<I extends Instance> extends AbstractInstancer<I>
 		}
 	}
 
-	public void bindToVAO(GlVertexArray vao, int startAttrib) {
+	/**
+	 * Bind this instancer's vbo to the given vao if it hasn't already been bound.
+	 * @param vao The vao to bind to.
+	 * @param startAttrib The first attribute to bind. This method will bind attributes in the half open range
+	 * 		{@code [startAttrib, startAttrib + instanceFormat.getAttributeCount())}.
+	 */
+	public void bindIfNeeded(GlVertexArray vao, int startAttrib) {
 		if (!boundTo.add(vao)) {
 			return;
 		}
@@ -87,6 +93,13 @@ public class InstancedInstancer<I extends Instance> extends AbstractInstancer<I>
 		bindRaw(vao, startAttrib, 0);
 	}
 
+	/**
+	 * Bind this instancer's vbo to the given vao with the given base instance to calculate the binding offset.
+	 * @param vao The vao to bind to.
+	 * @param startAttrib The first attribute to bind. This method will bind attributes in the half open range
+	 * 		{@code [startAttrib, startAttrib + instanceFormat.getAttributeCount())}.
+	 * @param baseInstance The base instance to calculate the binding offset from.
+	 */
 	public void bindRaw(GlVertexArray vao, int startAttrib, int baseInstance) {
 		long offset = (long) baseInstance * instanceStride;
 		vao.bindVertexBuffer(1, vbo.handle(), offset, instanceStride);

src/main/java/com/jozufozu/flywheel/backend/engine/instancing/InstancingEngine.java

@@ -1,7 +1,11 @@
 package com.jozufozu.flywheel.backend.engine.instancing;
 
+import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
+import org.jetbrains.annotations.NotNull;
 import org.lwjgl.opengl.GL32;
 
 import com.jozufozu.flywheel.api.context.Context;
@@ -83,40 +87,70 @@ public class InstancingEngine extends AbstractEngine {
 
 	@Override
 	public void renderCrumblingInstances(TaskExecutor executor, RenderContext context, List<Instance> instances, int progress) {
-		// TODO: optimize
+		if (instances.isEmpty()) {
+			return;
+		}
 
+		if (progress < 0 || progress >= ModelBakery.DESTROY_TYPES.size()) {
+			return;
+		}
+
+		// Need to wait for flush before we can inspect instancer state.
 		executor.syncUntil(flushFlag::isRaised);
 
-		var type = ModelBakery.DESTROY_TYPES.get(progress);
+		// Sort draw calls into buckets, so we don't have to do as many shader binds.
+		var drawMap = getDrawsForInstances(instances);
+
+		if (drawMap.isEmpty()) {
+			return;
+		}
 
 		try (var state = GlStateTracker.getRestoreState()) {
-			type.setupRenderState();
+			ModelBakery.DESTROY_TYPES.get(progress)
+					.setupRenderState();
 
-			Textures.bindActiveTextures();
+			for (var entry : drawMap.entrySet()) {
+				setup(entry.getKey(), Contexts.CRUMBLING);
 
-			for (Instance instance : instances) {
-				if (!(instance.handle() instanceof InstanceHandleImpl impl)) {
-					continue;
-				}
-				if (!(impl.instancer instanceof InstancedInstancer<?> instancer)) {
-					continue;
-				}
-
-				List<DrawCall> draws = drawManager.drawCallsForInstancer(instancer);
-
-				draws.removeIf(DrawCall::isInvalid);
-
-				for (DrawCall draw : draws) {
-					var shader = draw.shaderState;
-
-					setup(shader, Contexts.CRUMBLING);
-
-					draw.renderOne(impl.index);
+				for (Runnable draw : entry.getValue()) {
+					draw.run();
 				}
 			}
 		}
 	}
 
+	/**
+	 * Get all draw calls for the given instances, grouped by shader state.
+	 * @param instances The instances to draw.
+	 * @return A mapping of shader states to many runnable draw calls.
+	 */
+	@NotNull
+	private Map<ShaderState, List<Runnable>> getDrawsForInstances(List<Instance> instances) {
+		Map<ShaderState, List<Runnable>> out = new HashMap<>();
+
+		for (Instance instance : instances) {
+			// Filter out instances that weren't created by this engine.
+			// If all is well, we probably shouldn't take the `continue`
+			// branches but better to do checked casts.
+			if (!(instance.handle() instanceof InstanceHandleImpl impl)) {
+				continue;
+			}
+			if (!(impl.instancer instanceof InstancedInstancer<?> instancer)) {
+				continue;
+			}
+
+			List<DrawCall> draws = drawManager.drawCallsForInstancer(instancer);
+
+			draws.removeIf(DrawCall::isInvalid);
+
+			for (DrawCall draw : draws) {
+				out.computeIfAbsent(draw.shaderState, $ -> new ArrayList<>())
+						.add(() -> draw.renderOne(impl));
+			}
+		}
+		return out;
+	}
+
 	private void setup() {
 		GlTextureUnit.T2.makeActive();
 		Minecraft.getInstance().gameRenderer.lightTexture().turnOnLightLayer();

src/main/java/com/jozufozu/flywheel/impl/visualization/VisualizationManagerImpl.java

@@ -1,5 +1,7 @@
 package com.jozufozu.flywheel.impl.visualization;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.SortedSet;
 
 import org.jetbrains.annotations.Nullable;
@@ -8,6 +10,7 @@ import com.jozufozu.flywheel.api.backend.BackendManager;
 import com.jozufozu.flywheel.api.backend.Engine;
 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.task.Plan;
 import com.jozufozu.flywheel.api.task.TaskExecutor;
 import com.jozufozu.flywheel.api.visual.DynamicVisual;
@@ -28,6 +31,8 @@ import com.jozufozu.flywheel.lib.task.NamedFlag;
 import com.jozufozu.flywheel.lib.task.RaisePlan;
 import com.jozufozu.flywheel.lib.util.LevelAttached;
 
+import it.unimi.dsi.fastutil.ints.Int2ObjectArrayMap;
+import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
 import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
 import net.minecraft.client.Minecraft;
 import net.minecraft.client.multiplayer.ClientLevel;
@@ -205,6 +210,8 @@ public class VisualizationManagerImpl implements VisualizationManager {
 	public void renderCrumbling(RenderContext context, Long2ObjectMap<SortedSet<BlockDestructionProgress>> destructionProgress) {
 		taskExecutor.syncUntil(frameVisualsFlag::isRaised);
 
+		Int2ObjectMap<List<Instance>> progress2instances = new Int2ObjectArrayMap<>();
+
 		for (var entry : destructionProgress.long2ObjectEntrySet()) {
 			var set = entry.getValue();
 			if (set == null || set.isEmpty()) {
@@ -231,7 +238,12 @@ public class VisualizationManagerImpl implements VisualizationManager {
 			int progress = set.last()
 					.getProgress();
 
-			engine.renderCrumblingInstances(taskExecutor, context, instances, progress);
+			progress2instances.computeIfAbsent(progress, $ -> new ArrayList<>())
+					.addAll(instances);
+		}
+
+		for (var entry : progress2instances.int2ObjectEntrySet()) {
+			engine.renderCrumblingInstances(taskExecutor, context, entry.getValue(), entry.getIntKey());
 		}
 	}