Fast Canvas Text in React Three Fiber with WebGPU

'use client'

import {
  type CreatorState,
  useFrame,
  useLocalNodes,
  useThree,
  useUniforms,
} from '@react-three/fiber/webgpu'
import { type FC, useEffect } from 'react'
import { mix, texture, uv, vec4 } from 'three/tsl'
import { type Color, type UniformNode, type Vector3Tuple } from 'three/webgpu'

import { TRANSPARENT_TEXTURE, useFastTextCanvas } from './useFastTextCanvas'

const COUNTDOWN_CENTISECONDS = 10_000
const COUNTDOWN_START = (COUNTDOWN_CENTISECONDS - 1) / 100
const FAST_CANVAS_TEXT_UNIFORM_SCOPE = 'fastCanvasText'
const TEXT_ASPECT_RATIO = 5

type FastCanvasTextUniforms = {
  uColorFrom: UniformNode<'color', Color>
  uColorTo: UniformNode<'color', Color>
}

function formatCountdown(value: number): string {
  return value.toFixed(2).padStart(5, '0')
}

type Props = {
  position?: Vector3Tuple
  height?: number
  canvasScale?: number
  textAlign?: 'left' | 'center' | 'right'
}

const FastCanvasText: FC<Props> = ({
  position = [0, 0, 0],
  height = 120,
  canvasScale = 1,
  textAlign = 'center',
}) => {
  const dpr = useThree((s) => s.viewport.dpr ?? 1)
  const canvasHeight = Math.max(1, Math.round(height * dpr * canvasScale))
  const canvasWidth = Math.max(1, Math.round(height * TEXT_ASPECT_RATIO * dpr * canvasScale))

  const { canvasState, onTextValueChange } = useFastTextCanvas({
    width: canvasWidth,
    height: canvasHeight,
    textAlign,
    fontSize: Math.floor(canvasHeight),
    fontWeight: 900,
    format: formatCountdown,
  })

  useUniforms(
    {
      uColorFrom: '#ffe071',
      uColorTo: '#ff4b93',
    },
    FAST_CANVAS_TEXT_UNIFORM_SCOPE,
  )

  const { colorNode, emissiveNode } = useLocalNodes(({ uniforms }: CreatorState) => {
    const { uColorFrom, uColorTo } = uniforms.scope<FastCanvasTextUniforms>(
      FAST_CANVAS_TEXT_UNIFORM_SCOPE,
    )
    const alphaNode = texture(canvasState?.texture ?? TRANSPARENT_TEXTURE).r
    const colorNode = vec4(mix(uColorFrom, uColorTo, uv().y), alphaNode)
    const emissiveNode = colorNode.mul(uv().x.oneMinus()).mul(alphaNode)

    return {
      colorNode,
      emissiveNode,
    }
  })

  useEffect(() => {
    onTextValueChange(COUNTDOWN_START)
  }, [onTextValueChange])

  useFrame(
    ({ elapsed }) => {
      const elapsedCentiseconds = Math.floor(elapsed * 100) % COUNTDOWN_CENTISECONDS
      const countdownValue = (COUNTDOWN_CENTISECONDS - 1 - elapsedCentiseconds) / 100
      onTextValueChange(countdownValue)
    },
    { fps: 30 },
  )

  return (
    <mesh position={position} frustumCulled={false}>
      <planeGeometry args={[height * TEXT_ASPECT_RATIO, height, 1, 1]} />
      <meshLambertNodeMaterial
        key={colorNode.uuid}
        lights={false}
        colorNode={colorNode}
        emissiveNode={emissiveNode}
        transparent={true}
      />
    </mesh>
  )
}

export default FastCanvasText

'use client'

import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
import {
  CanvasTexture,
  ClampToEdgeWrapping,
  DataTexture,
  LinearFilter,
  NoColorSpace,
  type Texture,
} from 'three/webgpu'

export type CanvasState = {
  canvas: HTMLCanvasElement
  context: CanvasRenderingContext2D
  texture: CanvasTexture
}
function configureCanvasTexture(
  texture: CanvasTexture,
  colorSpace: Texture['colorSpace'] = NoColorSpace,
): void {
  texture.colorSpace = colorSpace
  texture.generateMipmaps = false
  texture.minFilter = LinearFilter
  texture.magFilter = LinearFilter
  texture.wrapS = ClampToEdgeWrapping
  texture.wrapT = ClampToEdgeWrapping
}

export function setupCanvasTexture(
  width: number,
  height: number,
  alpha = false,
  colorSpace: Texture['colorSpace'] = NoColorSpace,
): CanvasState {
  const canvas = document.createElement('canvas')
  canvas.width = width
  canvas.height = height
  canvas.style.display = 'none'

  const context = canvas.getContext('2d', { alpha })
  if (!context) throw new Error('Failed to get canvas 2D context')

  const texture = new CanvasTexture(canvas)
  configureCanvasTexture(texture, colorSpace)

  return { canvas, context, texture }
}

const transparentPixel = new Uint8Array([0, 0, 0, 0])
export const TRANSPARENT_TEXTURE: Texture = new DataTexture(transparentPixel, 1, 1)
TRANSPARENT_TEXTURE.generateMipmaps = false
TRANSPARENT_TEXTURE.needsUpdate = true
TRANSPARENT_TEXTURE.minFilter = LinearFilter
TRANSPARENT_TEXTURE.magFilter = LinearFilter
TRANSPARENT_TEXTURE.wrapS = ClampToEdgeWrapping
TRANSPARENT_TEXTURE.wrapT = ClampToEdgeWrapping

type FastTextCanvasOptions = {
  width: number
  height: number
  fontWeight?: number | string
  fontSize?: number
  textAlign?: CanvasTextAlign
  format?: (value: number) => string
}

type UseFastTextCanvasResult = {
  canvasState: CanvasState | null
  onTextValueChange: (value: number) => void
}

type NormalizedOptions = {
  width: number
  height: number
  fontWeight: number | string
  fontSize: number
  textAlign: CanvasTextAlign
  formatter: (value: number) => string
}

const FONT_FAMILY = 'monospace'
const TEXT_COLOR = '#ffffff'

const defaultFormatter = (value: number): string => value.toString()

const DEFAULTS: Omit<NormalizedOptions, 'width' | 'height'> = {
  fontWeight: 900,
  fontSize: 128,
  textAlign: 'center',
  formatter: defaultFormatter,
}

function drawCanvasText(
  context: CanvasRenderingContext2D,
  text: string,
  options: NormalizedOptions,
): void {
  const { width, height, textAlign, fontWeight, fontSize } = options

  context.clearRect(0, 0, width, height)

  if (context.textAlign !== textAlign) {
    context.textAlign = textAlign
  }

  context.textBaseline = 'middle'

  const x =
    textAlign === 'left' || textAlign === 'start'
      ? 0
      : textAlign === 'right' || textAlign === 'end'
        ? width
        : width / 2

  const nextFont = `${fontWeight} ${fontSize}px ${FONT_FAMILY}`
  if (context.font !== nextFont) context.font = nextFont

  if (context.fillStyle !== TEXT_COLOR) {
    context.fillStyle = TEXT_COLOR
  }

  context.fillText(text, x, height / 2)
}

export function useFastTextCanvas(options: FastTextCanvasOptions): UseFastTextCanvasResult {
  const [canvasState, setCanvasState] = useState<CanvasState | null>(null)
  const canvasStateRef = useRef<CanvasState | null>(null)
  const lastValueRef = useRef<number | null>(null)
  const pendingTextRef = useRef<string | null>(null)
  const lastDrawnTextRef = useRef<string | null>(null)
  const frameIdRef = useRef<number | null>(null)

  const normalizedOptions = useMemo<NormalizedOptions>(
    () => ({
      width: options.width,
      height: options.height,
      fontWeight: options.fontWeight ?? DEFAULTS.fontWeight,
      fontSize: options.fontSize ?? DEFAULTS.fontSize,
      textAlign: options.textAlign ?? DEFAULTS.textAlign,
      formatter: options.format ?? DEFAULTS.formatter,
    }),
    [
      options.width,
      options.height,
      options.fontWeight,
      options.fontSize,
      options.textAlign,
      options.format,
    ],
  )

  const flushDraw = useCallback(() => {
    const nextCanvasState = canvasStateRef.current
    const nextText = pendingTextRef.current

    if (!nextCanvasState || nextText === null) return
    if (nextText === lastDrawnTextRef.current) return

    drawCanvasText(nextCanvasState.context, nextText, normalizedOptions)
    nextCanvasState.texture.needsUpdate = true
    lastDrawnTextRef.current = nextText
  }, [normalizedOptions])

  const scheduleDraw = useCallback(() => {
    if (frameIdRef.current !== null) return

    frameIdRef.current = requestAnimationFrame(() => {
      frameIdRef.current = null
      flushDraw()
    })
  }, [flushDraw])

  const onTextValueChange = useCallback(
    (value: number) => {
      lastValueRef.current = value

      const nextText = normalizedOptions.formatter(value)
      if (nextText === pendingTextRef.current) return

      pendingTextRef.current = nextText
      scheduleDraw()
    },
    [normalizedOptions, scheduleDraw],
  )

  useEffect(() => {
    const nextCanvasState = setupCanvasTexture(
      normalizedOptions.width,
      normalizedOptions.height,
    )
    canvasStateRef.current = nextCanvasState

    let isDisposed = false
    const frameId = requestAnimationFrame(() => {
      if (isDisposed) return
      setCanvasState(nextCanvasState)
    })

    lastDrawnTextRef.current = null
    scheduleDraw()

    return () => {
      isDisposed = true
      cancelAnimationFrame(frameId)

      if (frameIdRef.current !== null) {
        cancelAnimationFrame(frameIdRef.current)
        frameIdRef.current = null
      }

      if (canvasStateRef.current === nextCanvasState) {
        canvasStateRef.current = null
      }

      nextCanvasState.texture.dispose()
    }
  }, [
    normalizedOptions.fontSize,
    normalizedOptions.height,
    normalizedOptions.width,
    scheduleDraw,
  ])

  useEffect(() => {
    if (lastValueRef.current === null) return

    pendingTextRef.current = normalizedOptions.formatter(lastValueRef.current)
    lastDrawnTextRef.current = null
    scheduleDraw()
  }, [normalizedOptions, scheduleDraw])

  return { canvasState, onTextValueChange }
}