From 8eeef43bdbf48290310917603aa87c0bfd1f7acb Mon Sep 17 00:00:00 2001
From: Martin Valigursky <mvaligursky@snapchat.com>
Date: Wed, 23 Apr 2025 16:43:46 +0100
Subject: [PATCH] WGSL clustered cookies

---
 .../shader-lib/chunks-wgsl/chunks-wgsl.js     |  4 ++--
 .../lit/frag/clusteredLightCookies.js         | 19 +++++++++++++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)
 create mode 100644 src/scene/shader-lib/chunks-wgsl/lit/frag/clusteredLightCookies.js

diff --git a/src/scene/shader-lib/chunks-wgsl/chunks-wgsl.js b/src/scene/shader-lib/chunks-wgsl/chunks-wgsl.js
index d947d7f7fbc..0c46dda00af 100644
--- a/src/scene/shader-lib/chunks-wgsl/chunks-wgsl.js
+++ b/src/scene/shader-lib/chunks-wgsl/chunks-wgsl.js
@@ -12,7 +12,7 @@ import clearCoatPS from './standard/frag/clearCoat.js';
 import clearCoatGlossPS from './standard/frag/clearCoatGloss.js';
 import clearCoatNormalPS from './standard/frag/clearCoatNormal.js';
 import clusteredLightUtilsPS from './lit/frag/clusteredLightUtils.js';
-// import clusteredLightCookiesPS from './lit/frag/clusteredLightCookies.js';
+import clusteredLightCookiesPS from './lit/frag/clusteredLightCookies.js';
 import clusteredLightShadowsPS from './lit/frag/clusteredLightShadows.js';
 import clusteredLightPS from './lit/frag/clusteredLight.js';
 import combinePS from './lit/frag/combine.js';
@@ -221,7 +221,7 @@ const shaderChunksWGSL = {
     clearCoatPS,
     clearCoatGlossPS,
     clearCoatNormalPS,
-    // clusteredLightCookiesPS,
+    clusteredLightCookiesPS,
     clusteredLightShadowsPS,
     clusteredLightUtilsPS,
     clusteredLightPS,
diff --git a/src/scene/shader-lib/chunks-wgsl/lit/frag/clusteredLightCookies.js b/src/scene/shader-lib/chunks-wgsl/lit/frag/clusteredLightCookies.js
new file mode 100644
index 00000000000..8c58de66122
--- /dev/null
+++ b/src/scene/shader-lib/chunks-wgsl/lit/frag/clusteredLightCookies.js
@@ -0,0 +1,19 @@
+export default /* wgsl */`
+fn _getCookieClustered(tex: texture_2d<f32>, texSampler: sampler, uv: vec2f, intensity: f32, cookieChannel: vec4f) -> vec3f {
+    let pixel: vec4f = mix(vec4f(1.0), textureSampleLevel(tex, texSampler, uv, 0.0), intensity);
+    let isRgb: bool = dot(cookieChannel.rgb, vec3f(1.0)) == 3.0;
+    return select(vec3f(dot(pixel, cookieChannel)), pixel.rgb, isRgb);
+}
+
+// getCookie2D for clustered lighting including channel selector
+fn getCookie2DClustered(tex: texture_2d<f32>, texSampler: sampler, transform: mat4x4f, worldPosition: vec3f, intensity: f32, cookieChannel: vec4f) -> vec3f {
+    let projPos: vec4f = transform * vec4f(worldPosition, 1.0);
+    return _getCookieClustered(tex, texSampler, projPos.xy / projPos.w, intensity, cookieChannel);
+}
+
+// getCookie for clustered omni light with the cookie texture being stored in the cookie atlas
+fn getCookieCubeClustered(tex: texture_2d<f32>, texSampler: sampler, dir: vec3f, intensity: f32, cookieChannel: vec4f, shadowTextureResolution: f32, shadowEdgePixels: f32, omniAtlasViewport: vec3f) -> vec3f {
+    let uv: vec2f = getCubemapAtlasCoordinates(omniAtlasViewport, shadowEdgePixels, shadowTextureResolution, dir);
+    return _getCookieClustered(tex, texSampler, uv, intensity, cookieChannel);
+}
+`;