From dd7b867f8966dc9cf098958bab463b4ebe17b427 Mon Sep 17 00:00:00 2001
From: vis-prime <vishals1030@gmail.com>
Date: Tue, 5 Dec 2023 10:12:01 +0530
Subject: [PATCH] fix: shaderMaterial sync with drei

---
 src/core/shaderMaterial.ts | 93 ++++++++++++++++++--------------------
 1 file changed, 45 insertions(+), 48 deletions(-)

diff --git a/src/core/shaderMaterial.ts b/src/core/shaderMaterial.ts
index 02c0c76..cbbadce 100644
--- a/src/core/shaderMaterial.ts
+++ b/src/core/shaderMaterial.ts
@@ -1,61 +1,58 @@
 import * as THREE from 'three'
 
-type UniformValue =
-  | THREE.CubeTexture
-  | THREE.Texture
-  | Int32Array
-  | Float32Array
-  | THREE.Matrix4
-  | THREE.Matrix3
-  | THREE.Quaternion
-  | THREE.Vector4
-  | THREE.Vector3
-  | THREE.Vector2
-  | THREE.Color
-  | number
-  | boolean
-  | Array<any>
-  | null
-
-type UniformProps = { [name: string]: UniformValue }
-
-type ShaderMaterialInstance<TProps extends UniformProps> = THREE.ShaderMaterial & TProps
-
-type ShaderMaterialParameters<TProps extends UniformProps> = THREE.ShaderMaterialParameters & Partial<TProps>
-
-type ShaderMaterial<TProps extends UniformProps> = (new (
-  parameters?: ShaderMaterialParameters<TProps>
-) => ShaderMaterialInstance<TProps>) & { key: string }
-
-export function shaderMaterial<TProps extends UniformProps>(
-  uniforms: TProps,
+export function shaderMaterial(
+  uniforms: {
+    [name: string]:
+      | THREE.CubeTexture
+      | THREE.Texture
+      | Int32Array
+      | Float32Array
+      | THREE.Matrix4
+      | THREE.Matrix3
+      | THREE.Quaternion
+      | THREE.Vector4
+      | THREE.Vector3
+      | THREE.Vector2
+      | THREE.Color
+      | number
+      | boolean
+      | Array<any>
+      | null
+  },
   vertexShader: string,
   fragmentShader: string,
-  onInit?: (material: ShaderMaterialInstance<TProps>) => void
+  onInit?: (material?: THREE.ShaderMaterial) => void
 ) {
-  const entries = Object.entries(uniforms)
-  const uniformDefs = Object.fromEntries(entries.map(([name, value]) => [name, { value }])) as {
-    [K in keyof TProps]: { value: TProps[K] }
-  }
-
-  class Material extends THREE.ShaderMaterial {
-    static key = THREE.MathUtils.generateUUID()
-
-    constructor(parameters?: ShaderMaterialParameters<TProps>) {
-      super({ uniforms: uniformDefs, vertexShader, fragmentShader })
-
-      for (const [name] of entries) {
+  const material = class extends THREE.ShaderMaterial {
+    public key: string = ''
+    constructor(parameters = {}) {
+      const entries = Object.entries(uniforms)
+      // Create uniforms and shaders
+      super({
+        uniforms: entries.reduce((acc, [name, value]) => {
+          const uniform = THREE.UniformsUtils.clone({ [name]: { value } })
+          return {
+            ...acc,
+            ...uniform,
+          }
+        }, {}),
+        vertexShader,
+        fragmentShader,
+      })
+      // Create getter/setters
+      entries.forEach(([name]) =>
         Object.defineProperty(this, name, {
           get: () => this.uniforms[name].value,
           set: (v) => (this.uniforms[name].value = v),
         })
-      }
+      )
 
+      // Assign parameters, this might include uniforms
       Object.assign(this, parameters)
-
-      onInit?.(this as unknown as ShaderMaterialInstance<TProps>)
+      // Call onInit
+      if (onInit) onInit(this)
     }
-  }
-
-  return Material as ShaderMaterial<TProps>
+  } as unknown as typeof THREE.ShaderMaterial & { key: string }
+  material.key = THREE.MathUtils.generateUUID()
+  return material
 }