Shader stage

import {useEffect, useLayoutEffect, useRef, useState, type CSSProperties, type ReactNode} from "react";
import {useFrame, useReadiness, useVideo} from "odori";
import {effect, frame as gpuFrame, init, surface, target, type Effect, type Gpu, type Surface, type Target} from "vgpu";

export type ShaderUniformState = {
  /** The current frame. */
  frame: number;
  /** The frame as seconds, which is what most shaders animate on. */
  time: number;
  /** The canvas's backing size in pixels. */
  width: number;
  height: number;
};

export type ShaderStageProps = {
  /**
   * A WGSL fragment shader as a plain string. The fullscreen vertex stage is
   * provided, so the fragment entry receives `@location(0) uv: vec2f` and
   * nothing else is required. Uniforms are ordinary `var<uniform>`
   * declarations, addressed by name from `uniforms`.
   */
  wgsl: string;
  /**
   * The uniform values for a frame. Called with the frame clock, so a shader
   * animates on `time` the way a component animates on `useFrame`: the value
   * is a pure function of the frame, which is what keeps the export
   * deterministic. Return only names the shader declares.
   */
  uniforms?: (state: ShaderUniformState) => Record<string, unknown>;
  /**
   * Drawn when the browser has no WebGPU. Studio in an older browser shows
   * this instead of a blank; the render pipeline always has WebGPU when the
   * video declares `webgpu: true`.
   */
  fallback?: ReactNode;
  /**
   * Draw on a transparent surface, compositing over whatever is underneath.
   * The shader then returns premultiplied alpha: `vec4f(rgb * a, a)`. This is
   * what makes a fullscreen pass a filter rather than a background.
   */
  alpha?: boolean;
  style?: CSSProperties;
};

/* The device type comes through vgpu rather than from the ambient WebGPU
   globals, so installing this file does not oblige a project to also carry
   @webgpu/types. */
type Stage = {gpu: Gpu; surface: Surface; effect: Effect; device: Gpu["device"]["gpu"]; offscreen?: Target};

/**
 * How `odori test` sees this canvas's pixels.
 *
 * A WebGPU canvas cannot be read back reliably from outside: `drawImage`
 * reads the current texture, which the compositor expires the moment it
 * presents, so an external reader wins or loses a race it cannot control.
 * The stage owns the shader, so it can do what no reader can — re-encode the
 * current frame into an offscreen target and read genuine pixels, on demand.
 * The hook only runs when something calls it; exports and Studio never do.
 */
type ReadbackCanvas = HTMLCanvasElement & {__odoriReadback?: () => Promise<{hash: string; blank: boolean}>};

/**
 * One GPU context for the whole page, acquired on first use and kept for the
 * page's life. A stage per device reads naturally but does not survive
 * contact with real pages: a scene that stacks a stage and two filters, plus
 * the runtime's hidden audio-compile mount, is six devices, and a browser is
 * allowed to start destroying them. Surfaces, effects, and targets are all
 * per-stage on the one shared device, which is also what lets the pipeline
 * caches actually cache.
 */
let sharedGpu: Promise<Gpu> | null = null;
const acquireGpu = (): Promise<Gpu> => {
  if (sharedGpu) return sharedGpu;
  const acquired = init().then(
    (created) => {
      // The platform may take the device away at any moment, and the spec
      // says so. When it happens the shared handle is retired, so the next
      // acquisition builds a fresh device rather than encoding into a dead
      // one; the stages watching this device rebuild themselves on it.
      created.device.gpu.lost.then(() => {
        if (sharedGpu === acquired) sharedGpu = null;
      });
      return created;
    },
    (error: unknown) => {
      // A failed acquisition does not poison the page: the next mount retries.
      if (sharedGpu === acquired) sharedGpu = null;
      throw error;
    },
  );
  sharedGpu = acquired;
  return acquired;
};

/**
 * A WebGPU shader as a component on the frame clock.
 *
 * The two clocks are the whole problem. A GPU wants to draw on
 * `requestAnimationFrame`, which is wall time; a video is a pure function of
 * the frame. So there is no loop here: every change of frame encodes exactly
 * one GPU frame, and the readiness handshake holds the capture until the
 * queue reports that frame's work done. Scrubbing backwards, exporting in
 * parallel chunks, and rendering the same frame twice all produce the same
 * pixels, which is what makes a shader a citizen of the video rather than a
 * live widget embedded in one.
 *
 * The canvas backs at the composition's own resolution and stretches to its
 * container, so a full-bleed stage samples one fragment per output pixel.
 */
export const ShaderStage = ({wgsl, uniforms, fallback, alpha = false, style}: ShaderStageProps) => {
  const canvas = useRef<HTMLCanvasElement | null>(null);
  const stage = useRef<Stage | null>(null);
  const frame = useFrame();
  const {width, height, fps} = useVideo();
  const readiness = useReadiness();
  const [ready, setReady] = useState(false);
  const [unsupported, setUnsupported] = useState(false);
  // Bumped when the device is lost: the init effect depends on it, so a bump
  // is a rebuild on a fresh device. Capped, because a platform that loses
  // every device it grants is a platform without WebGPU.
  const [epoch, setEpoch] = useState(0);
  const latestUniforms = useRef(uniforms);
  latestUniforms.current = uniforms;

  useEffect(() => {
    const element = canvas.current;
    if (!element) return;
    if (typeof navigator === "undefined" || !(navigator as Navigator & {gpu?: unknown}).gpu) {
      setUnsupported(true);
      return;
    }
    if (epoch >= 3) {
      // A platform that loses every device it grants is a platform without
      // WebGPU, whatever navigator.gpu says.
      setUnsupported(true);
      return;
    }

    // Held across the async device acquisition, so the first frame cannot be
    // captured before the shader exists.
    const release = readiness.hold();
    let disposed = false;

    (async () => {
      const created = await acquireGpu();
      if (disposed) return;
      created.device.gpu.lost.then(() => {
        if (disposed) return;
        console.warn("ShaderStage: the GPU device was lost; rebuilding on a fresh one.");
        setReady(false);
        setEpoch((count) => count + 1);
      });
      stage.current = {
        gpu: created,
        surface: surface(created, element, {
          autoResize: false,
          dpr: 1,
          size: [width, height],
          ...(alpha ? {alphaMode: "premultiplied" as const, clearColor: [0, 0, 0, 0] as const} : {}),
        }),
        effect: effect(created, wgsl),
        device: created.device.gpu,
      };
      (element as ReadbackCanvas).__odoriReadback = async () => {
        const current = stage.current;
        if (!current) return {hash: "no-context", blank: true};
        try {
          current.offscreen ??= target(current.gpu, {size: [width, height], clearColor: [0, 0, 0, 0]});
          const offscreen = current.offscreen;
          gpuFrame(current.gpu, (pass) => pass.pass(offscreen, current.effect));
          const pixels = await offscreen.read();
          // The same prime-stride digest the test computes for a 2d canvas.
          const count = width * height;
          let hash = 5381;
          let opaque = 0;
          for (let pixel = 0; pixel < count; pixel += 97) {
            const offset = pixel * 4;
            if (pixels[offset + 3] > 8) opaque += 1;
            hash = ((hash << 5) + hash + pixels[offset] + pixels[offset + 1] * 3 + pixels[offset + 2] * 7 + pixels[offset + 3] * 11) | 0;
          }
          return {hash: String(hash), blank: opaque === 0};
        } catch (error) {
          // A lost device drew nothing, and blank is the honest report.
          console.error("ShaderStage readback failed.", error);
          return {hash: `gpu-lost: ${String(error).slice(0, 140)}`, blank: true};
        }
      };
      setReady(true);
    })()
      .catch((error) => {
        // A machine whose WebGPU exists but cannot produce a device gets the
        // fallback rather than a hung frame marker.
        console.error("ShaderStage could not acquire a WebGPU device.", error);
        setUnsupported(true);
      })
      .finally(release);

    return () => {
      disposed = true;
      delete (element as ReadbackCanvas).__odoriReadback;
      // The device is the page's; only this stage's hold on the canvas ends.
      // Disposing a surface on an already-lost device has nothing to free.
      try {
        stage.current?.surface.dispose();
      } catch {
        // Nothing to release: the device the surface lived on is gone.
      }
      stage.current = null;
      setReady(false);
    };
    // The shader is the identity of the stage: a new source is a new pipeline.
  }, [wgsl, alpha, width, height, epoch]);

  useLayoutEffect(() => {
    const current = stage.current;
    if (!ready || !current) return;

    // Held until the queue has finished this frame's work: the capture then
    // sees the drawn pixels, never the previous frame's or an empty buffer.
    const release = readiness.hold();
    (async () => {
      const state: ShaderUniformState = {frame, time: frame / fps, width, height};
      const values = latestUniforms.current?.(state);
      if (values) current.effect.set(values);
      gpuFrame(current.gpu, (pass) => pass.pass(current.surface, current.effect));
      await current.device.queue.onSubmittedWorkDone();
    })()
      .catch((error) => console.error("ShaderStage failed to draw.", error))
      .finally(release);
  }, [ready, frame, width, height, fps]);

  if (unsupported) return <>{fallback}</>;

  return (
    <canvas
      ref={canvas}
      width={width}
      height={height}
      style={{display: "block", height: "100%", width: "100%", ...style}}
    />
  );
};

import {defineComponentPreview} from "odori/preview";
import {ShaderStage} from "./shader-stage";

const AURORA = /* wgsl */ `
struct Params { time: f32, aspect: f32 }
@group(0) @binding(0) var<uniform> params: Params;

@fragment fn main(@location(0) uv: vec2f) -> @location(0) vec4f {
  let p = vec2f((uv.x - 0.5) * params.aspect, uv.y - 0.62);
  var v = 0.0;
  var amp = 0.6;
  var q = p * 3.0;
  for (var i = 0; i < 5; i++) {
    v += amp * sin(q.x + params.time * (0.4 + f32(i) * 0.13) + sin(q.y * 1.7 + params.time * 0.35));
    q = vec2f(q.y * 1.6 - q.x * 0.4, q.x * 1.6 + q.y * 0.4) + vec2f(1.3, -0.7);
    amp *= 0.55;
  }
  let band = 0.5 + 0.5 * sin(v * 2.2 + p.y * 4.0);
  let glow = pow(band, 3.0);
  let falloff = exp(-2.6 * abs(p.y + 0.18));
  let color = vec3f(0.015) + vec3f(0.93) * mix(vec3f(1.0), vec3f(0.72, 0.82, 1.0), 0.4) * glow * falloff;
  return vec4f(color, 1.0);
}
`;

const RINGS = /* wgsl */ `
struct Params { time: f32, aspect: f32 }
@group(0) @binding(0) var<uniform> params: Params;

@fragment fn main(@location(0) uv: vec2f) -> @location(0) vec4f {
  let p = vec2f((uv.x - 0.5) * params.aspect, uv.y - 0.5);
  let d = length(p);
  let wave = sin(d * 26.0 - params.time * 2.4) * 0.5 + 0.5;
  let ring = pow(wave, 6.0) * exp(-3.0 * d);
  return vec4f(vec3f(0.02) + vec3f(0.9) * ring, 1.0);
}
`;

export default defineComponentPreview({
  title: "Shader stage",
  category: "Shaders",
  description: "A WGSL fragment shader drawn through WebGPU on the frame clock.",
  component: ShaderStage,
  canvas: {width: 1920, height: 1080, duration: "8s"},
  examples: [
    {
      name: "Aurora",
      props: {wgsl: AURORA, uniforms: ({time, width, height}) => ({time, aspect: width / height})},
    },
    {
      name: "Rings",
      props: {wgsl: RINGS, uniforms: ({time, width, height}) => ({time, aspect: width / height})},
    },
  ],
});

import {Scene, Video} from "odori";
import {ShaderStage} from "../components/shader-stage/shader-stage";

<Video>
  <Scene id="shader-stage" duration="5s">
    <ShaderStage />
  </Scene>
</Video>