From b2c03a4069c4a904ae05c8f4bb5f81edc5f13c5f Mon Sep 17 00:00:00 2001
From: Jozufozu <jozsefaug@gmail.com>
Date: Sat, 8 Jan 2022 16:48:42 -0800
Subject: [PATCH] Conditional state restore and BufferUploader sanity check

 - Only restore state that has changed
 - Sanity check to fix error condition on reset
---
 .../flywheel/backend/gl/GlStateTracker.java   | 13 +++-
 .../flywheel/mixin/BufferUploaderMixin.java   | 29 ++++++++
 src/main/resources/flywheel.mixins.json       | 67 ++++++++++---------
 3 files changed, 73 insertions(+), 36 deletions(-)
 create mode 100644 src/main/java/com/jozufozu/flywheel/mixin/BufferUploaderMixin.java

diff --git a/src/main/java/com/jozufozu/flywheel/backend/gl/GlStateTracker.java b/src/main/java/com/jozufozu/flywheel/backend/gl/GlStateTracker.java
index f8ad22cd6..eec0a4cf8 100644
--- a/src/main/java/com/jozufozu/flywheel/backend/gl/GlStateTracker.java
+++ b/src/main/java/com/jozufozu/flywheel/backend/gl/GlStateTracker.java
@@ -45,11 +45,18 @@ public class GlStateTracker {
 			GlBufferType[] values = GlBufferType.values();
 
 			for (int i = 0; i < values.length; i++) {
-				GlStateManager._glBindBuffer(values[i].glEnum, buffers[i]);
+				if (buffers[i] != GlStateTracker.buffers[i]) {
+					GlStateManager._glBindBuffer(values[i].glEnum, buffers[i]);
+				}
 			}
 
-			GlStateManager._glBindVertexArray(vao);
-			GlStateManager._glUseProgram(program);
+			if (vao != GlStateTracker.vao) {
+				GlStateManager._glBindVertexArray(vao);
+			}
+
+			if (program != GlStateTracker.program) {
+				GlStateManager._glUseProgram(program);
+			}
 		}
 	}
 }
diff --git a/src/main/java/com/jozufozu/flywheel/mixin/BufferUploaderMixin.java b/src/main/java/com/jozufozu/flywheel/mixin/BufferUploaderMixin.java
new file mode 100644
index 000000000..7cd63581c
--- /dev/null
+++ b/src/main/java/com/jozufozu/flywheel/mixin/BufferUploaderMixin.java
@@ -0,0 +1,29 @@
+package com.jozufozu.flywheel.mixin;
+
+import javax.annotation.Nullable;
+
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+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.mojang.blaze3d.vertex.BufferUploader;
+import com.mojang.blaze3d.vertex.VertexFormat;
+
+@Mixin(BufferUploader.class)
+public class BufferUploaderMixin {
+
+	@Shadow
+	@Nullable
+	private static VertexFormat lastFormat;
+
+	@Inject(method = "reset", at = @At("HEAD"))
+	private static void stopBufferUploaderFromClearingBufferStateIfNothingIsBound(CallbackInfo ci) {
+		// Trust our tracker over BufferUploader's.
+		if (GlStateTracker.getVertexArray() == 0) {
+			lastFormat = null;
+		}
+	}
+}
diff --git a/src/main/resources/flywheel.mixins.json b/src/main/resources/flywheel.mixins.json
index 8b0e770d7..bbf0648db 100644
--- a/src/main/resources/flywheel.mixins.json
+++ b/src/main/resources/flywheel.mixins.json
@@ -1,36 +1,37 @@
 {
-    "required": true,
-    "minVersion": "0.8",
-    "package": "com.jozufozu.flywheel.mixin",
-    "compatibilityLevel": "JAVA_17",
-    "refmap": "flywheel.refmap.json",
-    "client": [
-        "BlockEntityTypeMixin",
-        "BufferBuilderMixin",
-        "CameraMixin",
-        "CancelEntityRenderMixin",
-        "ChunkRebuildHooksMixin",
-        "EntityTypeMixin",
-        "FixFabulousDepthMixin",
-        "FrustumMixin",
-        "GlStateManagerMixin",
-        "InstanceAddMixin",
-        "InstanceRemoveMixin",
-        "LevelRendererAccessor",
-        "LevelRendererMixin",
-        "PausedPartialTickAccessor",
-        "RenderTexturesMixin",
-        "RenderTypeMixin",
-        "ShaderCloseMixin",
-        "atlas.AtlasDataMixin",
-        "atlas.SheetDataAccessor",
-        "light.LightUpdateMixin",
-        "light.NetworkLightUpdateMixin",
-        "matrix.Matrix3fMixin",
-        "matrix.Matrix4fMixin",
-        "matrix.PoseStackMixin"
-    ],
-    "injectors": {
-        "defaultRequire": 1
+  "required": true,
+  "minVersion": "0.8",
+  "package": "com.jozufozu.flywheel.mixin",
+  "compatibilityLevel": "JAVA_17",
+  "refmap": "flywheel.refmap.json",
+  "client": [
+    "BlockEntityTypeMixin",
+    "BufferBuilderMixin",
+    "BufferUploaderMixin",
+    "CameraMixin",
+    "CancelEntityRenderMixin",
+    "ChunkRebuildHooksMixin",
+    "EntityTypeMixin",
+    "FixFabulousDepthMixin",
+    "FrustumMixin",
+    "GlStateManagerMixin",
+    "InstanceAddMixin",
+    "InstanceRemoveMixin",
+    "LevelRendererAccessor",
+    "LevelRendererMixin",
+    "PausedPartialTickAccessor",
+    "RenderTexturesMixin",
+    "RenderTypeMixin",
+    "ShaderCloseMixin",
+    "atlas.AtlasDataMixin",
+    "atlas.SheetDataAccessor",
+    "light.LightUpdateMixin",
+    "light.NetworkLightUpdateMixin",
+    "matrix.Matrix3fMixin",
+    "matrix.Matrix4fMixin",
+    "matrix.PoseStackMixin"
+  ],
+  "injectors": {
+    "defaultRequire": 1
   }
 }