DynamicEnvMap.tsx

import {
  Group,
  AmbientLight,
  CubeCamera,
  HalfFloatType,
  Mesh,
  MeshStandardMaterial,
  Texture,
  Vector3,
  WebGLCubeRenderTarget,
  MeshPhysicalMaterial,
  Object3D,
  MeshLambertMaterial,
  MeshBasicMaterial,
  ShaderMaterial,
  Material,
} from 'three'
import { useMemo, useLayoutEffect, useRef, useEffect, useDeferredValue, useState } from 'react'
import { useFrame, useThree } from '@react-three/fiber'

function canUseEnvMap(material: Material) {
  if (material instanceof MeshStandardMaterial || material instanceof MeshPhysicalMaterial || material instanceof MeshLambertMaterial || material instanceof MeshBasicMaterial) {
    return true
  }

  // Check for custom ShaderMaterial with an envMap uniform
  if (material instanceof ShaderMaterial) {
    return material.uniforms && 'envMap' in material.uniforms
  }

  return false
}

function canUseEnvMapIntensity(material: Material) {
  if (material instanceof MeshStandardMaterial || material instanceof MeshPhysicalMaterial) {
    return true
  }

  // Check for custom ShaderMaterial with an envMap uniform
  if (material instanceof ShaderMaterial) {
    return material.uniforms && 'envMapIntensity' in material.uniforms
  }

  return false
}

// the returned function will return true at p% of the time it gets called
const createFrequencyCheck = () => {
  // we start with a random value, this way if we create multiple components
  // with the frame update frequency (e.g. 0.5) in the same frame they won't
  // all simulatenously update till the end of the application, but smear
  // out over several frames, reducing the per-frame update pressure
  let sum = Math.random()
  return (p: number) => {
    sum += p
    if (sum >= 1) {
      sum = sum - Math.trunc(sum)
      return true
    }
    return false
  }
}

const defaultSettings = {
  // cube render target resolution
  resolution: 64,
  // % of frames we want to update
  updateFrequency: 1,
  // cube camera near
  near: 0.1,
  // cube camera far
  far: 25,
  // optional ambient light used during rendering, helpful for dark scenes
  ambientIntensity: 1,
  // optional ambient light color, tints reflections
  ambientTint: 'white',
}

export function useDynamicEnvMap(object: Object3D | Mesh, settings: Partial<typeof defaultSettings>) {
  // state & preparation
  const { gl, scene } = useThree(({ gl, scene }) => ({ gl, scene }))
  const frequencyCheck = useMemo(createFrequencyCheck, [])
  const worldPos = useMemo(() => new Vector3(), [])
  const mergedSettings = useMemo(() => Object.assign({}, defaultSettings, settings), [settings])

  const renderedOnce = useRef(false)
  const boundingSphereRef = useRef()
  const ambientLight = useMemo(() => new AmbientLight(mergedSettings.ambientTint, mergedSettings.ambientIntensity), [])

  const { renderTarget, cubeCamera } = useMemo(() => {
    const renderTarget = new WebGLCubeRenderTarget(mergedSettings.resolution)
    renderTarget.texture.type = HalfFloatType
    const cubeCamera = new CubeCamera(mergedSettings.near, mergedSettings.far, renderTarget)

    return { renderTarget, cubeCamera }
  }, [mergedSettings.resolution, mergedSettings.near, mergedSettings.far])

  // =============================================================================================

  // if the object is a mesh, we use the bounding sphere to get a better cam position
  useLayoutEffect(() => {
    if (!object) return
    if (object instanceof Mesh) {
      object.geometry.computeBoundingSphere()
      boundingSphereRef.current = object.geometry.boundingSphere
    }
  }, [object])

  // control the extra ambient light used during cube map rendering
  useLayoutEffect(() => {
    ambientLight.color.set(mergedSettings.ambientTint)
    ambientLight.intensity = mergedSettings.ambientIntensity
    return () => {
      // making sure the light is really removed when the component unmounts – this is mostly for hmr
      return () => scene.remove(ambientLight)
    }
  }, [ambientLight, mergedSettings.ambientTint, mergedSettings.ambientIntensity, scene])

  // make sure we render at least once (even if frequency is set to 0) and we stay reactive to settings changing
  useEffect(() => {
    renderedOnce.current = false
  }, [settings])

  // =============================================================================================
  // rendering

  useFrame(() => {
    if (!object) return

    if (mergedSettings.updateFrequency === 1 || !renderedOnce.current || (mergedSettings.updateFrequency > 0 && frequencyCheck(mergedSettings.updateFrequency))) {
      // hide the mesh before rendering
      object.visible = false

      // add an extra ambient light to make reflections more visible (more important for dark scenes)
      scene.add(ambientLight)

      // position camera (if it's a mesh we take the bounding sphere into account)
      object.getWorldPosition(worldPos)
      if (boundingSphereRef.current) worldPos.add(boundingSphereRef.current.center)
      cubeCamera.position.copy(worldPos)

      // render the new env map
      cubeCamera.update(gl, scene)

      // remove extra light
      scene.remove(ambientLight)

      // add mesh back
      object.visible = true

      // we enable the user to render the map just a single time
      renderedOnce.current = true
    }
  })

  return renderTarget
}

export function DynamicEnvMap(
  props: Partial<typeof defaultSettings> & {
    // sets envMapIntensity – only works on PBR materials
    intensity?: number
  },
) {
  const [group, setGroup] = useState(null)

  // =============================================================================================
  // effects (assigning env maps, swapping back to the original on unmount etc.)

  const renderTarget = useDynamicEnvMap(group, props)

  const originalEnvMap = useRef<Texture | null>(null)
  const originalEnvIntensity = useRef<number | null>(null)

  useLayoutEffect(() => {
    if (!group) return
    const parent = group.parent

    if (!(parent instanceof Mesh)) {
      throw new Error('<DynamicEnvMap> can only be used as a direct child of a mesh!')
    }

    if (!canUseEnvMap(parent.material)) {
      throw new Error('<DynamicEnvMap> can only be used with meshes whose material supports environment maps!')
    }

    // save the original env map
    originalEnvMap.current = parent.material.envMap

    // if this is a pbr material we also save the original env map intensity
    if (canUseEnvMapIntensity(parent.material)) {
      originalEnvIntensity.current = parent.material.envMapIntensity
    }

    return () => {
      // restore original env map and env map intensity if applicable
      parent.material.envMap = originalEnvMap.current
      if (originalEnvIntensity.current !== null) parent.material.envMapIntensity = originalEnvIntensity.current
    }
  }, [group])

  useLayoutEffect(() => {
    if (!group) return
    // update the env map to the new render target
    const parent = group.parent
    parent.material.envMap = renderTarget.texture
  }, [renderTarget, group])

  useLayoutEffect(() => {
    if (!group) return
    if (props.intensity === undefined) return

    // update the env map intensity if applicable
    const parent = group.parent
    if (!canUseEnvMapIntensity(parent.material)) {
      throw new Error(`<DynamicEnvMap>: the intensity prop can only be used with Materials that support 'envMapIntensity' like MeshStandardMaterial or MeshPhysicalMaterial`)
    }
    parent.material.envMapIntensity = props.intensity
  }, [props.intensity, group])

  return <group ref={setGroup} />
}