From e86480de37425245c3b1bf2719e4d60a53996a57 Mon Sep 17 00:00:00 2001
From: Kevin Rajan <7121943+kvnloo@users.noreply.github.com>
Date: Tue, 1 Sep 2026 21:40:01 +0000
Subject: [PATCH] Key patched materials by injection source

three.js caches programs on onBeforeCompile.toString(). Wrappers from
patchOnBeforeCompile are identically worded, so later patches reused the
first compiled program. customProgramCacheKey now includes fn source.
---
 src/utils/shaderPatch.js      |  9 +++++++++
 src/utils/shaderPatch.test.js | 15 +++++++++++++++
 2 files changed, 24 insertions(+)
 create mode 100644 src/utils/shaderPatch.test.js

diff --git a/src/utils/shaderPatch.js b/src/utils/shaderPatch.js
index e6e52df..59b0cf4 100644
--- a/src/utils/shaderPatch.js
+++ b/src/utils/shaderPatch.js
@@ -8,10 +8,19 @@
  */
 export function patchOnBeforeCompile(material, fn) {
   const previous = material.onBeforeCompile;
+  const previousKey = material.customProgramCacheKey?.bind(material);
+  const injectionKey = Function.prototype.toString.call(fn);
   material.onBeforeCompile = function (shader, renderer) {
     if (previous) previous.call(this, shader, renderer);
     fn.call(this, shader, renderer);
   };
+  // three.js keys the program cache on onBeforeCompile.toString() by default.
+  // Nested wrappers are identically worded, so the closed-over injection never
+  // appears in that string and every patched material would share one program.
+  material.customProgramCacheKey = function () {
+    const prior = previousKey ? previousKey.call(this) : '';
+    return `${prior}\n${injectionKey}`;
+  };
   return material;
 }
 
diff --git a/src/utils/shaderPatch.test.js b/src/utils/shaderPatch.test.js
new file mode 100644
index 0000000..b2d2ac4
--- /dev/null
+++ b/src/utils/shaderPatch.test.js
@@ -0,0 +1,15 @@
+import assert from 'node:assert/strict';
+import test from 'node:test';
+import { patchOnBeforeCompile } from './shaderPatch.js';
+
+test('composed patches get distinct program cache keys', () => {
+  const material = {};
+  patchOnBeforeCompile(material, function icePatch() {});
+  const iceKey = material.customProgramCacheKey();
+  patchOnBeforeCompile(material, function firePatch() {});
+  const fireKey = material.customProgramCacheKey();
+  assert.notEqual(iceKey, fireKey);
+  assert.match(iceKey, /icePatch/);
+  assert.match(fireKey, /firePatch/);
+  assert.match(fireKey, /icePatch/);
+});