feuerwerk/src/viewport/shaders.ts

import {
  Color,
  GlProgram,
  Matrix,
  type Mesh,
  type ParticleContainer,
  Shader,
  Texture,
  TextureStyle,
  UniformGroup,
} from 'pixi.js';
import type { PropertyBag } from '../core/model/Properties';
import { shaderOf } from '../core/properties/shaders';
import type { CompiledShader } from '../core/shader/compile';
import { TIME_UNIFORM, uniformName } from '../core/shader/nodes';

/** Where a node's shader comes from: the compiled form of the asset the sprite names, if it compiles. */
export interface ShaderSource {
  compiled(assetId: string): CompiledShader | null;
  /** The shader clock at this redraw, in seconds – what a Time node reads, the same on every sprite. */
  now: number;
}

/** The compiled shader the bag names, if there is one and it compiles. */
export function wornShader(props: PropertyBag, shaders: ShaderSource): CompiledShader | null {
  const id = shaderOf(props);
  return id ? shaders.compiled(id) : null;
}

/**
 * Pixi's own mesh vertex stage, minus what a sprite quad does not need: a frame is a whole texture
 * (`TextureLibrary`), so `aUV` maps 1:1 and there is no texture matrix. `uRound` is left out –
 * Pixi's uniform sync skips uniforms the program does not declare.
 */
const SPRITE_VERTEX = `
in vec2 aPosition;
in vec2 aUV;
out vec2 vTextureCoord;
out vec4 vColor;
uniform mat3 uProjectionMatrix;
uniform mat3 uWorldTransformMatrix;
uniform vec4 uWorldColorAlpha;
uniform mat3 uTransformMatrix;
uniform vec4 uColor;
void main()
{
    mat3 m = uProjectionMatrix * uWorldTransformMatrix * uTransformMatrix;
    gl_Position = vec4((m * vec3(aPosition, 1.0)).xy, 0.0, 1.0);
    vTextureCoord = aUV;
    vColor = uColor * uWorldColorAlpha;
}
`;

/**
 * Pixi's own particle vertex stage in GLSL ES 3.00, which is what the compiled fragment is written
 * in: `attribute`/`varying` become `in`/`out` and `vUV` takes the name the fragment reads. The
 * rounding is kept so a particle sits on the same pixel shaded as unshaded.
 */
const PARTICLE_VERTEX = `
in vec2 aVertex;
in vec2 aUV;
in vec4 aColor;
in vec2 aPosition;
in float aRotation;
out vec2 vTextureCoord;
out vec4 vColor;
uniform mat3 uTranslationMatrix;
uniform float uRound;
uniform vec2 uResolution;
uniform vec4 uColor;
vec2 roundPixels(vec2 position, vec2 targetSize)
{
    return (floor(((position * 0.5 + 0.5) * targetSize) + 0.5) / targetSize) * 2.0 - 1.0;
}
void main(void)
{
    float c = cos(aRotation);
    float s = sin(aRotation);
    vec2 v = vec2(aVertex.x * c - aVertex.y * s, aVertex.x * s + aVertex.y * c) + aPosition;
    gl_Position = vec4((uTranslationMatrix * vec3(v, 1.0)).xy, 0.0, 1.0);
    if (uRound == 1.0) gl_Position.xy = roundPixels(gl_Position.xy, uResolution);
    vTextureCoord = aUV;
    vColor = vec4(aColor.rgb * aColor.a, aColor.a) * uColor;
}
`;

/** One float uniform per graph parameter, plus the clock where a Time node reads one. */
function uniformGroupOf(compiled: CompiledShader): UniformGroup<Record<string, UniformFloat>> | null {
  const entries: [string, UniformFloat][] = compiled.uniforms.map((u) => [
    uniformName(u.name),
    { value: u.default, type: 'f32' },
  ]);
  if (compiled.time) entries.push([TIME_UNIFORM, { value: 0, type: 'f32' }]);
  return entries.length > 0 ? new UniformGroup(Object.fromEntries(entries)) : null;
}

type UniformFloat = { value: number; type: 'f32' };

function readUniforms(
  into: Record<string, number>,
  compiled: CompiledShader,
  props: PropertyBag,
  now: number,
): void {
  for (const u of compiled.uniforms) {
    const n = Number(props[u.name]);
    into[uniformName(u.name)] = Number.isFinite(n) ? n : u.default;
  }
  if (compiled.time) into[TIME_UNIFORM] = now;
}

/**
 * A compiled graph as a Pixi program. The program is shared (`GlProgram.from` caches by source),
 * the uniforms are this wearer's; `read` takes them from the item's bag. `resources` are the slots
 * Pixi fills for the vertex stage in question – `Shader.resources` only accepts the slots its
 * constructor declared, so a subclass names them all up front.
 */
abstract class CompiledGraphShader extends Shader {
  private readonly uniforms: Record<string, number> | null;
  private readonly compiled: CompiledShader;
  readonly fragment: string;

  constructor(
    readonly assetId: string,
    compiled: CompiledShader,
    program: { vertex: string; name: string },
    resources: Record<string, unknown>,
  ) {
    const group = uniformGroupOf(compiled);
    super({
      glProgram: GlProgram.from({ vertex: program.vertex, fragment: compiled.fragment, name: program.name }),
      resources: { ...resources, ...(group ? { shaderUniforms: group } : {}) },
    });
    this.fragment = compiled.fragment;
    this.compiled = compiled;
    this.uniforms = group ? group.uniforms : null;
  }

  /** True while this is still the shader `props` ask for, so it can be updated instead of rebuilt. */
  wears(id: string, compiled: CompiledShader): boolean {
    return this.assetId === id && this.fragment === compiled.fragment;
  }

  read(props: PropertyBag, now: number): void {
    if (this.uniforms) readUniforms(this.uniforms, this.compiled, props, now);
  }
}

/**
 * A shader asset worn by one sprite mesh: the fragment replaces Pixi's in the sprite's own draw,
 * so there is no filter and no extra pass (#77).
 */
class GraphShader extends CompiledGraphShader {
  /** What `Mesh.shader` asks of a shader; kept in step with the mesh's own by `wear`. */
  texture: Texture = Texture.WHITE;

  constructor(assetId: string, compiled: CompiledShader) {
    super(
      assetId,
      compiled,
      { vertex: SPRITE_VERTEX, name: `shader-${assetId}` },
      { uTexture: Texture.WHITE.source, uSampler: Texture.WHITE.source.style },
    );
  }

  /** Pixi binds a mesh's texture only for its own shader, so a custom one takes it here. */
  wear(texture: Texture): void {
    this.texture = texture;
    this.resources.uTexture = texture.source;
    this.resources.uSampler = texture.source.style;
  }
}

function wornBy(mesh: Mesh): GraphShader | null {
  return mesh.shader instanceof GraphShader ? mesh.shader : null;
}

/** The shader asset a mesh draws with, `''` for Pixi's own. */
export function wornShaderId(mesh: Mesh): string {
  return wornBy(mesh)?.assetId ?? '';
}

/**
 * Makes the mesh's shader match `props`: updated in place while asset and fragment are the same,
 * rebuilt when either changed, dropped for no shader – and for one that does not compile, which
 * leaves the mesh on Pixi's batched draw.
 */
export function syncShader(mesh: Mesh, props: PropertyBag, shaders: ShaderSource): void {
  const id = shaderOf(props);
  const compiled = id ? shaders.compiled(id) : null;
  const worn = wornBy(mesh);
  if (worn && compiled && worn.wears(id, compiled)) {
    worn.read(props, shaders.now);
    worn.wear(mesh.texture);
    return;
  }
  disposeShader(mesh);
  if (compiled) {
    const shader = new GraphShader(id, compiled);
    shader.read(props, shaders.now);
    shader.wear(mesh.texture);
    mesh.shader = shader;
  }
  // Pixi caches whether a mesh batches (MeshPipe's meshData) and re-checks it only in
  // validateRenderable, which it skips in any frame whose structure changed – and the draw order
  // re-attaches every object per redraw. A mesh that ever drew batched would so keep drawing
  // batched, plain, its shader ignored. `unload` drops that cache; the next frame reads `batched` afresh.
  if (Boolean(worn) !== Boolean(compiled)) mesh.unload();
}

/**
 * A shader asset worn by every particle of one source (#88). It replaces the container's whole
 * program, so the fragment sees one particle's own quad – its picture, not its life: Pixi's particle
 * buffer carries five fixed attributes and takes no sixth, so age, index and seed stay on the graph.
 *
 * The pipe breaks the batch for every field anyway, so unlike a shaded sprite (#77) this costs no
 * draw call at all: five hundred shaded particles are the one call five hundred plain ones were.
 */
class ParticleGraphShader extends CompiledGraphShader {
  constructor(assetId: string, compiled: CompiledShader) {
    // `GlParticleContainerAdaptor` writes the field's texture and the pipe's own uniforms into
    // whichever shader the container names, so both slots have to exist, empty, for Pixi to fill.
    super(
      assetId,
      compiled,
      { vertex: PARTICLE_VERTEX, name: `particle-shader-${assetId}` },
      {
        uTexture: Texture.WHITE.source,
        uSampler: new TextureStyle({}),
        uniforms: {
          uTranslationMatrix: { value: new Matrix(), type: 'mat3x3<f32>' },
          uColor: { value: new Color(0xffffff), type: 'vec4<f32>' },
          uRound: { value: 1, type: 'f32' },
          uResolution: { value: [0, 0], type: 'vec2<f32>' },
        },
      },
    );
  }
}

function wornByField(container: ParticleContainer): ParticleGraphShader | null {
  return container.shader instanceof ParticleGraphShader ? container.shader : null;
}

/** The compiled shader a source's particles are drawn with, if it names one that compiles. */
export function syncParticleShader(
  container: ParticleContainer,
  props: PropertyBag,
  shaders: ShaderSource,
): void {
  const id = shaderOf(props);
  const compiled = id ? shaders.compiled(id) : null;
  const worn = wornByField(container);
  if (worn && compiled && worn.wears(id, compiled)) {
    worn.read(props, shaders.now);
    return;
  }
  disposeParticleShader(container);
  if (compiled) {
    const shader = new ParticleGraphShader(id, compiled);
    shader.read(props, shaders.now);
    container.shader = shader;
  }
}

export function disposeParticleShader(container: ParticleContainer): void {
  const worn = wornByField(container);
  if (!worn) return;
  // Typed as always present, but the adaptor reads `container.shader || defaultShader` – dropping it
  // is how a field says it wears none.
  (container as { shader?: Shader }).shader = undefined;
  worn.destroy();
}

/** `Mesh.destroy` only lets go of its shader, so a mesh going away releases it here. */
export function disposeShader(mesh: Mesh): void {
  const worn = wornBy(mesh);
  if (!worn) return;
  mesh.shader = null;
  worn.destroy();
}