Good enough for government work

- Move crumbling tbn calculation back to the fragment shader.
- Pass the center of the crumbling block position as a uniform.
- Calculate how to flip the crumbling texture such that it tiles well at
  corners in model space.
- Move ugly crumbling code to utility class and improve bucketing.
- Cache uniform locations in GlProgram.

src/main/java/com/jozufozu/flywheel/api/backend/Engine.java

@@ -10,6 +10,7 @@ import com.jozufozu.flywheel.api.task.Plan;
 import com.jozufozu.flywheel.api.task.TaskExecutor;
 
 import net.minecraft.client.Camera;
+import net.minecraft.core.BlockPos;
 import net.minecraft.core.Vec3i;
 
 public interface Engine extends InstancerProvider {
@@ -29,12 +30,12 @@ public interface Engine extends InstancerProvider {
 
 	/**
 	 * Render the given instances as a crumbling overlay.
-	 * @param executor The task executor running the frame plan.
-	 * @param context The render context for this frame.
-	 * @param instances The instances to render.
-	 * @param progress The progress of the crumbling animation, i.e. which texture to use.
+	 *
+	 * @param executor        The task executor running the frame plan.
+	 * @param context         The render context for this frame.
+	 * @param crumblingBlocks The instances to render.
 	 */
-	void renderCrumblingInstances(TaskExecutor executor, RenderContext context, List<Instance> instances, int progress);
+	void renderCrumblingInstances(TaskExecutor executor, RenderContext context, List<CrumblingBlock> crumblingBlocks);
 
 	/**
 	 * Maintain the render origin to be within a certain distance from the camera in all directions,
@@ -55,4 +56,13 @@ public interface Engine extends InstancerProvider {
 	 * This engine will not be used again after this method is called.
 	 */
 	void delete();
+
+	/**
+	 * A block to be rendered as a crumbling overlay.
+	 * @param progress The progress of the crumbling animation in the range [0, 10).
+	 * @param pos The position of the block.
+	 * @param instances The instances associated with the BE at this position.
+	 */
+	record CrumblingBlock(int progress, BlockPos pos, List<Instance> instances) {
+	}
 }

src/main/java/com/jozufozu/flywheel/backend/engine/batching/BatchingEngine.java

@@ -4,7 +4,6 @@ import java.util.List;
 
 import com.jozufozu.flywheel.api.event.RenderContext;
 import com.jozufozu.flywheel.api.event.RenderStage;
-import com.jozufozu.flywheel.api.instance.Instance;
 import com.jozufozu.flywheel.api.task.Plan;
 import com.jozufozu.flywheel.api.task.TaskExecutor;
 import com.jozufozu.flywheel.backend.engine.AbstractEngine;
@@ -45,7 +44,7 @@ public class BatchingEngine extends AbstractEngine {
 	}
 
 	@Override
-	public void renderCrumblingInstances(TaskExecutor executor, RenderContext context, List<Instance> instances, int progress) {
+	public void renderCrumblingInstances(TaskExecutor executor, RenderContext context, List<CrumblingBlock> crumblingBlocks) {
 		// TODO: implement
 	}
 

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

@@ -6,7 +6,6 @@ import org.lwjgl.opengl.GL32;
 
 import com.jozufozu.flywheel.api.event.RenderContext;
 import com.jozufozu.flywheel.api.event.RenderStage;
-import com.jozufozu.flywheel.api.instance.Instance;
 import com.jozufozu.flywheel.api.task.Plan;
 import com.jozufozu.flywheel.api.task.TaskExecutor;
 import com.jozufozu.flywheel.backend.engine.AbstractEngine;
@@ -62,7 +61,7 @@ public class IndirectEngine extends AbstractEngine {
     }
 
 	@Override
-	public void renderCrumblingInstances(TaskExecutor executor, RenderContext context, List<Instance> instances, int progress) {
+	public void renderCrumblingInstances(TaskExecutor executor, RenderContext context, List<CrumblingBlock> crumblingBlocks) {
 		// TODO: implement
 	}
 

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

@@ -0,0 +1,131 @@
+package com.jozufozu.flywheel.backend.engine.instancing;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.jetbrains.annotations.NotNull;
+import org.lwjgl.opengl.GL32;
+
+import com.jozufozu.flywheel.api.backend.Engine;
+import com.jozufozu.flywheel.api.instance.Instance;
+import com.jozufozu.flywheel.api.material.Material;
+import com.jozufozu.flywheel.backend.compile.InstancingPrograms;
+import com.jozufozu.flywheel.backend.engine.InstanceHandleImpl;
+import com.jozufozu.flywheel.backend.engine.UniformBuffer;
+import com.jozufozu.flywheel.gl.GlStateTracker;
+import com.jozufozu.flywheel.gl.GlTextureUnit;
+import com.jozufozu.flywheel.lib.context.Contexts;
+import com.mojang.blaze3d.systems.RenderSystem;
+
+import it.unimi.dsi.fastutil.ints.Int2ObjectArrayMap;
+import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
+import net.minecraft.client.resources.model.ModelBakery;
+import net.minecraft.core.BlockPos;
+
+public class InstancingCrumble {
+	@NotNull
+	public static Map<ShaderState, Int2ObjectMap<Map<BlockPos, List<Runnable>>>> doCrumblingSort(List<Engine.CrumblingBlock> instances) {
+		Map<ShaderState, Int2ObjectMap<Map<BlockPos, List<Runnable>>>> out = new HashMap<>();
+
+		for (Engine.CrumblingBlock triple : instances) {
+			int progress = triple.progress();
+
+			if (progress < 0 || progress >= ModelBakery.DESTROY_TYPES.size()) {
+				continue;
+			}
+
+			BlockPos position = triple.pos();
+
+			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;
+				}
+
+				List<DrawCall> draws = instancer.drawCalls();
+
+				draws.removeIf(DrawCall::isInvalid);
+
+				for (DrawCall draw : draws) {
+					out.computeIfAbsent(draw.shaderState, $ -> new Int2ObjectArrayMap<>())
+							.computeIfAbsent(progress, $ -> new HashMap<>())
+							.computeIfAbsent(position, $ -> new ArrayList<>())
+							.add(() -> draw.renderOne(impl));
+				}
+			}
+		}
+		return out;
+	}
+
+	public static void render(List<Engine.CrumblingBlock> crumblingBlocks, BlockPos renderOrigin) {
+		// Sort draw calls into buckets, so we don't have to do as many shader binds.
+		var byShaderState = doCrumblingSort(crumblingBlocks);
+
+		if (byShaderState.isEmpty()) {
+			return;
+		}
+
+		try (var state = GlStateTracker.getRestoreState()) {
+			for (var shaderStateEntry : byShaderState.entrySet()) {
+				var byProgress = shaderStateEntry.getValue();
+
+				if (byProgress.isEmpty()) {
+					continue;
+				}
+
+				ShaderState shader = shaderStateEntry.getKey();
+
+				var material = shader.material();
+
+				var program = InstancingPrograms.get()
+						.get(shader.vertexType(), shader.instanceType(), Contexts.CRUMBLING);
+				UniformBuffer.syncAndBind(program);
+
+				int crumblingBlockPosUniform = program.getUniformLocation("_flw_crumblingBlockPos");
+
+				InstancingEngine.uploadMaterialIDUniform(program, material);
+
+				int renderTex = getDiffuseTexture(material);
+
+				for (Int2ObjectMap.Entry<Map<BlockPos, List<Runnable>>> progressEntry : byProgress.int2ObjectEntrySet()) {
+					var byPos = progressEntry.getValue();
+
+					if (byPos.isEmpty()) {
+						continue;
+					}
+
+					var crumblingType = ModelBakery.DESTROY_TYPES.get(progressEntry.getIntKey());
+
+					crumblingType.setupRenderState();
+
+					RenderSystem.setShaderTexture(1, renderTex);
+					GlTextureUnit.T1.makeActive();
+					RenderSystem.bindTexture(renderTex);
+
+					for (var blockPosEntry : byPos.entrySet()) {
+						var center = blockPosEntry.getKey().getCenter();
+						GL32.glUniform3f(crumblingBlockPosUniform, (float) center.x - renderOrigin.getX(), (float) center.y - renderOrigin.getY(), (float) center.z - renderOrigin.getZ());
+
+						blockPosEntry.getValue().forEach(Runnable::run);
+					}
+				}
+			}
+		}
+	}
+
+	private static int getDiffuseTexture(Material material) {
+		material.setup();
+
+		int out = RenderSystem.getShaderTexture(0);
+
+		material.clear();
+		return out;
+	}
+}

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

@@ -1,27 +1,22 @@
 package com.jozufozu.flywheel.backend.engine.instancing;
 
-import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 
-import org.jetbrains.annotations.NotNull;
 import org.lwjgl.opengl.GL32;
 
-import com.jozufozu.flywheel.api.context.Context;
 import com.jozufozu.flywheel.api.event.RenderContext;
 import com.jozufozu.flywheel.api.event.RenderStage;
-import com.jozufozu.flywheel.api.instance.Instance;
+import com.jozufozu.flywheel.api.material.Material;
 import com.jozufozu.flywheel.api.task.Plan;
 import com.jozufozu.flywheel.api.task.TaskExecutor;
 import com.jozufozu.flywheel.backend.compile.InstancingPrograms;
 import com.jozufozu.flywheel.backend.engine.AbstractEngine;
 import com.jozufozu.flywheel.backend.engine.AbstractInstancer;
-import com.jozufozu.flywheel.backend.engine.InstanceHandleImpl;
 import com.jozufozu.flywheel.backend.engine.InstancerStorage;
 import com.jozufozu.flywheel.backend.engine.UniformBuffer;
 import com.jozufozu.flywheel.gl.GlStateTracker;
 import com.jozufozu.flywheel.gl.GlTextureUnit;
+import com.jozufozu.flywheel.gl.shader.GlProgram;
 import com.jozufozu.flywheel.lib.context.Contexts;
 import com.jozufozu.flywheel.lib.material.MaterialIndices;
 import com.jozufozu.flywheel.lib.task.Flag;
@@ -30,7 +25,6 @@ import com.jozufozu.flywheel.lib.task.SyncedPlan;
 import com.mojang.blaze3d.systems.RenderSystem;
 
 import net.minecraft.client.Minecraft;
-import net.minecraft.client.resources.model.ModelBakery;
 
 public class InstancingEngine extends AbstractEngine {
 	private final InstancedDrawManager drawManager = new InstancedDrawManager();
@@ -74,82 +68,25 @@ public class InstancingEngine extends AbstractEngine {
     }
 
 	@Override
-	public void renderCrumblingInstances(TaskExecutor executor, RenderContext context, List<Instance> instances, int progress) {
-		if (instances.isEmpty()) {
-			return;
-		}
-
-		if (progress < 0 || progress >= ModelBakery.DESTROY_TYPES.size()) {
+	public void renderCrumblingInstances(TaskExecutor executor, RenderContext context, List<CrumblingBlock> crumblingBlocks) {
+		if (crumblingBlocks.isEmpty()) {
 			return;
 		}
 
 		// Need to wait for flush before we can inspect instancer state.
 		executor.syncUntil(flushFlag::isRaised);
 
-		// Sort draw calls into buckets, so we don't have to do as many shader binds.
-		var drawMap = getDrawsForInstances(instances);
-
-		if (drawMap.isEmpty()) {
-			return;
-		}
-
-		try (var state = GlStateTracker.getRestoreState()) {
-			var crumblingType = ModelBakery.DESTROY_TYPES.get(progress);
-
-			for (var entry : drawMap.entrySet()) {
-				var shader = entry.getKey();
-
-				setup(shader, Contexts.CRUMBLING);
-
-				shader.material().setup();
-
-				int renderTex = RenderSystem.getShaderTexture(0);
-
-				shader.material().clear();
-
-				crumblingType.setupRenderState();
-
-				RenderSystem.setShaderTexture(1, renderTex);
-				GlTextureUnit.T1.makeActive();
-				RenderSystem.bindTexture(renderTex);
-
-				for (Runnable draw : entry.getValue()) {
-					draw.run();
-				}
-			}
-		}
+		InstancingCrumble.render(crumblingBlocks, this.renderOrigin);
 	}
 
-	/**
-	 * Get all draw calls for the given instances, grouped by shader state.
-	 * @param instances The instances to draw.
-	 * @return A mapping of shader states to many runnable draw calls.
-	 */
-	@NotNull
-	private Map<ShaderState, List<Runnable>> getDrawsForInstances(List<Instance> instances) {
-		Map<ShaderState, List<Runnable>> out = new HashMap<>();
+	@Override
+	protected InstancerStorage<? extends AbstractInstancer<?>> getStorage() {
+		return drawManager;
+	}
 
-		for (Instance instance : 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;
-			}
-
-			List<DrawCall> draws = instancer.drawCalls();
-
-			draws.removeIf(DrawCall::isInvalid);
-
-			for (DrawCall draw : draws) {
-				out.computeIfAbsent(draw.shaderState, $ -> new ArrayList<>())
-						.add(() -> draw.renderOne(impl));
-			}
-		}
-		return out;
+	@Override
+	public void delete() {
+		drawManager.invalidate();
 	}
 
 	private void setup() {
@@ -174,7 +111,11 @@ public class InstancingEngine extends AbstractEngine {
 				continue;
 			}
 
-			setup(shader, Contexts.WORLD);
+			var program = InstancingPrograms.get()
+					.get(shader.vertexType(), shader.instanceType(), Contexts.WORLD);
+			UniformBuffer.syncAndBind(program);
+
+			uploadMaterialIDUniform(program, shader.material());
 
 			shader.material().setup();
 
@@ -186,28 +127,10 @@ public class InstancingEngine extends AbstractEngine {
 		}
 	}
 
-	private void setup(ShaderState desc, Context context) {
-		var material = desc.material();
-		var vertexType = desc.vertexType();
-		var instanceType = desc.instanceType();
-
-		var program = InstancingPrograms.get()
-				.get(vertexType, instanceType, context);
-		UniformBuffer.syncAndBind(program);
-
-		var uniformLocation = program.getUniformLocation("_flw_materialID_instancing");
-		var vertexID = MaterialIndices.getVertexShaderIndex(material);
-		var fragmentID = MaterialIndices.getFragmentShaderIndex(material);
-		GL32.glUniform2ui(uniformLocation, vertexID, fragmentID);
-	}
-
-	@Override
-	protected InstancerStorage<? extends AbstractInstancer<?>> getStorage() {
-		return drawManager;
-	}
-
-	@Override
-	public void delete() {
-		drawManager.invalidate();
+	public static void uploadMaterialIDUniform(GlProgram program, Material material) {
+		int materialIDUniform = program.getUniformLocation("_flw_materialID_instancing");
+		int vertexID = MaterialIndices.getVertexShaderIndex(material);
+		int fragmentID = MaterialIndices.getFragmentShaderIndex(material);
+		GL32.glUniform2ui(materialIDUniform, vertexID, fragmentID);
 	}
 }

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

@@ -13,9 +13,14 @@ import com.jozufozu.flywheel.gl.GlObject;
 import com.mojang.blaze3d.shaders.ProgramManager;
 import com.mojang.logging.LogUtils;
 
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
+
 public class GlProgram extends GlObject {
 	private static final Logger LOGGER = LogUtils.getLogger();
 
+	private final Object2IntMap<String> uniformLocationCache = new Object2IntOpenHashMap<>();
+
 	public GlProgram(int handle) {
 		handle(handle);
 	}
@@ -35,13 +40,14 @@ public class GlProgram extends GlObject {
 	 * @return The uniform's index
 	 */
 	public int getUniformLocation(String uniform) {
-		int index = glGetUniformLocation(this.handle(), uniform);
+		return uniformLocationCache.computeIfAbsent(uniform, s -> {
+			int index = glGetUniformLocation(this.handle(), uniform);
 
-		if (index < 0) {
-			LOGGER.debug("No active uniform '{}' exists. Could be unused.", uniform);
-		}
-
-		return index;
+			if (index < 0) {
+				LOGGER.debug("No active uniform '{}' exists. Could be unused.", uniform);
+			}
+			return index;
+		});
 	}
 
 	/**

src/main/java/com/jozufozu/flywheel/impl/visualization/VisualizationManagerImpl.java

@@ -10,7 +10,6 @@ import com.jozufozu.flywheel.api.backend.BackendManager;
 import com.jozufozu.flywheel.api.backend.Engine;
 import com.jozufozu.flywheel.api.event.RenderContext;
 import com.jozufozu.flywheel.api.event.RenderStage;
-import com.jozufozu.flywheel.api.instance.Instance;
 import com.jozufozu.flywheel.api.task.Plan;
 import com.jozufozu.flywheel.api.task.TaskExecutor;
 import com.jozufozu.flywheel.api.visual.DynamicVisual;
@@ -31,8 +30,6 @@ import com.jozufozu.flywheel.lib.task.NamedFlag;
 import com.jozufozu.flywheel.lib.task.RaisePlan;
 import com.jozufozu.flywheel.lib.util.LevelAttached;
 
-import it.unimi.dsi.fastutil.ints.Int2ObjectArrayMap;
-import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
 import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
 import net.minecraft.client.Minecraft;
 import net.minecraft.client.multiplayer.ClientLevel;
@@ -208,9 +205,13 @@ public class VisualizationManagerImpl implements VisualizationManager {
 	}
 
 	public void renderCrumbling(RenderContext context, Long2ObjectMap<SortedSet<BlockDestructionProgress>> destructionProgress) {
+		if (destructionProgress.isEmpty()) {
+			return;
+		}
+
 		taskExecutor.syncUntil(frameVisualsFlag::isRaised);
 
-		Int2ObjectMap<List<Instance>> progress2instances = new Int2ObjectArrayMap<>();
+		List<Engine.CrumblingBlock> crumblingBlocks = new ArrayList<>();
 
 		for (var entry : destructionProgress.long2ObjectEntrySet()) {
 			var set = entry.getValue();
@@ -233,18 +234,12 @@ public class VisualizationManagerImpl implements VisualizationManager {
 				continue;
 			}
 
-			// now for the fun part
+			var maxDestruction = set.last();
 
-			int progress = set.last()
-					.getProgress();
-
-			progress2instances.computeIfAbsent(progress, $ -> new ArrayList<>())
-					.addAll(instances);
+			crumblingBlocks.add(new Engine.CrumblingBlock(maxDestruction.getProgress(), maxDestruction.getPos(), instances));
 		}
 
-		for (var entry : progress2instances.int2ObjectEntrySet()) {
-			engine.renderCrumblingInstances(taskExecutor, context, entry.getValue(), entry.getIntKey());
-		}
+		engine.renderCrumblingInstances(taskExecutor, context, crumblingBlocks);
 	}
 
 	/**

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

@@ -7,17 +7,39 @@ layout (depth_greater) out float gl_FragDepth;
 #endif
 #endif
 
+uniform vec3 _flw_crumblingBlockPos;
+
 uniform sampler2D flw_diffuseTex;
 uniform sampler2D flw_crumblingTex;
 
-in vec2 flw_crumblingTexCoord;
+in vec2 _flw_crumblingFlip;
 
 out vec4 fragColor;
 
 vec4 flw_crumblingSampleColor;
 
+vec2 flattenedPos(vec3 pos, vec3 normal) {
+    pos = pos - _flw_crumblingBlockPos;
+
+    // https://community.khronos.org/t/52861
+    vec3 Q1 = dFdx(pos);
+    vec3 Q2 = dFdy(pos);
+    vec2 st1 = dFdx(flw_vertexTexCoord);
+    vec2 st2 = dFdy(flw_vertexTexCoord);
+
+    vec3 T = normalize(Q1*st2.t - Q2*st1.t);
+    vec3 B = normalize(-Q1*st2.s + Q2*st1.s);
+
+    mat3 tbn = mat3(T, B, normal);
+
+    // transpose is the same as inverse for orthonormal matrices
+    return ((transpose(tbn) * pos).xy + vec2(0.5)) * _flw_crumblingFlip;
+}
+
 void flw_initFragment() {
-    flw_crumblingSampleColor = texture(flw_crumblingTex, flw_crumblingTexCoord);
+    vec2 crumblingTexCoord = flattenedPos(flw_vertexPos.xyz, flw_vertexNormal);
+
+    flw_crumblingSampleColor = texture(flw_crumblingTex, crumblingTexCoord);
     flw_sampleColor = texture(flw_diffuseTex, flw_vertexTexCoord);
 
     // Let the other components modify the diffuse color as they normally would.

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

@@ -1,33 +1,86 @@
 #include "flywheel:api/vertex.glsl"
 #include "flywheel:util/fog.glsl"
 
-out vec2 flw_crumblingTexCoord;
+out vec2 _flw_crumblingFlip;
 
-vec3 tangent(vec3 normal) {
-    float sinYRot = -normal.x;
-    vec2 XZ = normal.xz;
-    float sqLength = dot(XZ, XZ);
-    if (sqLength > 0) {
-        sinYRot *= inversesqrt(sqLength);
-        sinYRot = clamp(sinYRot, -1, 1);
+const int DOWN = 0;
+const int UP = 1;
+const int NORTH = 2;
+const int SOUTH = 3;
+const int WEST = 4;
+const int EAST = 5;
+
+const vec2 FLIPS_BY_FACE[6] = vec2[](
+    vec2(1., -1.),
+    vec2(-1., -1.),
+    vec2(-1., -1.),
+    vec2(-1., -1.),
+    vec2(1., -1.),
+    vec2(1., -1.)
+);
+
+// based on net.minecraftforge.client.ForgeHooksClient.getNearestStable
+int getNearestFacing(vec3 normal) {
+    float maxAlignment = -2;
+    int face = 2;
+
+    vec3 alignment = vec3(
+    dot(normal, vec3(1., 0., 0.)),
+    dot(normal, vec3(0., 1., 0.)),
+    dot(normal, vec3(0., 0., 1.))
+    );
+
+    if (-alignment.y > maxAlignment) {
+        maxAlignment = -alignment.y;
+        face = DOWN;
+    }
+    if (alignment.y > maxAlignment) {
+        maxAlignment = alignment.y;
+        face = UP;
+    }
+    if (-alignment.z > maxAlignment) {
+        maxAlignment = -alignment.z;
+        face = NORTH;
+    }
+    if (alignment.z > maxAlignment) {
+        maxAlignment = alignment.z;
+        face = SOUTH;
+    }
+    if (alignment.x > maxAlignment) {
+        maxAlignment = -alignment.x;
+        face = WEST;
     }
 
-    return vec3(sqrt(1 - sinYRot * sinYRot) * (normal.z < 0 ? -1 : 1), 0, sinYRot);
+    return face;
 }
 
-vec2 flattenedPos(vec3 pos, vec3 normal) {
-    pos -= vec3(0.5);
+vec2 calculateFlip(vec3 normal) {
+    int face = getNearestFacing(normal);
+    return FLIPS_BY_FACE[face];
+}
 
-    vec3 tangent = tangent(normal);
-    vec3 bitangent = cross(tangent, normal);
-    mat3 tbn = mat3(tangent, bitangent, normal);
+// This is disgusting so if an issue comes up just throw this away and fix the branching version above.
+vec2 calculateFlipBranchless(vec3 normal) {
+    vec3 alignment = vec3(
+    dot(normal, vec3(1., 0., 0.)),
+    dot(normal, vec3(0., 1., 0.)),
+    dot(normal, vec3(0., 0., 1.))
+    );
 
-    // transpose is the same as inverse for orthonormal matrices
-    return (transpose(tbn) * pos).xy + vec2(0.5);
+    vec3 absAlignment = abs(alignment);
+
+    // x is the max alignment that would cause U to be -1.
+    // y is the max alignment that would cause U to be 1.
+    vec2 maxNegativeMaxPositive = max(vec2(absAlignment.z, alignment.y), vec2(-alignment.y, absAlignment.x));
+
+    bool flipU = maxNegativeMaxPositive.x > maxNegativeMaxPositive.y;
+
+    return vec2(mix(1., -1., flipU), -1.);
 }
 
 void flw_initVertex() {
-    flw_crumblingTexCoord = flattenedPos(flw_vertexPos.xyz, flw_vertexNormal);
+    // Calculate the flips in model space so that the crumbling effect doesn't have discontinuities.
+    _flw_crumblingFlip = calculateFlipBranchless(flw_vertexNormal);
 }
 
 void flw_contextVertex() {