Taking textures back out of context

- Contexts shouldn't be responsible for binding the api defined textures
- Move light and overlay binding to TextureBinder
- Clean up & standardize material/texture setup/teardown calls
- Revert removal of MaterialRenderState texture stuff
- Add SimpleContext + Builder
- Rename Textures -> TextureSource
- Add API methods to set basic float uniforms, more to come later
- Throw errors if a context tries to set uniforms/textures with reserved
  names.

src/main/java/com/jozufozu/flywheel/api/context/Context.java

@@ -10,7 +10,7 @@ public interface Context {
 	 *
 	 * @param material The material about to be rendered.
 	 * @param shader   The shader to prepare.
-	 * @param textures Source of the textures to use.
+	 * @param textureSource Source of the textures to use.
 	 */
-	void prepare(Material material, Shader shader, Textures textures);
+	void prepare(Material material, Shader shader, TextureSource textureSource);
 }

src/main/java/com/jozufozu/flywheel/api/context/Shader.java

@@ -1,8 +1,20 @@
 package com.jozufozu.flywheel.api.context;
 
+import org.joml.Matrix4fc;
+
 import com.jozufozu.flywheel.api.BackendImplemented;
 
 @BackendImplemented
 public interface Shader {
 	void setTexture(String glslName, Texture texture);
+
+	void setFloat(String glslName, float value);
+
+	void setVec2(String glslName, float x, float y);
+
+	void setVec3(String glslName, float x, float y, float z);
+
+	void setVec4(String glslName, float x, float y, float z, float w);
+
+	void setMat4(String glslName, Matrix4fc matrix);
 }

src/main/java/com/jozufozu/flywheel/api/context/TextureSource.java

@@ -5,7 +5,7 @@ import com.jozufozu.flywheel.api.BackendImplemented;
 import net.minecraft.resources.ResourceLocation;
 
 @BackendImplemented
-public interface Textures {
+public interface TextureSource {
 	/**
 	 * Get a built-in texture by its resource location.
 	 *

src/main/java/com/jozufozu/flywheel/backend/compile/PipelineCompiler.java

@@ -5,6 +5,7 @@ import java.util.List;
 import com.jozufozu.flywheel.backend.InternalVertex;
 import com.jozufozu.flywheel.backend.compile.core.CompilationHarness;
 import com.jozufozu.flywheel.backend.compile.core.Compile;
+import com.jozufozu.flywheel.backend.gl.shader.GlProgram;
 import com.jozufozu.flywheel.backend.gl.shader.ShaderType;
 import com.jozufozu.flywheel.backend.glsl.ShaderSources;
 import com.jozufozu.flywheel.backend.glsl.SourceComponent;
@@ -52,6 +53,14 @@ public class PipelineCompiler {
 				.then((key, program) -> {
 					program.setUniformBlockBinding("_FlwFrameUniforms", 0);
 					program.setUniformBlockBinding("_FlwFogUniforms", 1);
+
+					program.bind();
+
+					program.setSamplerBinding("_flw_diffuseTex", 0);
+					program.setSamplerBinding("_flw_overlayTex", 1);
+					program.setSamplerBinding("_flw_lightTex", 2);
+
+					GlProgram.unbind();
 				})
 				.harness(pipeline.compilerMarker(), sources);
 	}

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

@@ -8,9 +8,13 @@ import com.jozufozu.flywheel.api.material.DepthTest;
 import com.jozufozu.flywheel.api.material.Material;
 import com.jozufozu.flywheel.api.material.Transparency;
 import com.jozufozu.flywheel.api.material.WriteMask;
+import com.jozufozu.flywheel.backend.gl.GlTextureUnit;
 import com.mojang.blaze3d.platform.GlStateManager;
 import com.mojang.blaze3d.systems.RenderSystem;
 
+import net.minecraft.client.Minecraft;
+import net.minecraft.client.renderer.texture.AbstractTexture;
+
 public final class MaterialRenderState {
 	public static final Comparator<Material> COMPARATOR = Comparator.comparing(Material::texture)
 			.thenComparing(Material::blur)
@@ -25,6 +29,7 @@ public final class MaterialRenderState {
 	}
 
 	public static void setup(Material material) {
+		setupTexture(material);
 		setupBackfaceCulling(material.backfaceCulling());
 		setupPolygonOffset(material.polygonOffset());
 		setupDepthTest(material.depthTest());
@@ -32,6 +37,17 @@ public final class MaterialRenderState {
 		setupWriteMask(material.writeMask());
 	}
 
+	private static void setupTexture(Material material) {
+		GlTextureUnit.T0.makeActive();
+		AbstractTexture texture = Minecraft.getInstance()
+				.getTextureManager()
+				.getTexture(material.texture());
+		texture.setFilter(material.blur(), material.mipmap());
+		var textureId = texture.getId();
+		RenderSystem.setShaderTexture(0, textureId);
+		RenderSystem.bindTexture(textureId);
+	}
+
 	private static void setupBackfaceCulling(boolean backfaceCulling) {
 		if (backfaceCulling) {
 			RenderSystem.enableCull();
@@ -125,6 +141,7 @@ public final class MaterialRenderState {
 	}
 
 	public static void reset() {
+		resetTexture();
 		resetBackfaceCulling();
 		resetPolygonOffset();
 		resetDepthTest();
@@ -132,6 +149,11 @@ public final class MaterialRenderState {
 		resetWriteMask();
 	}
 
+	private static void resetTexture() {
+		GlTextureUnit.T0.makeActive();
+		RenderSystem.setShaderTexture(0, 0);
+	}
+
 	private static void resetBackfaceCulling() {
 		RenderSystem.enableCull();
 	}

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

@@ -26,7 +26,7 @@ import com.jozufozu.flywheel.api.model.Model;
 import com.jozufozu.flywheel.backend.compile.IndirectPrograms;
 import com.jozufozu.flywheel.backend.engine.MaterialRenderState;
 import com.jozufozu.flywheel.backend.engine.textures.TextureBinder;
-import com.jozufozu.flywheel.backend.engine.textures.TexturesImpl;
+import com.jozufozu.flywheel.backend.engine.textures.TextureSourceImpl;
 import com.jozufozu.flywheel.backend.engine.uniform.Uniforms;
 import com.jozufozu.flywheel.backend.gl.Driver;
 import com.jozufozu.flywheel.backend.gl.GlCompat;
@@ -185,7 +185,7 @@ public class IndirectCullingGroup<I extends Instance> {
 		hasNewDraws = true;
 	}
 
-	public void submit(RenderStage stage, TexturesImpl textures) {
+	public void submit(RenderStage stage, TextureSourceImpl textures) {
 		if (nothingToDo(stage)) {
 			return;
 		}

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

@@ -23,7 +23,8 @@ import com.jozufozu.flywheel.backend.engine.InstancerKey;
 import com.jozufozu.flywheel.backend.engine.InstancerStorage;
 import com.jozufozu.flywheel.backend.engine.MaterialRenderState;
 import com.jozufozu.flywheel.backend.engine.textures.TextureBinder;
-import com.jozufozu.flywheel.backend.engine.textures.TexturesImpl;
+import com.jozufozu.flywheel.backend.engine.textures.TextureSourceImpl;
+import com.jozufozu.flywheel.backend.gl.GlStateTracker;
 import com.jozufozu.flywheel.backend.gl.buffer.GlBuffer;
 import com.jozufozu.flywheel.backend.gl.buffer.GlBufferType;
 import com.jozufozu.flywheel.lib.context.ContextShaders;
@@ -39,7 +40,7 @@ import net.minecraft.client.resources.model.ModelBakery;
 public class IndirectDrawManager extends InstancerStorage<IndirectInstancer<?>> {
 	private final IndirectPrograms programs;
 	private final StagingBuffer stagingBuffer;
-	private final TexturesImpl textures = new TexturesImpl();
+	private final TextureSourceImpl textures = new TextureSourceImpl();
 	private final Map<GroupKey<?>, IndirectCullingGroup<?>> cullingGroups = new HashMap<>();
 	private final GlBuffer crumblingDrawBuffer = new GlBuffer();
 
@@ -71,9 +72,14 @@ public class IndirectDrawManager extends InstancerStorage<IndirectInstancer<?>>
 	}
 
 	public void renderStage(RenderStage stage) {
+		TextureBinder.bindLightAndOverlay();
+
 		for (var group : cullingGroups.values()) {
 			group.submit(stage, textures);
 		}
+
+		MaterialRenderState.reset();
+		TextureBinder.resetLightAndOverlay();
 	}
 
 	@Override
@@ -117,6 +123,9 @@ public class IndirectDrawManager extends InstancerStorage<IndirectInstancer<?>>
 			return;
 		}
 
+		try (var state = GlStateTracker.getRestoreState()) {
+			TextureBinder.bindLightAndOverlay();
+
 			var crumblingMaterial = SimpleMaterial.builder();
 
 			// Scratch memory for writing draw commands.
@@ -141,10 +150,9 @@ public class IndirectDrawManager extends InstancerStorage<IndirectInstancer<?>>
 
 							// Transform the material to be suited for crumbling.
 							CommonCrumbling.applyCrumblingProperties(crumblingMaterial, draw.material());
-						Contexts.CRUMBLING.get(progressEntry.getIntKey())
-								.prepare(crumblingMaterial, program, textures);
+							var context = Contexts.CRUMBLING.get(progressEntry.getIntKey());
 
-						// Set up gl state for the draw.
+							context.prepare(crumblingMaterial, program, textures);
 							MaterialRenderState.setup(crumblingMaterial);
 
 							// Upload the draw command.
@@ -160,8 +168,12 @@ public class IndirectDrawManager extends InstancerStorage<IndirectInstancer<?>>
 				}
 			}
 
+			MaterialRenderState.reset();
+			TextureBinder.resetLightAndOverlay();
+
 			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<>();

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

@@ -10,7 +10,6 @@ import com.jozufozu.flywheel.backend.compile.IndirectPrograms;
 import com.jozufozu.flywheel.backend.engine.AbstractEngine;
 import com.jozufozu.flywheel.backend.engine.AbstractInstancer;
 import com.jozufozu.flywheel.backend.engine.InstancerStorage;
-import com.jozufozu.flywheel.backend.engine.MaterialRenderState;
 import com.jozufozu.flywheel.backend.engine.uniform.Uniforms;
 import com.jozufozu.flywheel.backend.gl.GlStateTracker;
 import com.jozufozu.flywheel.lib.task.Flag;
@@ -55,8 +54,6 @@ public class IndirectEngine extends AbstractEngine {
 
 		try (var restoreState = GlStateTracker.getRestoreState()) {
 			drawManager.renderStage(stage);
-
-			MaterialRenderState.reset();
 		}
 	}
 
@@ -64,11 +61,7 @@ public class IndirectEngine extends AbstractEngine {
 	public void renderCrumbling(TaskExecutor executor, RenderContext context, List<CrumblingBlock> crumblingBlocks) {
 		executor.syncUntil(flushFlag::isRaised);
 
-		try (var restoreState = GlStateTracker.getRestoreState()) {
 		drawManager.renderCrumbling(crumblingBlocks);
-
-			MaterialRenderState.reset();
-		}
 	}
 
 	@Override

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

@@ -6,7 +6,7 @@ import java.util.List;
 import java.util.Map;
 
 import com.jozufozu.flywheel.api.backend.Engine;
-import com.jozufozu.flywheel.api.context.Textures;
+import com.jozufozu.flywheel.api.context.TextureSource;
 import com.jozufozu.flywheel.api.instance.Instance;
 import com.jozufozu.flywheel.backend.compile.InstancingPrograms;
 import com.jozufozu.flywheel.backend.engine.CommonCrumbling;
@@ -24,7 +24,7 @@ import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
 import net.minecraft.client.resources.model.ModelBakery;
 
 public class InstancedCrumbling {
-	public static void render(List<Engine.CrumblingBlock> crumblingBlocks, InstancingPrograms programs, Textures textures) {
+	public static void render(List<Engine.CrumblingBlock> crumblingBlocks, InstancingPrograms programs, TextureSource textureSource) {
 		// Sort draw calls into buckets, so we don't have to do as many shader binds.
 		var byShaderState = doCrumblingSort(crumblingBlocks);
 
@@ -35,6 +35,8 @@ public class InstancedCrumbling {
 		var crumblingMaterial = SimpleMaterial.builder();
 
 		try (var state = GlStateTracker.getRestoreState()) {
+			TextureBinder.bindLightAndOverlay();
+
 			for (var shaderStateEntry : byShaderState.entrySet()) {
 				var byProgress = shaderStateEntry.getValue();
 
@@ -61,8 +63,8 @@ public class InstancedCrumbling {
 						continue;
 					}
 
-					Contexts.CRUMBLING.get(progressEntry.getIntKey())
-							.prepare(crumblingMaterial, program, textures);
+					var context = Contexts.CRUMBLING.get(progressEntry.getIntKey());
+					context.prepare(crumblingMaterial, program, textureSource);
 
 					drawCalls.forEach(Runnable::run);
 
@@ -71,6 +73,7 @@ public class InstancedCrumbling {
 			}
 
 			MaterialRenderState.reset();
+			TextureBinder.resetLightAndOverlay();
 		}
 	}
 

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

@@ -17,7 +17,7 @@ import com.jozufozu.flywheel.backend.engine.InstancerStorage;
 import com.jozufozu.flywheel.backend.engine.MaterialEncoder;
 import com.jozufozu.flywheel.backend.engine.MaterialRenderState;
 import com.jozufozu.flywheel.backend.engine.textures.TextureBinder;
-import com.jozufozu.flywheel.backend.engine.textures.TexturesImpl;
+import com.jozufozu.flywheel.backend.engine.textures.TextureSourceImpl;
 import com.jozufozu.flywheel.backend.engine.uniform.Uniforms;
 import com.jozufozu.flywheel.backend.gl.GlStateTracker;
 import com.jozufozu.flywheel.backend.gl.shader.GlProgram;
@@ -27,7 +27,7 @@ import com.jozufozu.flywheel.lib.task.SyncedPlan;
 
 public class InstancingEngine extends AbstractEngine {
 	private final InstancingPrograms programs;
-	private final TexturesImpl textures = new TexturesImpl();
+	private final TextureSourceImpl textures = new TextureSourceImpl();
 	private final InstancedDrawManager drawManager = new InstancedDrawManager();
 	private final Flag flushFlag = new NamedFlag("flushed");
 
@@ -88,6 +88,8 @@ public class InstancingEngine extends AbstractEngine {
 	}
 
 	private void render(InstancedDrawManager.DrawSet drawSet) {
+		TextureBinder.bindLightAndOverlay();
+
 		for (var entry : drawSet) {
 			var shader = entry.getKey();
 			var drawCalls = entry.getValue();
@@ -98,16 +100,17 @@ public class InstancingEngine extends AbstractEngine {
 				continue;
 			}
 
-			var program = programs.get(shader.instanceType(), shader.context()
-					.contextShader());
+			var context = shader.context();
+			var material = shader.material();
+
+			var program = programs.get(shader.instanceType(), context.contextShader());
 			program.bind();
 
 			Uniforms.bindForDraw();
-			uploadMaterialUniform(program, shader.material());
+			uploadMaterialUniform(program, material);
 
-			shader.context()
-					.prepare(shader.material(), program, textures);
-			MaterialRenderState.setup(shader.material());
+			context.prepare(material, program, textures);
+			MaterialRenderState.setup(material);
 
 			for (var drawCall : drawCalls) {
 				drawCall.render();
@@ -116,6 +119,7 @@ public class InstancingEngine extends AbstractEngine {
 		}
 
 		MaterialRenderState.reset();
+		TextureBinder.resetLightAndOverlay();
 	}
 
 	public static void uploadMaterialUniform(GlProgram program, Material material) {

src/main/java/com/jozufozu/flywheel/backend/engine/textures/TextureBinder.java

@@ -2,16 +2,23 @@ package com.jozufozu.flywheel.backend.engine.textures;
 
 import static org.lwjgl.opengl.GL13.GL_TEXTURE0;
 
+import com.jozufozu.flywheel.backend.gl.GlTextureUnit;
 import com.mojang.blaze3d.systems.RenderSystem;
 
 import it.unimi.dsi.fastutil.ints.Int2IntArrayMap;
 import it.unimi.dsi.fastutil.ints.Int2IntMap;
+import net.minecraft.client.Minecraft;
 
 public class TextureBinder {
 	// TODO: some kind of cache eviction when the program changes
 	//  so we don't always reset and bind the light and overlay textures?
 	private static final Int2IntMap texturesToSamplerUnits = new Int2IntArrayMap();
-	private static int nextSamplerUnit = 0;
+	// 0 is reserved for diffuse
+	// 1 is overlay
+	// 2 is light
+	// 3..n are for whatever else the context needs
+	private static final int baseSamplerUnit = 3;
+	private static int nextSamplerUnit = baseSamplerUnit;
 
 	/**
 	 * Binds the given texture to the next available texture unit, returning the unit it was bound to.
@@ -29,7 +36,7 @@ public class TextureBinder {
 	}
 
 	public static void resetTextureBindings() {
-		nextSamplerUnit = 0;
+		nextSamplerUnit = baseSamplerUnit;
 
 		for (Int2IntMap.Entry entry : texturesToSamplerUnits.int2IntEntrySet()) {
 			RenderSystem.activeTexture(GL_TEXTURE0 + entry.getIntValue());
@@ -38,4 +45,26 @@ public class TextureBinder {
 
 		texturesToSamplerUnits.clear();
 	}
+
+	public static void bindLightAndOverlay() {
+		var gameRenderer = Minecraft.getInstance().gameRenderer;
+		gameRenderer.overlayTexture()
+				.setupOverlayColor();
+		gameRenderer.lightTexture()
+				.turnOnLightLayer();
+
+		GlTextureUnit.T1.makeActive();
+		RenderSystem.bindTexture(RenderSystem.getShaderTexture(1));
+		GlTextureUnit.T2.makeActive();
+		RenderSystem.bindTexture(RenderSystem.getShaderTexture(2));
+	}
+
+	public static void resetLightAndOverlay() {
+		var gameRenderer = Minecraft.getInstance().gameRenderer;
+
+		gameRenderer.overlayTexture()
+				.teardownOverlayColor();
+		gameRenderer.lightTexture()
+				.turnOffLightLayer();
+	}
 }

src/main/java/com/jozufozu/flywheel/backend/engine/textures/TextureSourceImpl.java

@@ -4,7 +4,7 @@ import java.util.HashMap;
 import java.util.Map;
 
 import com.jozufozu.flywheel.api.context.Texture;
-import com.jozufozu.flywheel.api.context.Textures;
+import com.jozufozu.flywheel.api.context.TextureSource;
 import com.jozufozu.flywheel.backend.mixin.LightTextureAccessor;
 import com.jozufozu.flywheel.backend.mixin.OverlayTextureAccessor;
 
@@ -12,13 +12,13 @@ import net.minecraft.client.Minecraft;
 import net.minecraft.client.renderer.texture.AbstractTexture;
 import net.minecraft.resources.ResourceLocation;
 
-public class TexturesImpl implements Textures {
+public class TextureSourceImpl implements TextureSource {
 
 	private final DirectTexture lightTexture;
 	private final DirectTexture overlayTexture;
 	private final Map<ResourceLocation, WrappedTexture> wrappers = new HashMap<>();
 
-	public TexturesImpl() {
+	public TextureSourceImpl() {
 		var gameRenderer = Minecraft.getInstance().gameRenderer;
 
 		this.lightTexture = new DirectTexture(((LightTextureAccessor) gameRenderer.lightTexture()).flywheel$texture()

src/main/java/com/jozufozu/flywheel/backend/gl/shader/GlProgram.java

@@ -2,11 +2,17 @@ package com.jozufozu.flywheel.backend.gl.shader;
 
 import static org.lwjgl.opengl.GL20.glDeleteProgram;
 import static org.lwjgl.opengl.GL20.glGetUniformLocation;
+import static org.lwjgl.opengl.GL20.glUniform1f;
 import static org.lwjgl.opengl.GL20.glUniform1i;
+import static org.lwjgl.opengl.GL20.glUniform2f;
+import static org.lwjgl.opengl.GL20.glUniform3f;
+import static org.lwjgl.opengl.GL20.glUniform4f;
+import static org.lwjgl.opengl.GL20.glUniformMatrix4fv;
 import static org.lwjgl.opengl.GL31.GL_INVALID_INDEX;
 import static org.lwjgl.opengl.GL31.glGetUniformBlockIndex;
 import static org.lwjgl.opengl.GL31.glUniformBlockBinding;
 
+import org.joml.Matrix4fc;
 import org.slf4j.Logger;
 
 import com.jozufozu.flywheel.api.context.Shader;
@@ -39,6 +45,8 @@ public class GlProgram extends GlObject implements Shader {
 
 	@Override
 	public void setTexture(String glslName, Texture texture) {
+		throwIfReserved(glslName);
+
 		if (!(texture instanceof IdentifiedTexture identified)) {
 			return;
 		}
@@ -55,6 +63,71 @@ public class GlProgram extends GlObject implements Shader {
 		glUniform1i(uniform, binding);
 	}
 
+	@Override
+	public void setFloat(String glslName, float value) {
+		throwIfReserved(glslName);
+
+		int uniform = getUniformLocation(glslName);
+
+		if (uniform < 0) {
+			return;
+		}
+
+		glUniform1f(uniform, value);
+	}
+
+	@Override
+	public void setVec2(String glslName, float x, float y) {
+		throwIfReserved(glslName);
+
+		int uniform = getUniformLocation(glslName);
+
+		if (uniform < 0) {
+			return;
+		}
+
+		glUniform2f(uniform, x, y);
+	}
+
+	@Override
+	public void setVec3(String glslName, float x, float y, float z) {
+		throwIfReserved(glslName);
+
+		int uniform = getUniformLocation(glslName);
+
+		if (uniform < 0) {
+			return;
+		}
+
+		glUniform3f(uniform, x, y, z);
+	}
+
+	@Override
+	public void setVec4(String glslName, float x, float y, float z, float w) {
+		throwIfReserved(glslName);
+
+		int uniform = getUniformLocation(glslName);
+
+		if (uniform < 0) {
+			return;
+		}
+
+		glUniform4f(uniform, x, y, z, w);
+	}
+
+	@Override
+	public void setMat4(String glslName, Matrix4fc matrix) {
+		throwIfReserved(glslName);
+
+		int uniform = getUniformLocation(glslName);
+
+		if (uniform < 0) {
+			return;
+		}
+
+		glUniformMatrix4fv(uniform, false, matrix.get(new float[16]));
+	}
+
 	/**
 	 * Retrieves the index of the uniform with the given name.
 	 *
@@ -101,4 +174,18 @@ public class GlProgram extends GlObject implements Shader {
 	protected void deleteInternal(int handle) {
 		glDeleteProgram(handle);
 	}
+
+	public static void throwIfReserved(String glslName) {
+		if (glslName.startsWith("flw_")) {
+			throw new IllegalArgumentException("Uniform names starting with flw_are reserved");
+		}
+
+		if (glslName.startsWith("_flw_")) {
+			throw new IllegalArgumentException("Uniform names starting with _flw_ are reserved for internal use");
+		}
+
+		if (glslName.startsWith("gl_")) {
+			throw new IllegalArgumentException("Uniform names cannot start with gl_");
+		}
+	}
 }

src/main/java/com/jozufozu/flywheel/lib/context/Contexts.java

@@ -1,10 +1,6 @@
 package com.jozufozu.flywheel.lib.context;
 
 import com.jozufozu.flywheel.api.context.Context;
-import com.jozufozu.flywheel.api.context.ContextShader;
-import com.jozufozu.flywheel.api.context.Shader;
-import com.jozufozu.flywheel.api.context.Textures;
-import com.jozufozu.flywheel.api.material.Material;
 
 import it.unimi.dsi.fastutil.ints.Int2ObjectArrayMap;
 import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
@@ -12,22 +8,8 @@ import it.unimi.dsi.fastutil.ints.Int2ObjectMaps;
 import net.minecraft.client.resources.model.ModelBakery;
 
 public final class Contexts {
-	public static final Context DEFAULT = new Context() {
-		@Override
-		public ContextShader contextShader() {
-			return ContextShaders.DEFAULT;
-		}
-
-		@Override
-		public void prepare(Material material, Shader shader, Textures textures) {
-			var texture = textures.byName(material.texture());
-			texture.filter(material.blur(), material.mipmap());
-			shader.setTexture("_flw_diffuseTex", texture);
-
-			shader.setTexture("_flw_overlayTex", textures.overlay());
-			shader.setTexture("_flw_lightTex", textures.light());
-		}
-	};
+	public static final Context DEFAULT = SimpleContext.builder(ContextShaders.DEFAULT)
+			.build();
 
 	public static final Int2ObjectMap<Context> CRUMBLING;
 
@@ -35,7 +17,12 @@ public final class Contexts {
 		var map = new Int2ObjectArrayMap<Context>();
 
 		for (int i = 0; i < ModelBakery.BREAKING_LOCATIONS.size(); i++) {
-			map.put(i, new Crumbling(i));
+			var crumblingLocation = ModelBakery.BREAKING_LOCATIONS.get(i);
+			map.put(i, SimpleContext.builder(ContextShaders.CRUMBLING)
+					.preparation((material, shader, textureSource) -> {
+						shader.setTexture("crumblingTex", textureSource.byName(crumblingLocation));
+					})
+					.build());
 		}
 
 		CRUMBLING = Int2ObjectMaps.unmodifiable(map);
@@ -43,21 +30,4 @@ public final class Contexts {
 
 	private Contexts() {
 	}
-
-	private record Crumbling(int stage) implements Context {
-		@Override
-		public ContextShader contextShader() {
-			return ContextShaders.CRUMBLING;
-		}
-
-		@Override
-		public void prepare(Material material, Shader shader, Textures textures) {
-			var texture = textures.byName(material.texture());
-			texture.filter(material.blur(), material.mipmap());
-			shader.setTexture("_flw_diffuseTex", texture);
-
-			var crumblingTexture = textures.byName(ModelBakery.BREAKING_LOCATIONS.get(stage));
-			shader.setTexture("_flw_crumblingTex", crumblingTexture);
-		}
-	}
 }

src/main/java/com/jozufozu/flywheel/lib/context/SimpleContext.java

@@ -0,0 +1,61 @@
+package com.jozufozu.flywheel.lib.context;
+
+import org.jetbrains.annotations.Nullable;
+
+import com.jozufozu.flywheel.api.context.Context;
+import com.jozufozu.flywheel.api.context.ContextShader;
+import com.jozufozu.flywheel.api.context.Shader;
+import com.jozufozu.flywheel.api.context.TextureSource;
+import com.jozufozu.flywheel.api.material.Material;
+
+public class SimpleContext implements Context {
+	private final ContextShader contextShader;
+	private final Preparation preparation;
+
+	public SimpleContext(ContextShader contextShader, Preparation preparation) {
+		this.contextShader = contextShader;
+		this.preparation = preparation;
+	}
+
+	public static Builder builder(ContextShader contextShader) {
+		return new Builder(contextShader);
+	}
+
+	@Override
+	public ContextShader contextShader() {
+		return contextShader;
+	}
+
+	@Override
+	public void prepare(Material material, Shader shader, TextureSource textureSource) {
+		preparation.prepare(material, shader, textureSource);
+	}
+
+	@FunctionalInterface
+	public interface Preparation {
+		void prepare(Material material, Shader shader, TextureSource textureSource);
+	}
+
+	public static class Builder {
+		private final ContextShader contextShader;
+		@Nullable
+		private Preparation preparation;
+
+		public Builder(ContextShader contextShader) {
+			this.contextShader = contextShader;
+		}
+
+		public Builder preparation(Preparation preparation) {
+			this.preparation = preparation;
+			return this;
+		}
+
+		public SimpleContext build() {
+			if (preparation == null) {
+				preparation = (material, shader, textureSource) -> {
+				};
+			}
+			return new SimpleContext(contextShader, preparation);
+		}
+	}
+}

src/main/resources/assets/flywheel/flywheel/context/crumbling.frag

@@ -1,4 +1,4 @@
-uniform sampler2D _flw_crumblingTex;
+uniform sampler2D crumblingTex;
 
 in vec2 crumblingTexCoord;
 
@@ -8,7 +8,7 @@ void flw_beginFragment() {
 }
 
 void flw_endFragment() {
-    crumblingSampleColor = texture(_flw_crumblingTex, crumblingTexCoord);
+    crumblingSampleColor = texture(crumblingTex, crumblingTexCoord);
 
     // Make the crumbling overlay transparent when the fragment color after the material shader is transparent.
     flw_fragColor.rgb = crumblingSampleColor.rgb;