Mesh Surface Sampled Particles in React Three Fiber with TSL

'use client'

import { useGLTF } from '@react-three/drei'
import { useFrame, useLocalNodes, useThree } from '@react-three/fiber/webgpu'
import { type FC, useCallback, useMemo } from 'react'
import {
  Fn,
  If,
  color,
  deltaTime,
  hash,
  instanceIndex,
  instancedArray,
  shapeCircle,
  smoothstep,
  vec2,
  vertexStage,
} from 'three/tsl'
import { AdditiveBlending, type SkinnedMesh } from 'three/webgpu'

import { sampleMeshSurface } from './sampleMeshSurface'

const MODEL_URL = new URL('./phoenix_compressed.glb', import.meta.url).href
const PARTICLE_COUNT = 4096
const PARTICLE_FADE_START = 0.6
const PARTICLE_LOOP_DISTANCE = 0.2
const PARTICLE_OPACITY = 0.85
const PARTICLE_RISE_SPEED_MIN = 0.1
const PARTICLE_RISE_SPEED_MAX = 0.4

type PhoenixGLTF = {
  nodes: {
    Pheonix_Baked_Baked: SkinnedMesh
  }
}

const MeshSurfaceSampledParticles: FC = () => {
  const { nodes } = useGLTF(MODEL_URL) as unknown as PhoenixGLTF
  const particleBuffers = useMemo(() => {
    const sampledPositions = sampleMeshSurface(nodes.Pheonix_Baked_Baked, PARTICLE_COUNT)

    return {
      initialPositionBuffer: instancedArray(sampledPositions, 'vec3'),
      positionBuffer: instancedArray(sampledPositions.slice(), 'vec3'),
    }
  }, [nodes.Pheonix_Baked_Baked])

  const createParticleNodes = useCallback(() => {
    const { initialPositionBuffer, positionBuffer } = particleBuffers
    const updateParticles = Fn(() => {
      const initialPosition = initialPositionBuffer.element(instanceIndex)
      const position = positionBuffer.element(instanceIndex)
      const riseSpeed = hash(instanceIndex)
        .mul(PARTICLE_RISE_SPEED_MAX - PARTICLE_RISE_SPEED_MIN)
        .add(PARTICLE_RISE_SPEED_MIN)

      position.z.subAssign(deltaTime.mul(riseSpeed))

      If(position.z.lessThan(initialPosition.z.sub(PARTICLE_LOOP_DISTANCE)), () => {
        position.assign(initialPosition)
      })
    })().compute(PARTICLE_COUNT)
    const lifeProgress = initialPositionBuffer
      .element(instanceIndex)
      .z.sub(positionBuffer.element(instanceIndex).z)
      .div(PARTICLE_LOOP_DISTANCE)
    const opacityNode = vertexStage(
      smoothstep(PARTICLE_FADE_START, 1, lifeProgress).oneMinus().mul(PARTICLE_OPACITY),
    )

    return {
      colorNode: color('#ff4b93'),
      maskNode: shapeCircle(),
      opacityNode,
      positionNode: positionBuffer.toAttribute(),
      scaleNode: vec2(0.004),
      updateParticles,
    }
  }, [particleBuffers])
  const particles = useLocalNodes(createParticleNodes)
  const renderer = useThree((state) => state.renderer)

  useFrame(
    () => {
      renderer.compute(particles.updateParticles)
    },
    { phase: 'update', fps: 60, drop: true },
  )

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

export default MeshSurfaceSampledParticles

import { MeshSurfaceSampler } from 'three/addons/math/MeshSurfaceSampler.js'
import { type Mesh, Vector3 } from 'three/webgpu'

export const sampleMeshSurface = (mesh: Mesh, count: number): Float32Array => {
  const positions = new Float32Array(count * 3)
  const point = new Vector3()
  const sampler = new MeshSurfaceSampler(mesh).build()

  mesh.updateWorldMatrix(true, false)

  for (let index = 0; index < count; index += 1) {
    sampler.sample(point)
    point.applyMatrix4(mesh.matrixWorld).toArray(positions, index * 3)
  }

  return positions
}