Just not my type

- Remove VisualType
- Engines now render everything in a single pass, and instancers can be
  shared between entities, block entities, and effects

common/src/api/java/dev/engine_room/flywheel/api/backend/Engine.java

@@ -8,7 +8,6 @@ import org.jetbrains.annotations.Range;
 import dev.engine_room.flywheel.api.RenderContext;
 import dev.engine_room.flywheel.api.instance.Instance;
 import dev.engine_room.flywheel.api.task.Plan;
-import dev.engine_room.flywheel.api.visualization.VisualType;
 import dev.engine_room.flywheel.api.visualization.VisualizationContext;
 import it.unimi.dsi.fastutil.longs.LongSet;
 import net.minecraft.client.Camera;
@@ -22,10 +21,9 @@ public interface Engine {
 	/**
 	 * Create a visualization context that will be used to create visuals of the given type.
 	 *
-	 * @param visualType The type of visual.
 	 * @return A new visualization context.
 	 */
-	VisualizationContext createVisualizationContext(VisualType visualType);
+	VisualizationContext createVisualizationContext();
 
 	/**
 	 * Create a plan that will start execution after the start of the level render and
@@ -78,9 +76,8 @@ public interface Engine {
 	 * level render. This method is guaranteed to be called on the render thread.
 	 *
 	 * @param context The context for the current level render.
-	 * @param visualType The type of visual.
 	 */
-	void render(RenderContext context, VisualType visualType);
+	void render(RenderContext context);
 
 	/**
 	 * Render the given instances as a crumbling overlay.

common/src/api/java/dev/engine_room/flywheel/api/visualization/VisualType.java

@@ -1,7 +0,0 @@
-package dev.engine_room.flywheel.api.visualization;
-
-public enum VisualType {
-	BLOCK_ENTITY,
-	ENTITY,
-	EFFECT;
-}

common/src/api/java/dev/engine_room/flywheel/api/visualization/VisualizationManager.java

@@ -49,12 +49,8 @@ public interface VisualizationManager {
 	interface RenderDispatcher {
 		void onStartLevelRender(RenderContext ctx);
 
-		void afterBlockEntities(RenderContext ctx);
-
 		void afterEntities(RenderContext ctx);
 
 		void beforeCrumbling(RenderContext ctx, Long2ObjectMap<SortedSet<BlockDestructionProgress>> destructionProgress);
-
-		void afterParticles(RenderContext ctx);
 	}
 }

common/src/backend/java/dev/engine_room/flywheel/backend/engine/DrawManager.java

@@ -18,7 +18,6 @@ import dev.engine_room.flywheel.api.instance.Instance;
 import dev.engine_room.flywheel.api.instance.InstanceType;
 import dev.engine_room.flywheel.api.model.Model;
 import dev.engine_room.flywheel.api.task.Plan;
-import dev.engine_room.flywheel.api.visualization.VisualType;
 import dev.engine_room.flywheel.backend.FlwBackend;
 import dev.engine_room.flywheel.backend.engine.embed.Environment;
 import dev.engine_room.flywheel.backend.engine.embed.EnvironmentStorage;
@@ -43,8 +42,8 @@ public abstract class DrawManager<N extends AbstractInstancer<?>> {
 	 */
 	protected final Queue<UninitializedInstancer<N, ?>> initializationQueue = new ConcurrentLinkedQueue<>();
 
-	public <I extends Instance> AbstractInstancer<I> getInstancer(Environment environment, InstanceType<I> type, Model model, VisualType visualType, int bias) {
+	public <I extends Instance> AbstractInstancer<I> getInstancer(Environment environment, InstanceType<I> type, Model model, int bias) {
-		return getInstancer(new InstancerKey<>(environment, type, model, visualType, bias));
+		return getInstancer(new InstancerKey<>(environment, type, model, bias));
 	}
 
 	@SuppressWarnings("unchecked")
@@ -76,7 +75,7 @@ public abstract class DrawManager<N extends AbstractInstancer<?>> {
 				.forEach(AbstractInstancer::clear);
 	}
 
-	public abstract void render(VisualType visualType);
+	public abstract void render();
 
 	public abstract void renderCrumbling(List<Engine.CrumblingBlock> crumblingBlocks);
 

common/src/backend/java/dev/engine_room/flywheel/backend/engine/EngineImpl.java

@@ -13,7 +13,6 @@ import dev.engine_room.flywheel.api.instance.InstancerProvider;
 import dev.engine_room.flywheel.api.model.Model;
 import dev.engine_room.flywheel.api.task.Plan;
 import dev.engine_room.flywheel.api.visualization.VisualEmbedding;
-import dev.engine_room.flywheel.api.visualization.VisualType;
 import dev.engine_room.flywheel.api.visualization.VisualizationContext;
 import dev.engine_room.flywheel.backend.FlwBackend;
 import dev.engine_room.flywheel.backend.compile.core.ShaderException;
@@ -47,8 +46,8 @@ public class EngineImpl implements Engine {
 	}
 
 	@Override
-	public VisualizationContext createVisualizationContext(VisualType visualType) {
+	public VisualizationContext createVisualizationContext() {
-		return new VisualizationContextImpl(visualType);
+		return new VisualizationContextImpl();
 	}
 
 	@Override
@@ -104,9 +103,9 @@ public class EngineImpl implements Engine {
 	}
 
 	@Override
-	public void render(RenderContext context, VisualType visualType) {
+	public void render(RenderContext context) {
 		try (var state = GlStateTracker.getRestoreState()) {
-			drawManager.render(visualType);
+			drawManager.render();
 		} catch (ShaderException e) {
 			FlwBackend.LOGGER.error("Falling back", e);
 			triggerFallback();
@@ -134,8 +133,8 @@ public class EngineImpl implements Engine {
 		drawManager.triggerFallback();
 	}
 
-	public <I extends Instance> Instancer<I> instancer(Environment environment, InstanceType<I> type, Model model, VisualType visualType, int bias) {
+	public <I extends Instance> Instancer<I> instancer(Environment environment, InstanceType<I> type, Model model, int bias) {
-		return drawManager.getInstancer(environment, type, model, visualType, bias);
+		return drawManager.getInstancer(environment, type, model, bias);
 	}
 
 	public EnvironmentStorage environmentStorage() {
@@ -148,11 +147,9 @@ public class EngineImpl implements Engine {
 
 	private class VisualizationContextImpl implements VisualizationContext {
 		private final InstancerProviderImpl instancerProvider;
-		private final VisualType visualType;
 
-		public VisualizationContextImpl(VisualType visualType) {
+		public VisualizationContextImpl() {
-			instancerProvider = new InstancerProviderImpl(EngineImpl.this, visualType);
+			instancerProvider = new InstancerProviderImpl(EngineImpl.this);
-			this.visualType = visualType;
 		}
 
 		@Override
@@ -167,7 +164,7 @@ public class EngineImpl implements Engine {
 
 		@Override
 		public VisualEmbedding createEmbedding(Vec3i renderOrigin) {
-			var out = new EmbeddedEnvironment(EngineImpl.this, visualType, renderOrigin);
+			var out = new EmbeddedEnvironment(EngineImpl.this, renderOrigin);
 			environmentStorage.track(out);
 			return out;
 		}

common/src/backend/java/dev/engine_room/flywheel/backend/engine/InstancerKey.java

@@ -3,9 +3,7 @@ package dev.engine_room.flywheel.backend.engine;
 import dev.engine_room.flywheel.api.instance.Instance;
 import dev.engine_room.flywheel.api.instance.InstanceType;
 import dev.engine_room.flywheel.api.model.Model;
-import dev.engine_room.flywheel.api.visualization.VisualType;
 import dev.engine_room.flywheel.backend.engine.embed.Environment;
 
-public record InstancerKey<I extends Instance>(Environment environment, InstanceType<I> type, Model model,
+public record InstancerKey<I extends Instance>(Environment environment, InstanceType<I> type, Model model, int bias) {
-											   VisualType visualType, int bias) {
 }

common/src/backend/java/dev/engine_room/flywheel/backend/engine/InstancerProviderImpl.java

@@ -5,12 +5,11 @@ import dev.engine_room.flywheel.api.instance.InstanceType;
 import dev.engine_room.flywheel.api.instance.Instancer;
 import dev.engine_room.flywheel.api.instance.InstancerProvider;
 import dev.engine_room.flywheel.api.model.Model;
-import dev.engine_room.flywheel.api.visualization.VisualType;
 import dev.engine_room.flywheel.backend.engine.embed.GlobalEnvironment;
 
-public record InstancerProviderImpl(EngineImpl engine, VisualType visualType) implements InstancerProvider {
+public record InstancerProviderImpl(EngineImpl engine) implements InstancerProvider {
 	@Override
 	public <I extends Instance> Instancer<I> instancer(InstanceType<I> type, Model model, int bias) {
-		return engine.instancer(GlobalEnvironment.INSTANCE, type, model, visualType, bias);
+		return engine.instancer(GlobalEnvironment.INSTANCE, type, model, bias);
 	}
 }

common/src/backend/java/dev/engine_room/flywheel/backend/engine/embed/EmbeddedEnvironment.java

@@ -12,7 +12,6 @@ import dev.engine_room.flywheel.api.instance.Instancer;
 import dev.engine_room.flywheel.api.instance.InstancerProvider;
 import dev.engine_room.flywheel.api.model.Model;
 import dev.engine_room.flywheel.api.visualization.VisualEmbedding;
-import dev.engine_room.flywheel.api.visualization.VisualType;
 import dev.engine_room.flywheel.backend.compile.ContextShader;
 import dev.engine_room.flywheel.backend.engine.EngineImpl;
 import dev.engine_room.flywheel.backend.gl.shader.GlProgram;
@@ -21,7 +20,6 @@ import net.minecraft.core.Vec3i;
 
 public class EmbeddedEnvironment implements VisualEmbedding, Environment {
 	private final EngineImpl engine;
-	private final VisualType visualType;
 	private final Vec3i renderOrigin;
 	@Nullable
 	private final EmbeddedEnvironment parent;
@@ -36,9 +34,8 @@ public class EmbeddedEnvironment implements VisualEmbedding, Environment {
 
 	private boolean deleted = false;
 
-	public EmbeddedEnvironment(EngineImpl engine, VisualType visualType, Vec3i renderOrigin, @Nullable EmbeddedEnvironment parent) {
+	public EmbeddedEnvironment(EngineImpl engine, Vec3i renderOrigin, @Nullable EmbeddedEnvironment parent) {
 		this.engine = engine;
-		this.visualType = visualType;
 		this.renderOrigin = renderOrigin;
 		this.parent = parent;
 
@@ -46,13 +43,13 @@ public class EmbeddedEnvironment implements VisualEmbedding, Environment {
 			@Override
 			public <I extends Instance> Instancer<I> instancer(InstanceType<I> type, Model model, int bias) {
 				// Kinda cursed usage of anonymous classes here, but it does the job.
-				return engine.instancer(EmbeddedEnvironment.this, type, model, visualType, bias);
+				return engine.instancer(EmbeddedEnvironment.this, type, model, bias);
 			}
 		};
 	}
 
-	public EmbeddedEnvironment(EngineImpl engine, VisualType visualType, Vec3i renderOrigin) {
+	public EmbeddedEnvironment(EngineImpl engine, Vec3i renderOrigin) {
-		this(engine, visualType, renderOrigin, null);
+		this(engine, renderOrigin, null);
 	}
 
 	@Override
@@ -73,7 +70,7 @@ public class EmbeddedEnvironment implements VisualEmbedding, Environment {
 
 	@Override
 	public VisualEmbedding createEmbedding(Vec3i renderOrigin) {
-		var out = new EmbeddedEnvironment(engine, visualType, renderOrigin, this);
+		var out = new EmbeddedEnvironment(engine, renderOrigin, this);
 		engine.environmentStorage()
 				.track(out);
 		return out;

common/src/backend/java/dev/engine_room/flywheel/backend/engine/indirect/IndirectCullingGroup.java

@@ -8,15 +8,12 @@ import static org.lwjgl.opengl.GL43.glDispatchCompute;
 
 import java.util.ArrayList;
 import java.util.Comparator;
-import java.util.EnumMap;
 import java.util.List;
-import java.util.Map;
 
 import dev.engine_room.flywheel.api.instance.Instance;
 import dev.engine_room.flywheel.api.instance.InstanceType;
 import dev.engine_room.flywheel.api.material.Material;
 import dev.engine_room.flywheel.api.model.Model;
-import dev.engine_room.flywheel.api.visualization.VisualType;
 import dev.engine_room.flywheel.backend.compile.ContextShader;
 import dev.engine_room.flywheel.backend.compile.IndirectPrograms;
 import dev.engine_room.flywheel.backend.engine.InstancerKey;
@@ -28,8 +25,7 @@ import dev.engine_room.flywheel.backend.gl.shader.GlProgram;
 import dev.engine_room.flywheel.lib.math.MoreMath;
 
 public class IndirectCullingGroup<I extends Instance> {
-	private static final Comparator<IndirectDraw> DRAW_COMPARATOR = Comparator.comparing(IndirectDraw::visualType)
+	private static final Comparator<IndirectDraw> DRAW_COMPARATOR = Comparator.comparing(IndirectDraw::isEmbedded)
-			.thenComparing(IndirectDraw::isEmbedded)
 			.thenComparing(IndirectDraw::bias)
 			.thenComparing(IndirectDraw::indexOfMeshInModel)
 			.thenComparing(IndirectDraw::material, MaterialRenderState.COMPARATOR);
@@ -39,7 +35,7 @@ public class IndirectCullingGroup<I extends Instance> {
 	private final IndirectBuffers buffers;
 	private final List<IndirectInstancer<I>> instancers = new ArrayList<>();
 	private final List<IndirectDraw> indirectDraws = new ArrayList<>();
-	private final Map<VisualType, List<MultiDraw>> multiDraws = new EnumMap<>(VisualType.class);
+	private final List<MultiDraw> multiDraws = new ArrayList<>();
 
 	private final IndirectPrograms programs;
 	private final GlProgram cullProgram;
@@ -132,10 +128,6 @@ public class IndirectCullingGroup<I extends Instance> {
 		return indirectDraws.isEmpty() || instanceCountThisFrame == 0;
 	}
 
-	private boolean nothingToDo(VisualType visualType) {
-		return nothingToDo() || !multiDraws.containsKey(visualType);
-	}
-
 	private void sortDraws() {
 		multiDraws.clear();
 		// sort by visual type, then material
@@ -146,28 +138,19 @@ public class IndirectCullingGroup<I extends Instance> {
 
 			// if the next draw call has a different VisualType or Material, start a new MultiDraw
 			if (i == indirectDraws.size() - 1 || incompatibleDraws(draw1, indirectDraws.get(i + 1))) {
-				multiDraws.computeIfAbsent(draw1.visualType(), s -> new ArrayList<>())
+				multiDraws.add(new MultiDraw(draw1.material(), draw1.isEmbedded(), start, i + 1));
-						.add(new MultiDraw(draw1.material(), draw1.isEmbedded(), start, i + 1));
 				start = i + 1;
 			}
 		}
 	}
 
 	private boolean incompatibleDraws(IndirectDraw draw1, IndirectDraw draw2) {
-		if (draw1.visualType() != draw2.visualType()) {
-			return true;
-		}
-
 		if (draw1.isEmbedded() != draw2.isEmbedded()) {
 			return true;
 		}
 		return !MaterialRenderState.materialEquals(draw1.material(), draw2.material());
 	}
 
-	public boolean hasVisualType(VisualType visualType) {
-		return multiDraws.containsKey(visualType);
-	}
-
 	public void add(IndirectInstancer<I> instancer, InstancerKey<I> key, MeshPool meshPool) {
 		instancer.mapping = buffers.objectStorage.createMapping();
 		instancer.update(instancers.size(), -1);
@@ -180,7 +163,7 @@ public class IndirectCullingGroup<I extends Instance> {
             var entry = meshes.get(i);
 
             MeshPool.PooledMesh mesh = meshPool.alloc(entry.mesh());
-            var draw = new IndirectDraw(instancer, entry.material(), mesh, key.visualType(), key.bias(), i);
+			var draw = new IndirectDraw(instancer, entry.material(), mesh, key.bias(), i);
             indirectDraws.add(draw);
             instancer.addDraw(draw);
         }
@@ -188,8 +171,8 @@ public class IndirectCullingGroup<I extends Instance> {
 		needsDrawSort = true;
 	}
 
-	public void submit(VisualType visualType) {
+	public void submit() {
-		if (nothingToDo(visualType)) {
+		if (nothingToDo()) {
 			return;
 		}
 
@@ -199,7 +182,7 @@ public class IndirectCullingGroup<I extends Instance> {
 
 		GlProgram lastProgram = null;
 
-		for (var multiDraw : multiDraws.get(visualType)) {
+		for (var multiDraw : multiDraws) {
 			var drawProgram = programs.getIndirectProgram(instanceType, multiDraw.embedded ? ContextShader.EMBEDDED : ContextShader.DEFAULT, multiDraw.material);
 			if (drawProgram != lastProgram) {
 				lastProgram = drawProgram;

common/src/backend/java/dev/engine_room/flywheel/backend/engine/indirect/IndirectDraw.java

@@ -3,7 +3,6 @@ package dev.engine_room.flywheel.backend.engine.indirect;
 import org.lwjgl.system.MemoryUtil;
 
 import dev.engine_room.flywheel.api.material.Material;
-import dev.engine_room.flywheel.api.visualization.VisualType;
 import dev.engine_room.flywheel.backend.engine.MaterialEncoder;
 import dev.engine_room.flywheel.backend.engine.MeshPool;
 import dev.engine_room.flywheel.backend.engine.embed.EmbeddedEnvironment;
@@ -12,7 +11,6 @@ public class IndirectDraw {
 	private final IndirectInstancer<?> instancer;
 	private final Material material;
 	private final MeshPool.PooledMesh mesh;
-	private final VisualType visualType;
 	private final int bias;
 	private final int indexOfMeshInModel;
 
@@ -20,11 +18,10 @@ public class IndirectDraw {
 	private final int packedMaterialProperties;
 	private boolean deleted;
 
-	public IndirectDraw(IndirectInstancer<?> instancer, Material material, MeshPool.PooledMesh mesh, VisualType visualType, int bias, int indexOfMeshInModel) {
+	public IndirectDraw(IndirectInstancer<?> instancer, Material material, MeshPool.PooledMesh mesh, int bias, int indexOfMeshInModel) {
 		this.instancer = instancer;
 		this.material = material;
 		this.mesh = mesh;
-		this.visualType = visualType;
 		this.bias = bias;
 		this.indexOfMeshInModel = indexOfMeshInModel;
 
@@ -50,10 +47,6 @@ public class IndirectDraw {
 		return mesh;
 	}
 
-	public VisualType visualType() {
-		return visualType;
-	}
-
 	public int bias() {
 		return bias;
 	}

common/src/backend/java/dev/engine_room/flywheel/backend/engine/indirect/IndirectDrawManager.java

@@ -15,7 +15,6 @@ import java.util.Map;
 import dev.engine_room.flywheel.api.backend.Engine;
 import dev.engine_room.flywheel.api.instance.Instance;
 import dev.engine_room.flywheel.api.instance.InstanceType;
-import dev.engine_room.flywheel.api.visualization.VisualType;
 import dev.engine_room.flywheel.backend.Samplers;
 import dev.engine_room.flywheel.backend.compile.IndirectPrograms;
 import dev.engine_room.flywheel.backend.engine.AbstractInstancer;
@@ -79,20 +78,7 @@ public class IndirectDrawManager extends DrawManager<IndirectInstancer<?>> {
 		group.add((IndirectInstancer<I>) instancer, key, meshPool);
 	}
 
-	public boolean hasVisualType(VisualType visualType) {
+	public void render() {
-		for (var group : cullingGroups.values()) {
-			if (group.hasVisualType(visualType)) {
-				return true;
-			}
-		}
-		return false;
-	}
-
-	public void render(VisualType visualType) {
-		if (!hasVisualType(visualType)) {
-			return;
-		}
-
 		TextureBinder.bindLightAndOverlay();
 
 		vertexArray.bindForDraw();
@@ -106,7 +92,7 @@ public class IndirectDrawManager extends DrawManager<IndirectInstancer<?>> {
 		}
 
 		for (var group : cullingGroups.values()) {
-			group.submit(visualType);
+			group.submit();
 		}
 
 		MaterialRenderState.reset();

common/src/backend/java/dev/engine_room/flywheel/backend/engine/instancing/InstancedDrawManager.java

@@ -1,13 +1,10 @@
 package dev.engine_room.flywheel.backend.engine.instancing;
 
-import java.util.EnumMap;
 import java.util.List;
-import java.util.Map;
 
 import dev.engine_room.flywheel.api.backend.Engine;
 import dev.engine_room.flywheel.api.instance.Instance;
 import dev.engine_room.flywheel.api.material.Material;
-import dev.engine_room.flywheel.api.visualization.VisualType;
 import dev.engine_room.flywheel.backend.Samplers;
 import dev.engine_room.flywheel.backend.compile.ContextShader;
 import dev.engine_room.flywheel.backend.compile.InstancingPrograms;
@@ -31,10 +28,7 @@ import net.minecraft.client.Minecraft;
 import net.minecraft.client.resources.model.ModelBakery;
 
 public class InstancedDrawManager extends DrawManager<InstancedInstancer<?>> {
-	/**
+	private final InstancedRenderStage draws = new InstancedRenderStage();
-	 * The set of draw calls to make for each {@link VisualType}.
-	 */
-	private final Map<VisualType, InstancedRenderStage> stages = new EnumMap<>(VisualType.class);
 	private final InstancingPrograms programs;
 	/**
 	 * A map of vertex types to their mesh pools.
@@ -71,10 +65,8 @@ public class InstancedDrawManager extends DrawManager<InstancedInstancer<?>> {
 			}
 		});
 
-		for (InstancedRenderStage stage : stages.values()) {
+		// Remove the draw calls for any instancers we deleted.
-			// Remove the draw calls for any instancers we deleted.
+		draws.flush();
-			stage.flush();
-		}
 
 		meshPool.flush();
 
@@ -82,10 +74,10 @@ public class InstancedDrawManager extends DrawManager<InstancedInstancer<?>> {
 	}
 
 	@Override
-	public void render(VisualType visualType) {
+	public void render() {
-		var stage = stages.get(visualType);
+		var stage = draws;
 
-		if (stage == null || stage.isEmpty()) {
+		if (stage.isEmpty()) {
 			return;
 		}
 
@@ -105,9 +97,7 @@ public class InstancedDrawManager extends DrawManager<InstancedInstancer<?>> {
 		instancers.values()
 				.forEach(InstancedInstancer::delete);
 
-		stages.values()
+		draws.delete();
-				.forEach(InstancedRenderStage::delete);
-		stages.clear();
 
 		meshPool.delete();
 		instanceTexture.delete();
@@ -128,8 +118,6 @@ public class InstancedDrawManager extends DrawManager<InstancedInstancer<?>> {
 	protected <I extends Instance> void initialize(InstancerKey<I> key, InstancedInstancer<?> instancer) {
 		instancer.init();
 
-		InstancedRenderStage stage = stages.computeIfAbsent(key.visualType(), $ -> new InstancedRenderStage());
-
 		var meshes = key.model()
 				.meshes();
 		for (int i = 0; i < meshes.size(); i++) {
@@ -139,7 +127,7 @@ public class InstancedDrawManager extends DrawManager<InstancedInstancer<?>> {
 			GroupKey<?> groupKey = new GroupKey<>(key.type(), key.environment());
 			InstancedDraw instancedDraw = new InstancedDraw(instancer, mesh, groupKey, entry.material(), key.bias(), i);
 
-			stage.put(groupKey, instancedDraw);
+			draws.put(groupKey, instancedDraw);
 			instancer.addDrawCall(instancedDraw);
 		}
 	}

common/src/main/java/dev/engine_room/flywheel/impl/mixin/LevelRendererMixin.java

@@ -9,8 +9,6 @@ import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
 import org.spongepowered.asm.mixin.Unique;
 import org.spongepowered.asm.mixin.injection.At;
-import org.spongepowered.asm.mixin.injection.At.Shift;
-import org.spongepowered.asm.mixin.injection.Group;
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 
@@ -82,16 +80,6 @@ abstract class LevelRendererMixin {
 		}
 	}
 
-	@Inject(method = "renderLevel", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/renderer/OutlineBufferSource;endOutlineBatch()V", ordinal = 0))
-	private void flywheel$afterBlockEntities(CallbackInfo ci) {
-		if (flywheel$renderContext != null) {
-			VisualizationManager manager = VisualizationManager.get(level);
-			if (manager != null) {
-				manager.renderDispatcher().afterBlockEntities(flywheel$renderContext);
-			}
-		}
-	}
-
 	@Inject(method = "renderLevel", at = @At(value = "INVOKE_STRING", target = "Lnet/minecraft/util/profiling/ProfilerFiller;popPush(Ljava/lang/String;)V", args = "ldc=destroyProgress"))
 	private void flywheel$beforeRenderCrumbling(CallbackInfo ci) {
 		if (flywheel$renderContext != null) {
@@ -102,28 +90,6 @@ abstract class LevelRendererMixin {
 		}
 	}
 
-	@Group(name = "afterParticles", min = 2, max = 2)
-	@Inject(method = "renderLevel", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/particle/ParticleEngine;render(Lcom/mojang/blaze3d/vertex/PoseStack;Lnet/minecraft/client/renderer/MultiBufferSource$BufferSource;Lnet/minecraft/client/renderer/LightTexture;Lnet/minecraft/client/Camera;F)V", shift = Shift.AFTER))
-	private void flywheel$afterParticles$fabric(CallbackInfo ci) {
-		if (flywheel$renderContext != null) {
-			VisualizationManager manager = VisualizationManager.get(level);
-			if (manager != null) {
-				manager.renderDispatcher().afterParticles(flywheel$renderContext);
-			}
-		}
-	}
-
-	@Group(name = "afterParticles")
-	@Inject(method = "renderLevel", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/particle/ParticleEngine;render(Lcom/mojang/blaze3d/vertex/PoseStack;Lnet/minecraft/client/renderer/MultiBufferSource$BufferSource;Lnet/minecraft/client/renderer/LightTexture;Lnet/minecraft/client/Camera;FLnet/minecraft/client/renderer/culling/Frustum;)V", shift = Shift.AFTER))
-	private void flywheel$afterParticles$forge(CallbackInfo ci) {
-		if (flywheel$renderContext != null) {
-			VisualizationManager manager = VisualizationManager.get(level);
-			if (manager != null) {
-				manager.renderDispatcher().afterParticles(flywheel$renderContext);
-			}
-		}
-	}
-
 	@Inject(method = "renderEntity", at = @At("HEAD"), cancellable = true)
 	private void flywheel$decideNotToRenderEntity(Entity entity, double camX, double camY, double camZ, float partialTick, PoseStack poseStack, MultiBufferSource bufferSource, CallbackInfo ci) {
 		if (VisualizationManager.supportsVisualization(entity.level()) && VisualizationHelper.skipVanillaRender(entity)) {

common/src/main/java/dev/engine_room/flywheel/impl/visualization/VisualizationManagerImpl.java

@@ -18,7 +18,6 @@ import dev.engine_room.flywheel.api.visual.DynamicVisual;
 import dev.engine_room.flywheel.api.visual.Effect;
 import dev.engine_room.flywheel.api.visual.TickableVisual;
 import dev.engine_room.flywheel.api.visualization.VisualManager;
-import dev.engine_room.flywheel.api.visualization.VisualType;
 import dev.engine_room.flywheel.api.visualization.VisualizationLevel;
 import dev.engine_room.flywheel.api.visualization.VisualizationManager;
 import dev.engine_room.flywheel.impl.FlwConfig;
@@ -82,9 +81,10 @@ public class VisualizationManagerImpl implements VisualizationManager {
 				.createEngine(level);
 		frameLimiter = createUpdateLimiter();
 
-		var blockEntitiesStorage = new BlockEntityStorage(engine.createVisualizationContext(VisualType.BLOCK_ENTITY));
+		var visualizationContext = engine.createVisualizationContext();
-		var entitiesStorage = new EntityStorage(engine.createVisualizationContext(VisualType.ENTITY));
+		var blockEntitiesStorage = new BlockEntityStorage(visualizationContext);
-		var effectsStorage = new EffectStorage(engine.createVisualizationContext(VisualType.EFFECT));
+		var entitiesStorage = new EntityStorage(visualizationContext);
+		var effectsStorage = new EffectStorage(visualizationContext);
 
 		blockEntities = new VisualManagerImpl<>(blockEntitiesStorage);
 		entities = new VisualManagerImpl<>(entitiesStorage);
@@ -257,9 +257,9 @@ public class VisualizationManagerImpl implements VisualizationManager {
 	/**
 	 * Draw all visuals of the given type.
 	 */
-	private void render(RenderContext context, VisualType visualType) {
+	private void render(RenderContext context) {
 		ensureCanRender(context);
-		engine.render(context, visualType);
+		engine.render(context);
 	}
 
 	private void renderCrumbling(RenderContext context, Long2ObjectMap<SortedSet<BlockDestructionProgress>> destructionProgress) {
@@ -337,25 +337,15 @@ public class VisualizationManagerImpl implements VisualizationManager {
 			beginFrame(ctx);
 		}
 
-		@Override
-		public void afterBlockEntities(RenderContext ctx) {
-			render(ctx, VisualType.BLOCK_ENTITY);
-		}
-
 		@Override
 		public void afterEntities(RenderContext ctx) {
-			render(ctx, VisualType.ENTITY);
+			render(ctx);
 		}
 
 		@Override
 		public void beforeCrumbling(RenderContext ctx, Long2ObjectMap<SortedSet<BlockDestructionProgress>> destructionProgress) {
 			renderCrumbling(ctx, destructionProgress);
 		}
-
-		@Override
-		public void afterParticles(RenderContext ctx) {
-			render(ctx, VisualType.EFFECT);
-		}
 	}
 
 	private record CrumblingBlockImpl(BlockPos pos, int progress, List<Instance> instances) implements Engine.CrumblingBlock {