WebGPU Firework Particles in React Three Fiber with TSL Sprite Sheets

'use client'

import {
  type CreatorState,
  useFrame,
  useLocalNodes,
  useTexture,
  useThree,
  useUniforms,
} from '@react-three/fiber/webgpu'
import { type FC, useCallback, useEffect, useMemo, useRef } from 'react'
import {
  Fn,
  If,
  PI,
  Return,
  array,
  color,
  cos,
  float,
  hash,
  instanceIndex,
  instancedArray,
  mix,
  sin,
  smoothstep,
  spritesheetUV,
  texture,
  time,
  uv,
  vec2,
  vec3,
  vec4,
  vertexStage,
} from 'three/tsl'
import { NearestFilter } from 'three/webgpu'
import type { Node, UniformNode } from 'three/webgpu'
import * as THREE from 'three/webgpu'

export const DEFAULT_FPS = 60
export const DEFAULT_PARTICLE_COUNT = 2 ** 12
export const DEFAULT_BATCH_COUNT = 8
export const DEFAULT_BURST_DELAY_MIN = 300
export const DEFAULT_BURST_DELAY_MAX = 1000
const EXPLOSION_TRIGGER_LIFE = 1.15
const EXPLOSION_SCALE_MULTIPLIER = 1.6
const FIREWORK_PHASE = {
  launch: 0,
  explosion: 1,
} as const

const FIREWORKS_SPRITE_URL = new URL('./fireworks-sprites.webp', import.meta.url).href

const PARTICLE_PALETTE = [
  '#ffe071',
  '#ff9183',
  '#ff4b93',
  '#5f00aa',
  '#9a00a3',
  '#fff636',
  '#f6cc00',
  '#ffdf3a',
  '#f8d200',
  '#ffe81f',
  '#fff132',
  '#fffc5f',
  '#edc740',
  '#d5c8a3',
  '#ff9b91',
  '#ffa798',
  '#ff978b',
  '#ffbaa7',
  '#ff957a',
  '#ffa59d',
  '#ffc3a6',
  '#ffceb2',
  '#bf897e',
  '#ff9be2',
  '#ff85d4',
  '#ffb1ff',
  '#ff8ddc',
  '#ff6ecd',
  '#ff6fc7',
  '#ffb1f2',
  '#ff96e4',
  '#ca92a3',
  '#ffbfff',
  '#faa7ff',
  '#ffb6ff',
  '#ffc8ff',
  '#ffbcff',
  '#ffccff',
  '#e2c0ff',
  '#ffc9ff',
  '#ccb0d9',
  '#ffc3ff',
  '#ffb0ff',
  '#ffafff',
  '#ffb5ff',
  '#ff90ff',
  '#ffd5ff',
  '#ffa9ff',
  '#d6abda',
] as const

const PARTICLE_PALETTE_SIZE = PARTICLE_PALETTE.length
const particlePaletteNodes = array(PARTICLE_PALETTE.map((hex) => color(hex)))

const getParticleSeed = (index: number): number => {
  const value = Math.sin((index + 1) * 12.9898) * 43758.5453
  return value - Math.floor(value)
}

type Props = {
  batchCount?: number
  burstDelayMax?: number
  burstDelayMin?: number
  fps?: number
  particleCount?: number
}

export const FireworksParticles: FC<Props> = ({
  batchCount = DEFAULT_BATCH_COUNT,
  burstDelayMax = DEFAULT_BURST_DELAY_MAX,
  burstDelayMin = DEFAULT_BURST_DELAY_MIN,
  fps = DEFAULT_FPS,
  particleCount = DEFAULT_PARTICLE_COUNT,
}) => {
  const renderer = useThree((s) => s.renderer)
  const activeBatchRef = useRef(0)

  useUniforms({ uActiveBatch: 0, uDeltaTime: 0 }, 'fireworks')

  const spriteTexture = useTexture(FIREWORKS_SPRITE_URL, (textureValue) => {
    textureValue.generateMipmaps = false
    textureValue.minFilter = NearestFilter
    textureValue.magFilter = NearestFilter
  })

  const particleBuffers = useMemo(() => {
    // metaBuffer: x = batch index, y = deterministic seed, z = phase,
    // w = explosion scale per unit of remaining life.
    const metaValues = new Float32Array(particleCount * 4)

    for (let index = 0; index < particleCount; index += 1) {
      const metaOffset = index * 4

      metaValues[metaOffset] = Math.min(
        batchCount - 1,
        Math.floor((index / particleCount) * batchCount),
      )
      metaValues[metaOffset + 1] = getParticleSeed(index)
    }

    return {
      colorBuffer: instancedArray(particleCount, 'vec4'),
      metaBuffer: instancedArray(metaValues, 'vec4'),
      positionBuffer: instancedArray(particleCount, 'vec4'),
      velocityBuffer: instancedArray(particleCount, 'vec4'),
    }
  }, [batchCount, particleCount])

  const createSimulationNodes = useCallback(
    ({ uniforms }: CreatorState) => {
      const { colorBuffer, metaBuffer, positionBuffer, velocityBuffer } = particleBuffers
      const fireworkUniforms = uniforms.scope<{
        uActiveBatch: UniformNode<'float', number>
        uDeltaTime: UniformNode<'float', number>
      }>('fireworks')
      const activeBatch = fireworkUniforms.uActiveBatch
      const simulationDelta = fireworkUniforms.uDeltaTime
      let previousSimulationTime: number | null = null

      // TSL deltaTime only covers the latest render frame, not skipped compute updates.
      simulationDelta.onRenderUpdate(({ deltaTime: frameDelta, time: frameTime }) => {
        const value =
          previousSimulationTime === null ? frameDelta : frameTime - previousSimulationTime

        previousSimulationTime = frameTime
        return value
      })

      const launchPhaseNode = float(FIREWORK_PHASE.launch).toConst()
      const explosionPhaseNode = float(FIREWORK_PHASE.explosion).toConst()
      const explosionTriggerLifeNode = float(EXPLOSION_TRIGGER_LIFE).toConst()
      const explosionScaleMultiplierNode = float(EXPLOSION_SCALE_MULTIPLIER).toConst()

      const resetParticles = Fn(() => {
        positionBuffer.element(instanceIndex).assign(vec4(0, -100, 0, 0))
        velocityBuffer.element(instanceIndex).assign(vec4(0))
        metaBuffer.element(instanceIndex).z.assign(launchPhaseNode)
        metaBuffer.element(instanceIndex).w.assign(0)

        const paletteIndex = hash(instanceIndex.add(11)).mul(PARTICLE_PALETTE_SIZE).floor()
        const paletteColor = particlePaletteNodes.element(
          paletteIndex,
        ) as unknown as Node<'vec3'>
        const frameIndex = hash(instanceIndex.add(23)).mul(4).floor()

        colorBuffer.element(instanceIndex).assign(vec4(paletteColor, frameIndex))
      })().compute(particleCount)

      const spawnBatch = Fn(() => {
        const particleIndex = instanceIndex.toFloat()
        const meta = metaBuffer.element(instanceIndex)
        const batch = meta.x
        const position = positionBuffer.element(instanceIndex)
        const life = position.w
        const positionValue = position.xyz
        const velocityValue = velocityBuffer.element(instanceIndex)
        const seed = meta.y
        const phase = meta.z
        const scaleFadeRate = meta.w
        const colorValue = colorBuffer.element(instanceIndex)
        const frameIndex = colorValue.a

        If(batch.equal(activeBatch).and(life.lessThanEqual(0)), () => {
          const batchSeed = activeBatch.add(time.mul(0.1)).toVar()
          const particleTime = particleIndex.add(time).toVar()

          life.assign(mix(1.2, 2, hash(batchSeed.add(7.3))))

          const colorIndex = hash(particleTime.mul(3.7)).mul(PARTICLE_PALETTE_SIZE).floor()
          colorValue.xyz.assign(
            particlePaletteNodes.element(colorIndex) as unknown as Node<'vec3'>,
          )
          frameIndex.assign(hash(particleTime.add(seed)).mul(4).floor())

          const originX = mix(float(-2), float(2), hash(batchSeed.add(1.1)))
          const originY = mix(float(-6), float(-5), hash(batchSeed.add(2.3)))
          const originZ = mix(float(-4), float(0), hash(batchSeed.add(3.7)))

          positionValue.assign(vec3(originX, originY, originZ))

          const batchSpeed = mix(float(20), float(40), hash(batchSeed.add(7.9)))
          const batchAngle = mix(float(-0.7), float(0.7), hash(batchSeed.add(5.1)))
          const batchVelocityX = sin(batchAngle).mul(batchSpeed)
          const batchVelocityY = cos(batchAngle).mul(batchSpeed)
          const jitterX = mix(float(-0.05), float(0.05), hash(particleIndex.add(time.mul(5.1))))
          const jitterY = mix(float(-0.05), float(0.05), hash(particleIndex.add(time.mul(6.3))))
          const baseScale = mix(float(0.18), float(0.4), seed)

          velocityValue.assign(
            vec4(batchVelocityX.add(jitterX), batchVelocityY.add(jitterY), float(0), baseScale),
          )
          phase.assign(launchPhaseNode)
          scaleFadeRate.assign(0)
        })
      })().compute(particleCount)

      const updateParticles = Fn(() => {
        const particleIndex = instanceIndex.toFloat()
        const position = positionBuffer.element(instanceIndex)
        const life = position.w

        If(life.lessThanEqual(0), () => {
          Return()
        })

        const positionValue = position.xyz
        const velocityValue = velocityBuffer.element(instanceIndex)
        const movementVelocity = velocityValue.xyz
        const meta = metaBuffer.element(instanceIndex)
        const seed = meta.y
        const phase = meta.z
        const scaleFadeRate = meta.w
        const dt = simulationDelta
        const explosionDrag = float(1).sub(float(0.14).mul(dt))

        If(phase.lessThan(explosionPhaseNode), () => {
          velocityValue.y.subAssign(float(4.2).mul(dt))
          velocityValue.x.addAssign(
            sin(seed.mul(30).add(life.mul(4)))
              .mul(0.24)
              .mul(dt),
          )
          positionValue.addAssign(movementVelocity.mul(dt))

          If(life.lessThan(explosionTriggerLifeNode), () => {
            const directionCount = 16
            const particleTime = particleIndex.add(time).toVar()
            const angle = hash(particleTime.mul(13.7))
              .mul(directionCount)
              .floor()
              .mul(PI.div(directionCount / 2))
            const magnitude = mix(float(0.5), float(8), hash(particleTime.mul(19.3)))

            phase.assign(explosionPhaseNode)
            velocityValue.x.assign(cos(angle).mul(magnitude))
            velocityValue.y.assign(sin(angle).mul(magnitude))
            velocityValue.z.assign(mix(float(-5), float(0), hash(particleTime.mul(23.1))))
            velocityValue.w.mulAssign(explosionScaleMultiplierNode)
            scaleFadeRate.assign(velocityValue.w.div(life))
            Return()
          })
        })

        If(phase.greaterThanEqual(explosionPhaseNode), () => {
          velocityValue.y.subAssign(float(7).mul(dt))
          velocityValue.x.mulAssign(explosionDrag)
          velocityValue.y.mulAssign(explosionDrag)
          velocityValue.z.mulAssign(explosionDrag)
          positionValue.addAssign(movementVelocity.mul(dt))
        })

        life.subAssign(dt.mul(0.8))

        If(life.lessThanEqual(0), () => {
          life.assign(0)
          velocityValue.w.assign(0)
          Return()
        })

        If(phase.greaterThanEqual(explosionPhaseNode), () => {
          velocityValue.w.assign(scaleFadeRate.mul(life))
        })
      })().compute(particleCount)

      const colorNode = colorBuffer.toAttribute()
      const positionNode = positionBuffer.toAttribute()
      const spriteUV = spritesheetUV(vec2(2), uv(), colorNode.a)
      const maskNode = texture(spriteTexture, spriteUV).r.greaterThan(float(0))
      const opacityNode = vertexStage(
        smoothstep(float(0), float(0.18), positionNode.w).mul(
          mix(float(0.35), float(0.95), metaBuffer.toAttribute().y),
        ),
      )

      return {
        colorNode,
        maskNode,
        opacityNode,
        positionNode,
        resetParticles,
        scaleNode: vec2(velocityBuffer.toAttribute().w),
        setActiveBatch: (batch: number) => {
          activeBatch.value = batch
        },
        spawnBatch,
        updateParticles,
      }
    },
    [particleBuffers, particleCount, spriteTexture],
  )

  const simulation = useLocalNodes(createSimulationNodes)

  useEffect(() => {
    renderer.compute(simulation.resetParticles)
  }, [renderer, simulation])

  useEffect(() => {
    activeBatchRef.current = 0

    const triggerBurst = () => {
      simulation.setActiveBatch(activeBatchRef.current)
      activeBatchRef.current = (activeBatchRef.current + 1) % batchCount
      renderer.compute(simulation.spawnBatch)
    }
    let timeoutId = 0

    const scheduleBurst = () => {
      const delay = burstDelayMin + Math.random() * (burstDelayMax - burstDelayMin)

      timeoutId = window.setTimeout(() => {
        triggerBurst()
        scheduleBurst()
      }, delay)
    }

    triggerBurst()
    scheduleBurst()

    return () => window.clearTimeout(timeoutId)
  }, [batchCount, burstDelayMax, burstDelayMin, renderer, simulation])

  useFrame(
    () => {
      renderer.compute(simulation.updateParticles)
    },
    { fps },
  )

  return (
    <instancedMesh args={[undefined, undefined, particleCount]} frustumCulled={false}>
      <planeGeometry args={[0.8, 0.8, 1, 1]} />
      <spriteNodeMaterial
        key={simulation.colorNode.uuid}
        blending={THREE.AdditiveBlending}
        colorNode={simulation.colorNode}
        depthTest={false}
        depthWrite={false}
        maskNode={simulation.maskNode}
        opacityNode={simulation.opacityNode}
        positionNode={simulation.positionNode}
        scaleNode={simulation.scaleNode}
        transparent={true}
      />
    </instancedMesh>
  )
}

export default FireworksParticles