GlStateTracker for better state restore

- Replaces both ShaderInstanceAccessor and BufferUploaderAccessor

src/main/java/com/jozufozu/flywheel/backend/gl/GlStateTracker.java

@@ -0,0 +1,55 @@
+package com.jozufozu.flywheel.backend.gl;
+
+import com.jozufozu.flywheel.backend.gl.buffer.GlBufferType;
+import com.mojang.blaze3d.platform.GlStateManager;
+
+/**
+ * Tracks bound buffers/vbos because GlStateManager doesn't do that for us.
+ */
+public class GlStateTracker {
+
+	private static final int[] buffers = new int[GlBufferType.values().length];
+	private static int vao;
+	private static int program;
+
+	public static int getBuffer(GlBufferType type) {
+		return buffers[type.ordinal()];
+	}
+
+	public static int getVertexArray() {
+		return vao;
+	}
+
+	public static int getProgram() {
+		return program;
+	}
+
+	public static void _setBuffer(GlBufferType type, int buffer) {
+		buffers[type.ordinal()] = buffer;
+	}
+
+	public static void _setProgram(int id) {
+		program = id;
+	}
+
+	public static void _setVertexArray(int id) {
+		vao = id;
+	}
+
+	public static State getRestoreState() {
+		return new State(buffers.clone(), vao, program);
+	}
+
+	public static record State(int[] buffers, int vao, int program) {
+		public void restore() {
+			GlBufferType[] values = GlBufferType.values();
+
+			for (int i = 0; i < values.length; i++) {
+				GlStateManager._glBindBuffer(values[i].glEnum, buffers[i]);
+			}
+
+			GlStateManager._glBindVertexArray(vao);
+			GlStateManager._glUseProgram(program);
+		}
+	}
+}

src/main/java/com/jozufozu/flywheel/backend/gl/GlVertexArray.java

@@ -4,7 +4,6 @@ import org.lwjgl.opengl.GL20;
 
 import com.jozufozu.flywheel.core.layout.BufferLayout;
 import com.jozufozu.flywheel.core.layout.LayoutItem;
-import com.jozufozu.flywheel.mixin.BufferUploaderAccessor;
 import com.mojang.blaze3d.platform.GlStateManager;
 
 public class GlVertexArray extends GlObject {
@@ -14,20 +13,14 @@ public class GlVertexArray extends GlObject {
 
 	public static void bind(int vao) {
 		GlStateManager._glBindVertexArray(vao);
-		BufferUploaderAccessor.flywheel$setLastVAO(vao);
 	}
 
 	public void bind() {
 		bind(handle());
 	}
 
-	public static int getBoundVertexArray() {
-		return BufferUploaderAccessor.flywheel$getLastVAO();
-	}
-
 	public static void unbind() {
 		GlStateManager._glBindVertexArray(0);
-		BufferUploaderAccessor.flywheel$setLastVAO(0);
 	}
 
 	public void enableArrays(int count) {

src/main/java/com/jozufozu/flywheel/backend/gl/buffer/GlBufferType.java

@@ -8,7 +8,7 @@ import org.lwjgl.opengl.GL40;
 import org.lwjgl.opengl.GL42;
 import org.lwjgl.opengl.GL43;
 
-import com.jozufozu.flywheel.mixin.BufferUploaderAccessor;
+import com.jozufozu.flywheel.backend.gl.GlStateTracker;
 import com.mojang.blaze3d.platform.GlStateManager;
 
 public enum GlBufferType {
@@ -33,29 +33,34 @@ public enum GlBufferType {
 		this.glEnum = glEnum;
 	}
 
-	public void bind(int buffer) {
-		GlStateManager._glBindBuffer(glEnum, buffer);
+	public static GlBufferType fromTarget(int pTarget) {
+		return switch (pTarget) {
+			case GL15C.GL_ARRAY_BUFFER -> ARRAY_BUFFER;
+			case GL15C.GL_ELEMENT_ARRAY_BUFFER -> ELEMENT_ARRAY_BUFFER;
+			case GL21.GL_PIXEL_PACK_BUFFER -> PIXEL_PACK_BUFFER;
+			case GL21.GL_PIXEL_UNPACK_BUFFER -> PIXEL_UNPACK_BUFFER;
+			case GL30.GL_TRANSFORM_FEEDBACK_BUFFER -> TRANSFORM_FEEDBACK_BUFFER;
+			case GL31.GL_UNIFORM_BUFFER -> UNIFORM_BUFFER;
+			case GL31.GL_TEXTURE_BUFFER -> TEXTURE_BUFFER;
+			case GL31.GL_COPY_READ_BUFFER -> COPY_READ_BUFFER;
+			case GL31.GL_COPY_WRITE_BUFFER -> COPY_WRITE_BUFFER;
+			case GL40.GL_DRAW_INDIRECT_BUFFER -> DRAW_INDIRECT_BUFFER;
+			case GL42.GL_ATOMIC_COUNTER_BUFFER -> ATOMIC_COUNTER_BUFFER;
+			case GL43.GL_DISPATCH_INDIRECT_BUFFER -> DISPATCH_INDIRECT_BUFFER;
+			case GL43.GL_SHADER_STORAGE_BUFFER -> SHADER_STORAGE_BUFFER;
+			default -> throw new IllegalArgumentException("Unknown target: " + pTarget);
+		};
+	}
 
-		switch (this.glEnum) {
-		case GL15C.GL_ELEMENT_ARRAY_BUFFER -> BufferUploaderAccessor.flywheel$setLastEBO(buffer);
-		case GL15C.GL_ARRAY_BUFFER -> BufferUploaderAccessor.flywheel$setLastVBO(buffer);
-		}
+    public void bind(int buffer) {
+		GlStateManager._glBindBuffer(glEnum, buffer);
 	}
 
 	public void unbind() {
 		GlStateManager._glBindBuffer(glEnum, 0);
-
-		switch (this.glEnum) {
-		case GL15C.GL_ELEMENT_ARRAY_BUFFER -> BufferUploaderAccessor.flywheel$setLastEBO(0);
-		case GL15C.GL_ARRAY_BUFFER -> BufferUploaderAccessor.flywheel$setLastVBO(0);
-		}
 	}
 
 	public int getBoundBuffer() {
-		return switch (this.glEnum) {
-		case GL15C.GL_ELEMENT_ARRAY_BUFFER -> BufferUploaderAccessor.flywheel$getLastEBO();
-		case GL15C.GL_ARRAY_BUFFER -> BufferUploaderAccessor.flywheel$getLastVBO();
-		default -> -1;
-		};
+		return GlStateTracker.getBuffer(this);
 	}
 }

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

@@ -11,7 +11,6 @@ import org.lwjgl.system.MemoryStack;
 
 import com.jozufozu.flywheel.backend.Backend;
 import com.jozufozu.flywheel.backend.gl.GlObject;
-import com.jozufozu.flywheel.mixin.ShaderInstanceAccessor;
 import com.mojang.blaze3d.shaders.ProgramManager;
 import com.mojang.math.Matrix4f;
 
@@ -29,14 +28,11 @@ public abstract class GlProgram extends GlObject {
 	}
 
 	public void bind() {
-		int handle = handle();
-		ProgramManager.glUseProgram(handle);
-		ShaderInstanceAccessor.flywheel$setLastProgramId(handle);
+		ProgramManager.glUseProgram(handle());
 	}
 
-	public void unbind() {
+	public static void unbind() {
 		ProgramManager.glUseProgram(0);
-		ShaderInstanceAccessor.flywheel$setLastProgramId(0);
 	}
 
 	/**

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

@@ -10,8 +10,6 @@ import javax.annotation.Nullable;
 
 import com.jozufozu.flywheel.api.MaterialGroup;
 import com.jozufozu.flywheel.backend.RenderLayer;
-import com.jozufozu.flywheel.backend.gl.GlVertexArray;
-import com.jozufozu.flywheel.backend.gl.buffer.GlBufferType;
 import com.jozufozu.flywheel.backend.instancing.Engine;
 import com.jozufozu.flywheel.backend.instancing.TaskEngine;
 import com.jozufozu.flywheel.core.WorldContext;
@@ -79,10 +77,6 @@ public class InstancingEngine<P extends WorldProgram> implements Engine {
 	 */
 	@Override
 	public void render(TaskEngine taskEngine, RenderLayerEvent event) {
-		int ebo = GlBufferType.ELEMENT_ARRAY_BUFFER.getBoundBuffer();
-		int vbo = GlBufferType.ARRAY_BUFFER.getBoundBuffer();
-		int vao = GlVertexArray.getBoundVertexArray();
-
 		double camX;
 		double camY;
 		double camZ;
@@ -102,10 +96,6 @@ public class InstancingEngine<P extends WorldProgram> implements Engine {
 		}
 
 		getGroupsToRender(event.getLayer()).forEach(group -> group.render(viewProjection, camX, camY, camZ));
-
-		GlBufferType.ELEMENT_ARRAY_BUFFER.bind(ebo);
-		GlBufferType.ARRAY_BUFFER.bind(vbo);
-		GlVertexArray.bind(vao);
 	}
 
 	private Stream<InstancedMaterialGroup<P>> getGroupsToRender(@Nullable RenderLayer layer) {

src/main/java/com/jozufozu/flywheel/core/shader/WorldProgram.java

@@ -28,7 +28,7 @@ public class WorldProgram extends ExtensibleGlProgram {
 
 		super.bind();
 		registerSamplers();
-		super.unbind();
+		unbind();
 	}
 
 	protected void registerSamplers() {

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

@@ -1,39 +0,0 @@
-package com.jozufozu.flywheel.mixin;
-
-import org.spongepowered.asm.mixin.Mixin;
-import org.spongepowered.asm.mixin.gen.Accessor;
-
-import com.mojang.blaze3d.vertex.BufferUploader;
-
-@Mixin(BufferUploader.class)
-public interface BufferUploaderAccessor {
-	@Accessor("lastVertexArrayObject")
-	static void flywheel$setLastVAO(int id) {
-		throw new AssertionError();
-	}
-
-	@Accessor("lastVertexBufferObject")
-	static void flywheel$setLastVBO(int id) {
-		throw new AssertionError();
-	}
-
-	@Accessor("lastIndexBufferObject")
-	static void flywheel$setLastEBO(int id) {
-		throw new AssertionError();
-	}
-
-	@Accessor("lastIndexBufferObject")
-	static int flywheel$getLastEBO() {
-		throw new AssertionError();
-	}
-
-	@Accessor("lastVertexBufferObject")
-	static int flywheel$getLastVBO() {
-		throw new AssertionError();
-	}
-
-	@Accessor("lastVertexArrayObject")
-	static int flywheel$getLastVAO() {
-		throw new AssertionError();
-	}
-}

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

@@ -0,0 +1,29 @@
+package com.jozufozu.flywheel.mixin;
+
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+
+import com.jozufozu.flywheel.backend.gl.GlStateTracker;
+import com.jozufozu.flywheel.backend.gl.buffer.GlBufferType;
+import com.mojang.blaze3d.platform.GlStateManager;
+
+@Mixin(GlStateManager.class)
+public class GlStateManagerMixin {
+
+	@Inject(method = "_glBindBuffer", at = @At("TAIL"))
+	private static void onBindBuffer(int pTarget, int pBuffer, CallbackInfo ci) {
+		GlStateTracker._setBuffer(GlBufferType.fromTarget(pTarget), pBuffer);
+	}
+
+	@Inject(method = "_glBindVertexArray", at = @At("TAIL"))
+	private static void onBindVertexArray(int pArray, CallbackInfo ci) {
+		GlStateTracker._setVertexArray(pArray);
+	}
+
+	@Inject(method = "_glUseProgram", at = @At("TAIL"))
+	private static void onUseProgram(int pProgram, CallbackInfo ci) {
+		GlStateTracker._setProgram(pProgram);
+	}
+}

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

@@ -8,6 +8,7 @@ import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 
 import com.jozufozu.flywheel.backend.Backend;
+import com.jozufozu.flywheel.backend.gl.GlStateTracker;
 import com.jozufozu.flywheel.backend.instancing.InstancedRenderDispatcher;
 import com.jozufozu.flywheel.core.crumbling.CrumblingRenderer;
 import com.jozufozu.flywheel.event.BeginFrameEvent;
@@ -57,7 +58,11 @@ public class LevelRendererMixin {
 
 		RenderBuffers renderBuffers = this.renderBuffers;
 
+		GlStateTracker.State restoreState = GlStateTracker.getRestoreState();
+
 		MinecraftForge.EVENT_BUS.post(new RenderLayerEvent(level, type, stack, renderBuffers, camX, camY, camZ));
+
+		restoreState.restore();
 	}
 
 	@Inject(at = @At("TAIL"), method = "allChanged")
@@ -76,7 +81,9 @@ public class LevelRendererMixin {
 
 		Vec3 cameraPos = info.getPosition();
 
+		GlStateTracker.State restoreState = GlStateTracker.getRestoreState();
 		CrumblingRenderer.renderBreaking(new RenderLayerEvent(level, null, stack, null, cameraPos.x, cameraPos.y, cameraPos.z));
+		restoreState.restore();
 	}
 
 	// Instancing

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

@@ -1,14 +0,0 @@
-package com.jozufozu.flywheel.mixin;
-
-import org.spongepowered.asm.mixin.Mixin;
-import org.spongepowered.asm.mixin.gen.Accessor;
-
-import net.minecraft.client.renderer.ShaderInstance;
-
-@Mixin(ShaderInstance.class)
-public interface ShaderInstanceAccessor {
-	@Accessor("lastProgramId")
-	static void flywheel$setLastProgramId(int id) {
-		throw new AssertionError();
-	}
-}

src/main/resources/flywheel.mixins.json

@@ -7,13 +7,13 @@
     "client": [
         "BlockEntityTypeMixin",
         "BufferBuilderMixin",
-        "BufferUploaderAccessor",
         "CameraMixin",
         "CancelEntityRenderMixin",
         "ChunkRebuildHooksMixin",
         "EntityTypeMixin",
         "FixFabulousDepthMixin",
         "FrustumMixin",
+        "GlStateManagerMixin",
         "InstanceAddMixin",
         "InstanceRemoveMixin",
         "LevelRendererAccessor",
@@ -22,7 +22,6 @@
         "RenderTexturesMixin",
         "RenderTypeMixin",
         "ShaderCloseMixin",
-        "ShaderInstanceAccessor",
         "atlas.AtlasDataMixin",
         "atlas.SheetDataAccessor",
         "light.LightUpdateMixin",