fix(MeshRefractionMaterial): minimize shader bandwidth (#1303)

.storybook/stories/MeshRefractionMaterial.stories.tsx

@@ -0,0 +1,113 @@
+import * as React from 'react'
+import * as THREE from 'three'
+import { useLoader } from '@react-three/fiber'
+import { RGBELoader } from 'three-stdlib'
+import { Setup } from '../Setup'
+import {
+  MeshRefractionMaterial,
+  useGLTF,
+  Caustics,
+  CubeCamera,
+  Environment,
+  OrbitControls,
+  RandomizedLight,
+  AccumulativeShadows,
+  MeshTransmissionMaterial,
+} from '../../src'
+
+export default {
+  title: 'Shaders/MeshRefractionMaterial',
+  component: MeshRefractionMaterial,
+  decorators: [
+    (storyFn) => (
+      <Setup cameraFov={45} cameraPosition={new THREE.Vector3(-5, 0.5, 5)}>
+        {storyFn()}
+      </Setup>
+    ),
+  ],
+}
+
+function Diamond(props: any) {
+  const ref = React.useRef()
+  const { nodes } = useGLTF('/dflat.glb') as any
+  // Use a custom envmap/scene-backdrop for the diamond material
+  // This way we can have a clear BG while cube-cam can still film other objects
+  const texture = useLoader(
+    RGBELoader,
+    'https://dl.polyhaven.org/file/ph-assets/HDRIs/hdr/1k/aerodynamics_workshop_1k.hdr'
+  )
+  return (
+    <CubeCamera resolution={256} frames={1} envMap={texture}>
+      {(texture) => (
+        <Caustics
+          // @ts-ignore
+          backfaces
+          color="white"
+          position={[0, -0.5, 0]}
+          lightSource={[5, 5, -10]}
+          worldRadius={0.1}
+          ior={1.8}
+          backfaceIor={1.1}
+          intensity={0.1}
+        >
+          <mesh castShadow ref={ref} geometry={nodes.Diamond_1_0.geometry} {...props}>
+            <MeshRefractionMaterial
+              envMap={texture}
+              bounces={3}
+              aberrationStrength={0.01}
+              ior={2.75}
+              fresnel={1}
+              fastChroma
+              toneMapped={false}
+            />
+          </mesh>
+        </Caustics>
+      )}
+    </CubeCamera>
+  )
+}
+
+export const RefractionSt = () => (
+  <React.Suspense fallback={null}>
+    <color attach="background" args={['#f0f0f0']} />
+    <ambientLight intensity={0.5} />
+    <spotLight position={[5, 5, -10]} angle={0.15} penumbra={1} />
+    <pointLight position={[-10, -10, -10]} />
+    <Diamond rotation={[0, 0, 0.715]} position={[0, -0.175 + 0.5, 0]} />
+    {/* @ts-ignore */}
+    <Caustics
+      color="#FF8F20"
+      position={[0, -0.5, 0]}
+      lightSource={[5, 5, -10]}
+      worldRadius={0.003}
+      ior={1.16}
+      intensity={0.004}
+    >
+      <mesh castShadow receiveShadow position={[-2, 0.5, -1]} scale={0.5}>
+        <sphereGeometry args={[1, 64, 64]} />
+        {/* @ts-ignore */}
+        <MeshTransmissionMaterial resolution={1024} distortion={0.25} color="#FF8F20" thickness={1} anisotropy={1} />
+      </mesh>
+    </Caustics>
+    <mesh castShadow receiveShadow position={[1.75, 0.25, 1]} scale={0.75}>
+      <sphereGeometry args={[1, 64, 64]} />
+      <meshStandardMaterial color="hotpink" />
+    </mesh>
+    <AccumulativeShadows
+      temporal
+      frames={100}
+      color="orange"
+      colorBlend={2}
+      toneMapped={true}
+      alphaTest={0.8}
+      opacity={1}
+      scale={12}
+      position={[0, -0.5, 0]}
+    >
+      <RandomizedLight amount={8} radius={10} ambient={0.5} intensity={1} position={[5, 5, -10]} bias={0.001} />
+    </AccumulativeShadows>
+    <Environment files="https://dl.polyhaven.org/file/ph-assets/HDRIs/hdr/1k/aerodynamics_workshop_1k.hdr" />
+    <OrbitControls makeDefault autoRotate autoRotateSpeed={0.1} minPolarAngle={0} maxPolarAngle={Math.PI / 2} />
+  </React.Suspense>
+)
+RefractionSt.storyName = 'Default'

src/core/MeshRefractionMaterial.tsx

@@ -1,7 +1,7 @@
 import * as THREE from 'three'
 import * as React from 'react'
 import { useLayoutEffect, useMemo, useRef } from 'react'
-import { extend, ReactThreeFiber, useThree } from '@react-three/fiber'
+import { extend, ReactThreeFiber, useThree, useFrame } from '@react-three/fiber'
 import { MeshBVH, SAH } from 'three-mesh-bvh'
 import { MeshRefractionMaterial as MeshRefractionMaterial_ } from '../materials/MeshRefractionMaterial'
 
@@ -68,12 +68,18 @@ export function MeshRefractionMaterial({
     // Get the geometry of this materials parent
     const geometry = (material.current as any)?.__r3f?.parent?.geometry
     // Update the BVH
-    if (geometry)
-      (material.current as any).bvh.updateFrom(
+    if (geometry) {
+      ;(material.current as any).bvh.updateFrom(
         new MeshBVH(geometry.toNonIndexed(), { lazyGeneration: false, strategy: SAH })
       )
+    }
   }, [])
 
+  useFrame(({ camera }) => {
+    ;(material.current as any)!.viewMatrixInverse = camera.matrixWorld
+    ;(material.current as any)!.projectionMatrixInverse = camera.projectionMatrixInverse
+  })
+
   return (
     <meshRefractionMaterial
       // @ts-ignore

src/materials/MeshRefractionMaterial.tsx

@@ -16,41 +16,40 @@ export const MeshRefractionMaterial = shaderMaterial(
     bvh: new MeshBVHUniformStruct(),
     color: new THREE.Color('white'),
     resolution: new THREE.Vector2(),
+    viewMatrixInverse: new THREE.Matrix4(),
+    projectionMatrixInverse: new THREE.Matrix4(),
   },
   /*glsl*/ `
-  #ifndef USE_COLOR
-    uniform vec3 color;
-  #endif
+  uniform mat4 viewMatrixInverse;
+
   varying vec3 vWorldPosition;  
   varying vec3 vNormal;
-  varying mat4 projectionMatrixInv;
-  varying mat4 viewMatrixInv;
-  varying vec3 viewDirection;
-  varying mat4 vInstanceMatrix;
-  varying vec3 vColor;
-  
+  varying mat4 vModelMatrixInverse;
+
+  #ifdef USE_INSTANCING_COLOR
+    varying vec3 vInstanceColor;
+  #endif
+
   void main() {        
     vec4 transformedNormal = vec4(normal, 0.0);
     vec4 transformedPosition = vec4(position, 1.0);
     #ifdef USE_INSTANCING
-      vInstanceMatrix = instanceMatrix;
       transformedNormal = instanceMatrix * transformedNormal;
       transformedPosition = instanceMatrix * transformedPosition;
+    #endif
+
+    #ifdef USE_INSTANCING
+      vModelMatrixInverse = inverse(modelMatrix * instanceMatrix);
     #else
-      vInstanceMatrix = mat4(1.0);
+      vModelMatrixInverse = inverse(modelMatrix);
     #endif
 
-    vColor = color;
     #ifdef USE_INSTANCING_COLOR
-      vColor *= instanceColor.rgb;
+      vInstanceColor = instanceColor.rgb;
     #endif
-  
-    projectionMatrixInv = inverse(projectionMatrix);
-    viewMatrixInv = inverse(viewMatrix);
 
     vWorldPosition = (modelMatrix * transformedPosition).xyz;
-    vNormal = normalize((viewMatrixInv * vec4(normalMatrix * transformedNormal.xyz, 0.0)).xyz);
-    viewDirection = normalize(vWorldPosition - cameraPosition);
+    vNormal = normalize((viewMatrixInverse * vec4(normalMatrix * transformedNormal.xyz, 0.0)).xyz);
     gl_Position = projectionMatrix * viewMatrix * modelMatrix * transformedPosition;
   }`,
   /*glsl*/ `
@@ -59,6 +58,11 @@ export const MeshRefractionMaterial = shaderMaterial(
   precision highp usampler2D;
   varying vec3 vWorldPosition;
   varying vec3 vNormal;
+  varying mat4 vModelMatrixInverse;
+
+  #ifdef USE_INSTANCING_COLOR
+    varying vec3 vInstanceColor;
+  #endif
     
   #ifdef ENVMAP_TYPE_CUBEM
     uniform samplerCube envMap;
@@ -75,13 +79,10 @@ export const MeshRefractionMaterial = shaderMaterial(
   uniform vec2 resolution;
   uniform float fresnel;
   uniform mat4 modelMatrix;
-    
+  uniform mat4 projectionMatrixInverse;
+  uniform mat4 viewMatrixInverse;
   uniform float aberrationStrength;
-  varying mat4 projectionMatrixInv;
-  varying mat4 viewMatrixInv;
-  varying vec3 viewDirection;  
-  varying mat4 vInstanceMatrix;
-  varying vec3 vColor;
+  uniform vec3 color;
   
   float fresnelFunc(vec3 viewDirection, vec3 worldNormal) {
     return pow( 1.0 + dot( viewDirection, worldNormal), 10.0 );
@@ -130,33 +131,38 @@ export const MeshRefractionMaterial = shaderMaterial(
   #endif
   
   void main() {
-    mat4 modelMatrixInverse = inverse(modelMatrix * vInstanceMatrix);
     vec2 uv = gl_FragCoord.xy / resolution;
-    vec3 directionCamPerfect = (projectionMatrixInv * vec4(uv * 2.0 - 1.0, 0.0, 1.0)).xyz;
-    directionCamPerfect = (viewMatrixInv * vec4(directionCamPerfect, 0.0)).xyz;
+    vec3 directionCamPerfect = (projectionMatrixInverse * vec4(uv * 2.0 - 1.0, 0.0, 1.0)).xyz;
+    directionCamPerfect = (viewMatrixInverse * vec4(directionCamPerfect, 0.0)).xyz;
     directionCamPerfect = normalize(directionCamPerfect);
     vec3 normal = vNormal;
     vec3 rayOrigin = cameraPosition;
     vec3 rayDirection = normalize(vWorldPosition - cameraPosition);
     vec3 finalColor;
     #ifdef CHROMATIC_ABERRATIONS
-      vec3 rayDirectionG = totalInternalReflection(rayOrigin, rayDirection, normal, max(ior, 1.0), modelMatrixInverse);
+      vec3 rayDirectionG = totalInternalReflection(rayOrigin, rayDirection, normal, max(ior, 1.0), vModelMatrixInverse);
       #ifdef FAST_CHROMA 
         vec3 rayDirectionR = normalize(rayDirectionG + 1.0 * vec3(aberrationStrength / 2.0));
         vec3 rayDirectionB = normalize(rayDirectionG - 1.0 * vec3(aberrationStrength / 2.0));
       #else
-        vec3 rayDirectionR = totalInternalReflection(rayOrigin, rayDirection, normal, max(ior * (1.0 - aberrationStrength), 1.0), modelMatrixInverse);
-        vec3 rayDirectionB = totalInternalReflection(rayOrigin, rayDirection, normal, max(ior * (1.0 + aberrationStrength), 1.0), modelMatrixInverse);
+        vec3 rayDirectionR = totalInternalReflection(rayOrigin, rayDirection, normal, max(ior * (1.0 - aberrationStrength), 1.0), vModelMatrixInverse);
+        vec3 rayDirectionB = totalInternalReflection(rayOrigin, rayDirection, normal, max(ior * (1.0 + aberrationStrength), 1.0), vModelMatrixInverse);
       #endif
       float finalColorR = textureGradient(envMap, rayDirectionR, directionCamPerfect).r;
       float finalColorG = textureGradient(envMap, rayDirectionG, directionCamPerfect).g;
       float finalColorB = textureGradient(envMap, rayDirectionB, directionCamPerfect).b;
-      finalColor = vec3(finalColorR, finalColorG, finalColorB) * vColor;
+      finalColor = vec3(finalColorR, finalColorG, finalColorB);
     #else
-      rayDirection = totalInternalReflection(rayOrigin, rayDirection, normal, max(ior, 1.0), modelMatrixInverse);
+      rayDirection = totalInternalReflection(rayOrigin, rayDirection, normal, max(ior, 1.0), vModelMatrixInverse);
       finalColor = textureGradient(envMap, rayDirection, directionCamPerfect).rgb;    
-      finalColor *= vColor;
     #endif
+
+    finalColor *= color;
+    #ifdef USE_INSTANCING_COLOR
+      finalColor *= vInstanceColor;
+    #endif
+
+    vec3 viewDirection = normalize(vWorldPosition - cameraPosition);
     float nFresnel = fresnelFunc(viewDirection, normal) * fresnel;
     gl_FragColor = vec4(mix(finalColor, vec3(1.0), nFresnel), 1.0);      
     #include <tonemapping_fragment>