High noon

- Draw! meshes in contractual order on instancing engine
- Rename DrawSet -> InstancedRenderStage and consolidate associated code
  inside the class
- Re-use GroupKey between engines
- Sort instanced draws based on material and mesh index in model
- Make doCrumblingSort abstract on the instancer class and re-use it
  between engines

src/main/java/com/jozufozu/flywheel/backend/engine/DrawManager.java

@@ -1,6 +1,7 @@
 package com.jozufozu.flywheel.backend.engine;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
@@ -13,6 +14,11 @@ import com.jozufozu.flywheel.api.instance.InstanceType;
 import com.jozufozu.flywheel.api.instance.Instancer;
 import com.jozufozu.flywheel.api.model.Model;
 import com.jozufozu.flywheel.backend.engine.embed.Environment;
+import com.jozufozu.flywheel.lib.util.Pair;
+
+import it.unimi.dsi.fastutil.ints.Int2ObjectArrayMap;
+import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
+import net.minecraft.client.resources.model.ModelBakery;
 
 public abstract class DrawManager<N extends AbstractInstancer<?>> {
 	/**
@@ -93,4 +99,36 @@ public abstract class DrawManager<N extends AbstractInstancer<?>> {
 
 		return false;
 	}
+
+	protected static <I extends AbstractInstancer<?>> Map<GroupKey<?>, Int2ObjectMap<List<Pair<I, InstanceHandleImpl>>>> doCrumblingSort(Class<I> clazz, List<Engine.CrumblingBlock> crumblingBlocks) {
+		Map<GroupKey<?>, Int2ObjectMap<List<Pair<I, InstanceHandleImpl>>>> byType = new HashMap<>();
+		for (Engine.CrumblingBlock block : crumblingBlocks) {
+			int progress = block.progress();
+
+			if (progress < 0 || progress >= ModelBakery.DESTROY_TYPES.size()) {
+				continue;
+			}
+
+			for (Instance instance : block.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;
+				}
+
+				AbstractInstancer<?> abstractInstancer = impl.instancer;
+				if (!clazz.isInstance(abstractInstancer)) {
+					continue;
+				}
+
+				var instancer = clazz.cast(abstractInstancer);
+
+				byType.computeIfAbsent(new GroupKey<>(instancer.type, instancer.environment), $ -> new Int2ObjectArrayMap<>())
+						.computeIfAbsent(progress, $ -> new ArrayList<>())
+						.add(Pair.of(instancer, impl));
+			}
+		}
+		return byType;
+	}
 }

src/main/java/com/jozufozu/flywheel/backend/engine/GroupKey.java

@@ -0,0 +1,8 @@
+package com.jozufozu.flywheel.backend.engine;
+
+import com.jozufozu.flywheel.api.instance.Instance;
+import com.jozufozu.flywheel.api.instance.InstanceType;
+import com.jozufozu.flywheel.backend.engine.embed.Environment;
+
+public record GroupKey<I extends Instance>(InstanceType<I> instanceType, Environment environment) {
+}

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

@@ -42,7 +42,7 @@ public class IndirectCullingGroup<I extends Instance> {
 	private final Environment environment;
 	private final long instanceStride;
 	private final IndirectBuffers buffers;
-	private final List<IndirectInstancer<?>> instancers = new ArrayList<>();
+	private final List<IndirectInstancer<I>> instancers = new ArrayList<>();
 	private final List<IndirectDraw> indirectDraws = new ArrayList<>();
 	private final Map<RenderStage, List<MultiDraw>> multiDraws = new EnumMap<>(RenderStage.class);
 
@@ -216,7 +216,7 @@ public class IndirectCullingGroup<I extends Instance> {
 		}
 	}
 
-	public GlProgram bindWithContextShader(ContextShader override) {
+	public void bindWithContextShader(ContextShader override) {
 		var program = programs.getIndirectProgram(instanceType, override);
 
 		program.bind();
@@ -227,8 +227,6 @@ public class IndirectCullingGroup<I extends Instance> {
 
 		var flwBaseDraw = drawProgram.getUniformLocation("_flw_baseDraw");
 		glUniform1ui(flwBaseDraw, 0);
-
-		return program;
 	}
 
 	private void drawBarrier() {

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

@@ -6,7 +6,6 @@ import static org.lwjgl.opengl.GL30.glBindBufferRange;
 import static org.lwjgl.opengl.GL40.glDrawElementsIndirect;
 import static org.lwjgl.opengl.GL43.GL_SHADER_STORAGE_BUFFER;
 
-import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -14,18 +13,16 @@ import java.util.Map;
 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.instance.InstanceType;
 import com.jozufozu.flywheel.backend.Samplers;
 import com.jozufozu.flywheel.backend.compile.ContextShader;
 import com.jozufozu.flywheel.backend.compile.IndirectPrograms;
 import com.jozufozu.flywheel.backend.engine.CommonCrumbling;
 import com.jozufozu.flywheel.backend.engine.DrawManager;
-import com.jozufozu.flywheel.backend.engine.InstanceHandleImpl;
+import com.jozufozu.flywheel.backend.engine.GroupKey;
 import com.jozufozu.flywheel.backend.engine.InstancerKey;
 import com.jozufozu.flywheel.backend.engine.MaterialRenderState;
 import com.jozufozu.flywheel.backend.engine.MeshPool;
 import com.jozufozu.flywheel.backend.engine.TextureBinder;
-import com.jozufozu.flywheel.backend.engine.embed.Environment;
 import com.jozufozu.flywheel.backend.engine.uniform.Uniforms;
 import com.jozufozu.flywheel.backend.gl.GlStateTracker;
 import com.jozufozu.flywheel.backend.gl.array.GlVertexArray;
@@ -33,10 +30,7 @@ import com.jozufozu.flywheel.backend.gl.buffer.GlBuffer;
 import com.jozufozu.flywheel.backend.gl.buffer.GlBufferType;
 import com.jozufozu.flywheel.lib.material.SimpleMaterial;
 import com.jozufozu.flywheel.lib.memory.MemoryBlock;
-import com.jozufozu.flywheel.lib.util.Pair;
 
-import it.unimi.dsi.fastutil.ints.Int2ObjectArrayMap;
-import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
 import net.minecraft.client.resources.model.ModelBakery;
 
 public class IndirectDrawManager extends DrawManager<IndirectInstancer<?>> {
@@ -68,7 +62,7 @@ public class IndirectDrawManager extends DrawManager<IndirectInstancer<?>> {
 	@Override
 	protected <I extends Instance> void initialize(InstancerKey<I> key, IndirectInstancer<?> instancer) {
 		var groupKey = new GroupKey<>(key.type(), key.environment());
-		var group = (IndirectCullingGroup<I>) cullingGroups.computeIfAbsent(groupKey, t -> new IndirectCullingGroup<>(t.type, t.environment, programs));
+		var group = (IndirectCullingGroup<I>) cullingGroups.computeIfAbsent(groupKey, t -> new IndirectCullingGroup<>(t.instanceType(), t.environment(), programs));
 		group.add((IndirectInstancer<I>) instancer, key.model(), key.stage(), meshPool);
 	}
 
@@ -148,7 +142,7 @@ public class IndirectDrawManager extends DrawManager<IndirectInstancer<?>> {
 	}
 
 	public void renderCrumbling(List<Engine.CrumblingBlock> crumblingBlocks) {
-		var byType = doCrumblingSort(crumblingBlocks);
+		var byType = doCrumblingSort(IndirectInstancer.class, crumblingBlocks);
 
 		if (byType.isEmpty()) {
 			return;
@@ -168,15 +162,13 @@ public class IndirectDrawManager extends DrawManager<IndirectInstancer<?>> {
 			GlBufferType.DRAW_INDIRECT_BUFFER.bind(crumblingDrawBuffer.handle());
 			glBindBufferRange(GL_SHADER_STORAGE_BUFFER, IndirectBuffers.DRAW_INDEX, crumblingDrawBuffer.handle(), 0, IndirectBuffers.DRAW_COMMAND_STRIDE);
 
-			for (var instanceTypeEntry : byType.entrySet()) {
-				var byProgress = instanceTypeEntry.getValue();
+			for (var groupEntry : byType.entrySet()) {
+				var byProgress = groupEntry.getValue();
 
 				// Set up the crumbling program buffers. Nothing changes here between draws.
-				var program = cullingGroups.get(instanceTypeEntry.getKey())
+				cullingGroups.get(groupEntry.getKey())
 						.bindWithContextShader(ContextShader.CRUMBLING);
 
-				program.setSamplerBinding("crumblingTex", Samplers.CRUMBLING);
-
 				for (var progressEntry : byProgress.int2ObjectEntrySet()) {
 					Samplers.CRUMBLING.makeActive();
 					TextureBinder.bind(ModelBakery.BREAKING_LOCATIONS.get(progressEntry.getIntKey()));
@@ -186,7 +178,6 @@ public class IndirectDrawManager extends DrawManager<IndirectInstancer<?>> {
 						int instanceIndex = instanceHandlePair.second().index;
 
 						for (IndirectDraw draw : instancer.draws()) {
-
 							// Transform the material to be suited for crumbling.
 							CommonCrumbling.applyCrumblingProperties(crumblingMaterial, draw.material());
 
@@ -210,35 +201,4 @@ public class IndirectDrawManager extends DrawManager<IndirectInstancer<?>> {
 			block.free();
 		}
 	}
-
-	private static Map<GroupKey<?>, Int2ObjectMap<List<Pair<IndirectInstancer<?>, InstanceHandleImpl>>>> doCrumblingSort(List<Engine.CrumblingBlock> crumblingBlocks) {
-		Map<GroupKey<?>, Int2ObjectMap<List<Pair<IndirectInstancer<?>, InstanceHandleImpl>>>> byType = new HashMap<>();
-		for (Engine.CrumblingBlock block : crumblingBlocks) {
-			int progress = block.progress();
-
-			if (progress < 0 || progress >= ModelBakery.DESTROY_TYPES.size()) {
-				continue;
-			}
-
-			for (Instance instance : block.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 IndirectInstancer<?> instancer)) {
-					continue;
-				}
-
-				byType.computeIfAbsent(new GroupKey<>(instancer.type, instancer.environment), $ -> new Int2ObjectArrayMap<>())
-						.computeIfAbsent(progress, $ -> new ArrayList<>())
-						.add(Pair.of(instancer, impl));
-			}
-		}
-		return byType;
-	}
-
-	public record GroupKey<I extends Instance>(InstanceType<I> type, Environment environment) {
-	}
 }

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

@@ -1,24 +1,38 @@
 package com.jozufozu.flywheel.backend.engine.instancing;
 
+import com.jozufozu.flywheel.api.material.Material;
+import com.jozufozu.flywheel.backend.engine.GroupKey;
 import com.jozufozu.flywheel.backend.engine.InstanceHandleImpl;
 import com.jozufozu.flywheel.backend.engine.MeshPool;
 import com.jozufozu.flywheel.backend.gl.TextureBuffer;
 
-public class DrawCall {
-	public final ShaderState shaderState;
+public class InstancedDraw {
+	public final GroupKey<?> groupKey;
 	private final InstancedInstancer<?> instancer;
 	private final MeshPool.PooledMesh mesh;
+	private final Material material;
+	private final int indexOfMeshInModel;
 
 	private boolean deleted;
 
-	public DrawCall(InstancedInstancer<?> instancer, MeshPool.PooledMesh mesh, ShaderState shaderState) {
+	public InstancedDraw(InstancedInstancer<?> instancer, MeshPool.PooledMesh mesh, GroupKey<?> groupKey, Material material, int indexOfMeshInModel) {
 		this.instancer = instancer;
 		this.mesh = mesh;
-		this.shaderState = shaderState;
+		this.groupKey = groupKey;
+		this.material = material;
+		this.indexOfMeshInModel = indexOfMeshInModel;
 
 		mesh.acquire();
 	}
 
+	public int indexOfMeshInModel() {
+		return indexOfMeshInModel;
+	}
+
+	public Material material() {
+		return material;
+	}
+
 	public boolean deleted() {
 		return deleted;
 	}

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

@@ -1,19 +1,11 @@
 package com.jozufozu.flywheel.backend.engine.instancing;
 
-import java.util.ArrayList;
-import java.util.Collection;
 import java.util.EnumMap;
-import java.util.HashMap;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
-import java.util.function.Consumer;
 
 import org.lwjgl.opengl.GL32;
 
-import com.google.common.collect.ArrayListMultimap;
-import com.google.common.collect.ImmutableListMultimap;
-import com.google.common.collect.ListMultimap;
 import com.jozufozu.flywheel.api.backend.Engine;
 import com.jozufozu.flywheel.api.event.RenderStage;
 import com.jozufozu.flywheel.api.instance.Instance;
@@ -24,7 +16,7 @@ import com.jozufozu.flywheel.backend.compile.ContextShader;
 import com.jozufozu.flywheel.backend.compile.InstancingPrograms;
 import com.jozufozu.flywheel.backend.engine.CommonCrumbling;
 import com.jozufozu.flywheel.backend.engine.DrawManager;
-import com.jozufozu.flywheel.backend.engine.InstanceHandleImpl;
+import com.jozufozu.flywheel.backend.engine.GroupKey;
 import com.jozufozu.flywheel.backend.engine.InstancerKey;
 import com.jozufozu.flywheel.backend.engine.MaterialEncoder;
 import com.jozufozu.flywheel.backend.engine.MaterialRenderState;
@@ -37,15 +29,13 @@ import com.jozufozu.flywheel.backend.gl.array.GlVertexArray;
 import com.jozufozu.flywheel.backend.gl.shader.GlProgram;
 import com.jozufozu.flywheel.lib.material.SimpleMaterial;
 
-import it.unimi.dsi.fastutil.ints.Int2ObjectArrayMap;
-import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
 import net.minecraft.client.resources.model.ModelBakery;
 
 public class InstancedDrawManager extends DrawManager<InstancedInstancer<?>> {
 	/**
 	 * The set of draw calls to make in each {@link RenderStage}.
 	 */
-	private final Map<RenderStage, DrawSet> drawSets = new EnumMap<>(RenderStage.class);
+	private final Map<RenderStage, InstancedRenderStage> stages = new EnumMap<>(RenderStage.class);
 	private final InstancingPrograms programs;
 	/**
 	 * A map of vertex types to their mesh pools.
@@ -65,6 +55,7 @@ public class InstancedDrawManager extends DrawManager<InstancedInstancer<?>> {
 		meshPool.bind(vao);
 	}
 
+	@Override
 	public void flush() {
 		super.flush();
 
@@ -81,33 +72,42 @@ public class InstancedDrawManager extends DrawManager<InstancedInstancer<?>> {
 			}
 		});
 
-		for (DrawSet drawSet : drawSets.values()) {
+		for (InstancedRenderStage instancedRenderStage : stages.values()) {
 			// Remove the draw calls for any instancers we deleted.
-			drawSet.prune();
+			instancedRenderStage.flush();
 		}
 
 		meshPool.flush();
 	}
 
+	@Override
 	public void renderStage(RenderStage stage) {
-		var drawSet = drawSets.getOrDefault(stage, DrawSet.EMPTY);
+		var drawSet = stages.get(stage);
 
-		if (drawSet.isEmpty()) {
+		if (drawSet == null || drawSet.isEmpty()) {
 			return;
 		}
 
 		try (var state = GlStateTracker.getRestoreState()) {
-			render(drawSet);
+			Uniforms.bindForDraw();
+			vao.bindForDraw();
+			TextureBinder.bindLightAndOverlay();
+
+			drawSet.draw(instanceTexture, programs);
+
+			MaterialRenderState.reset();
+			TextureBinder.resetLightAndOverlay();
 		}
 	}
 
+	@Override
 	public void delete() {
 		instancers.values()
 				.forEach(InstancedInstancer::delete);
 
-		drawSets.values()
-				.forEach(DrawSet::delete);
-		drawSets.clear();
+		stages.values()
+				.forEach(InstancedRenderStage::delete);
+		stages.clear();
 
 		meshPool.delete();
 		instanceTexture.delete();
@@ -117,42 +117,6 @@ public class InstancedDrawManager extends DrawManager<InstancedInstancer<?>> {
 		super.delete();
 	}
 
-	private void render(InstancedDrawManager.DrawSet drawSet) {
-		Uniforms.bindForDraw();
-		vao.bindForDraw();
-		TextureBinder.bindLightAndOverlay();
-
-		for (var entry : drawSet) {
-			var shader = entry.getKey();
-			var drawCalls = entry.getValue();
-
-			if (drawCalls.isEmpty()) {
-				continue;
-			}
-
-			var environment = shader.environment();
-			var material = shader.material();
-
-			var program = programs.get(shader.instanceType(), environment.contextShader());
-			program.bind();
-
-			environment.setupDraw(program);
-
-			uploadMaterialUniform(program, material);
-
-			MaterialRenderState.setup(material);
-
-			Samplers.INSTANCE_BUFFER.makeActive();
-
-			for (var drawCall : drawCalls) {
-				drawCall.render(instanceTexture);
-			}
-		}
-
-		MaterialRenderState.reset();
-		TextureBinder.resetLightAndOverlay();
-	}
-
 	@Override
 	protected <I extends Instance> InstancedInstancer<I> create(InstancerKey<I> key) {
 		return new InstancedInstancer<>(key.type(), key.environment());
@@ -162,26 +126,28 @@ public class InstancedDrawManager extends DrawManager<InstancedInstancer<?>> {
 	protected <I extends Instance> void initialize(InstancerKey<I> key, InstancedInstancer<?> instancer) {
 		instancer.init();
 
-		DrawSet drawSet = drawSets.computeIfAbsent(key.stage(), DrawSet::new);
+		InstancedRenderStage instancedRenderStage = stages.computeIfAbsent(key.stage(), $ -> new InstancedRenderStage());
 
 		var meshes = key.model()
 				.meshes();
-		for (var entry : meshes) {
+		for (int i = 0; i < meshes.size(); i++) {
+			var entry = meshes.get(i);
 			var mesh = meshPool.alloc(entry.mesh());
 
-			ShaderState shaderState = new ShaderState(entry.material(), key.type(), key.environment());
-			DrawCall drawCall = new DrawCall(instancer, mesh, shaderState);
+			GroupKey<?> groupKey = new GroupKey<>(key.type(), key.environment());
+			InstancedDraw instancedDraw = new InstancedDraw(instancer, mesh, groupKey, entry.material(), i);
 
-			drawSet.put(shaderState, drawCall);
-			instancer.addDrawCall(drawCall);
+			instancedRenderStage.put(groupKey, instancedDraw);
+			instancer.addDrawCall(instancedDraw);
 		}
 	}
 
+	@Override
 	public void renderCrumbling(List<Engine.CrumblingBlock> crumblingBlocks) {
 		// Sort draw calls into buckets, so we don't have to do as many shader binds.
-		var byShaderState = doCrumblingSort(crumblingBlocks);
+		var byType = doCrumblingSort(InstancedInstancer.class, crumblingBlocks);
 
-		if (byShaderState.isEmpty()) {
+		if (byType.isEmpty()) {
 			return;
 		}
 
@@ -192,38 +158,32 @@ public class InstancedDrawManager extends DrawManager<InstancedInstancer<?>> {
 			vao.bindForDraw();
 			TextureBinder.bindLightAndOverlay();
 
-			for (var shaderStateEntry : byShaderState.entrySet()) {
-				var byProgress = shaderStateEntry.getValue();
+			for (var groupEntry : byType.entrySet()) {
+				var byProgress = groupEntry.getValue();
 
-				if (byProgress.isEmpty()) {
-					continue;
-				}
-
-				ShaderState shader = shaderStateEntry.getKey();
-
-				CommonCrumbling.applyCrumblingProperties(crumblingMaterial, shader.material());
+				GroupKey<?> shader = groupEntry.getKey();
 
 				var program = programs.get(shader.instanceType(), ContextShader.CRUMBLING);
 				program.bind();
 
-				uploadMaterialUniform(program, crumblingMaterial);
-
-				MaterialRenderState.setup(crumblingMaterial);
-
 				for (var progressEntry : byProgress.int2ObjectEntrySet()) {
-					var drawCalls = progressEntry.getValue();
-
-					if (drawCalls.isEmpty()) {
-						continue;
-					}
-
 					Samplers.CRUMBLING.makeActive();
 					TextureBinder.bind(ModelBakery.BREAKING_LOCATIONS.get(progressEntry.getIntKey()));
 
-					Samplers.INSTANCE_BUFFER.makeActive();
+					for (var instanceHandlePair : progressEntry.getValue()) {
+						InstancedInstancer<?> instancer = instanceHandlePair.first();
+						var handle = instanceHandlePair.second();
 
-					for (Consumer<TextureBuffer> drawCall : drawCalls) {
-						drawCall.accept(instanceTexture);
+						for (InstancedDraw draw : instancer.draws()) {
+							CommonCrumbling.applyCrumblingProperties(crumblingMaterial, draw.material());
+							uploadMaterialUniform(program, crumblingMaterial);
+
+							MaterialRenderState.setup(crumblingMaterial);
+
+							Samplers.INSTANCE_BUFFER.makeActive();
+
+							draw.renderOne(instanceTexture, handle);
+						}
 					}
 				}
 			}
@@ -233,38 +193,6 @@ public class InstancedDrawManager extends DrawManager<InstancedInstancer<?>> {
 		}
 	}
 
-	private static Map<ShaderState, Int2ObjectMap<List<Consumer<TextureBuffer>>>> doCrumblingSort(List<Engine.CrumblingBlock> instances) {
-		Map<ShaderState, Int2ObjectMap<List<Consumer<TextureBuffer>>>> out = new HashMap<>();
-
-		for (Engine.CrumblingBlock triple : instances) {
-			int progress = triple.progress();
-
-			if (progress < 0 || progress >= ModelBakery.DESTROY_TYPES.size()) {
-				continue;
-			}
-
-			for (Instance instance : triple.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;
-				}
-
-				for (DrawCall draw : instancer.drawCalls()) {
-					out.computeIfAbsent(draw.shaderState, $ -> new Int2ObjectArrayMap<>())
-							.computeIfAbsent(progress, $ -> new ArrayList<>())
-							.add(buf -> draw.renderOne(buf, impl));
-				}
-			}
-		}
-
-		return out;
-	}
-
 	public static void uploadMaterialUniform(GlProgram program, Material material) {
 		int uniformLocation = program.getUniformLocation("_flw_packedMaterial");
 		int vertexIndex = ShaderIndices.getVertexShaderIndex(material.shaders());
@@ -273,44 +201,4 @@ public class InstancedDrawManager extends DrawManager<InstancedInstancer<?>> {
 		int packedMaterialProperties = MaterialEncoder.packProperties(material);
 		GL32.glUniform4ui(uniformLocation, vertexIndex, fragmentIndex, packedFogAndCutout, packedMaterialProperties);
 	}
-
-	public static class DrawSet implements Iterable<Map.Entry<ShaderState, Collection<DrawCall>>> {
-		public static final DrawSet EMPTY = new DrawSet(ImmutableListMultimap.of());
-
-		private final ListMultimap<ShaderState, DrawCall> drawCalls;
-
-		public DrawSet(RenderStage renderStage) {
-			drawCalls = ArrayListMultimap.create();
-		}
-
-		public DrawSet(ListMultimap<ShaderState, DrawCall> drawCalls) {
-			this.drawCalls = drawCalls;
-		}
-
-		private void delete() {
-			drawCalls.values()
-					.forEach(DrawCall::delete);
-			drawCalls.clear();
-		}
-
-		public void put(ShaderState shaderState, DrawCall drawCall) {
-			drawCalls.put(shaderState, drawCall);
-		}
-
-		public boolean isEmpty() {
-			return drawCalls.isEmpty();
-		}
-
-		@Override
-		public Iterator<Map.Entry<ShaderState, Collection<DrawCall>>> iterator() {
-			return drawCalls.asMap()
-					.entrySet()
-					.iterator();
-		}
-
-		public void prune() {
-            drawCalls.values()
-                    .removeIf(DrawCall::deleted);
-		}
-	}
 }

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

@@ -23,7 +23,7 @@ public class InstancedInstancer<I extends Instance> extends AbstractInstancer<I>
 	@Nullable
 	private GlBuffer vbo;
 
-	private final List<DrawCall> drawCalls = new ArrayList<>();
+	private final List<InstancedDraw> draws = new ArrayList<>();
 
 	public InstancedInstancer(InstanceType<I> type, Environment environment) {
 		super(type, environment);
@@ -33,6 +33,10 @@ public class InstancedInstancer<I extends Instance> extends AbstractInstancer<I>
 		writer = type.writer();
 	}
 
+	public List<InstancedDraw> draws() {
+		return draws;
+	}
+
 	public void init() {
 		if (vbo != null) {
 			return;
@@ -115,17 +119,13 @@ public class InstancedInstancer<I extends Instance> extends AbstractInstancer<I>
 		vbo.delete();
 		vbo = null;
 
-		for (DrawCall drawCall : drawCalls) {
-			drawCall.delete();
+		for (InstancedDraw instancedDraw : draws) {
+			instancedDraw.delete();
 		}
 	}
 
-	public void addDrawCall(DrawCall drawCall) {
-		drawCalls.add(drawCall);
-	}
-
-	public List<DrawCall> drawCalls() {
-		return drawCalls;
+	public void addDrawCall(InstancedDraw instancedDraw) {
+		draws.add(instancedDraw);
 	}
 
 	public void bind(TextureBuffer buffer) {

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

@@ -0,0 +1,102 @@
+package com.jozufozu.flywheel.backend.engine.instancing;
+
+import static com.jozufozu.flywheel.backend.engine.instancing.InstancedDrawManager.uploadMaterialUniform;
+
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import com.jozufozu.flywheel.backend.Samplers;
+import com.jozufozu.flywheel.backend.compile.InstancingPrograms;
+import com.jozufozu.flywheel.backend.engine.GroupKey;
+import com.jozufozu.flywheel.backend.engine.MaterialRenderState;
+import com.jozufozu.flywheel.backend.gl.TextureBuffer;
+
+public class InstancedRenderStage {
+	private static final Comparator<InstancedDraw> DRAW_COMPARATOR = Comparator.comparing(InstancedDraw::indexOfMeshInModel)
+			.thenComparing(InstancedDraw::material, MaterialRenderState.COMPARATOR);
+
+	private final Map<GroupKey<?>, DrawGroup> groups = new HashMap<>();
+
+	public InstancedRenderStage() {
+	}
+
+	public void delete() {
+		groups.values()
+				.forEach(DrawGroup::delete);
+		groups.clear();
+	}
+
+	public void put(GroupKey<?> groupKey, InstancedDraw instancedDraw) {
+		groups.computeIfAbsent(groupKey, $ -> new DrawGroup())
+				.put(instancedDraw);
+	}
+
+	public boolean isEmpty() {
+		return groups.isEmpty();
+	}
+
+	public void flush() {
+		groups.values()
+				.forEach(DrawGroup::flush);
+
+		groups.values()
+				.removeIf(DrawGroup::isEmpty);
+	}
+
+	public void draw(TextureBuffer instanceTexture, InstancingPrograms programs) {
+		for (var entry : groups.entrySet()) {
+			var shader = entry.getKey();
+			var drawCalls = entry.getValue();
+
+			var environment = shader.environment();
+
+			var program = programs.get(shader.instanceType(), environment.contextShader());
+			program.bind();
+
+			environment.setupDraw(program);
+
+			for (var drawCall : drawCalls.draws) {
+				var material = drawCall.material();
+
+				uploadMaterialUniform(program, material);
+
+				MaterialRenderState.setup(material);
+
+				Samplers.INSTANCE_BUFFER.makeActive();
+
+				drawCall.render(instanceTexture);
+			}
+		}
+	}
+
+	public static class DrawGroup {
+		private final List<InstancedDraw> draws = new ArrayList<>();
+		private boolean needSort = false;
+
+		public void put(InstancedDraw instancedDraw) {
+			draws.add(instancedDraw);
+			needSort = true;
+		}
+
+		public void delete() {
+			draws.forEach(InstancedDraw::delete);
+			draws.clear();
+		}
+
+		public void flush() {
+			needSort |= draws.removeIf(InstancedDraw::deleted);
+
+			if (needSort) {
+				draws.sort(DRAW_COMPARATOR);
+				needSort = false;
+			}
+		}
+
+		public boolean isEmpty() {
+			return draws.isEmpty();
+		}
+	}
+}

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

@@ -1,8 +0,0 @@
-package com.jozufozu.flywheel.backend.engine.instancing;
-
-import com.jozufozu.flywheel.api.instance.InstanceType;
-import com.jozufozu.flywheel.api.material.Material;
-import com.jozufozu.flywheel.backend.engine.embed.Environment;
-
-public record ShaderState(Material material, InstanceType<?> instanceType, Environment environment) {
-}