feuerwerk/src/core/shader/compile.ts

import { graphEmitter } from '../graph/emitGraph';
import { edgesInto, type Graph, type GraphNode } from '../graph/Graph';
import { nodeDef } from '../graph/NodeRegistry';
import { colorOf } from '../model/Properties';
import { isSpriteProp, isValidFieldName } from '../properties/PropertySchema';
import { isPropertyNode, propertyName, shaderParamDefs } from '../properties/shaders';
import {
  asVec3,
  BLUR_FN,
  DIVIDE_FN,
  type EdgeMode,
  edgeOf,
  type Expr,
  floatLiteral,
  isShaderNodeDef,
  paddingOf,
  SHADER_BLUR,
  SHADER_OUTPUT,
  SHADER_SAMPLE,
  SHADER_TIME,
  SHADER_VERTEX_OUTPUT,
  type ShaderNodeDef,
  SHEET_EDGE,
  STRAIGHT_FN,
  subdivisionsOf,
  TEXEL_ALPHA,
  TEXEL_AT,
  TEXEL_COLOR,
  texelAtFn,
  TIME_UNIFORM,
  uniformName,
  vec3Literal,
} from './nodes';

export interface ShaderUniform {
  name: string;
  /** A colour is a vec3 uniform and its default a packed 0xRRGGBB (#298). */
  kind: 'number' | 'color';
  default: number;
}

const glslOf = (u: ShaderUniform) => (u.kind === 'color' ? 'vec3' : 'float');

/**
 * What a Vertex output node adds to the wearer's vertex stage (#79). It reads the corner's
 * `vTextureCoord` and the same uniforms the fragment does.
 */
export interface ShaderVertex {
  /** Declarations ahead of `main`: the uniforms, at the fragment's precision, and the helpers. */
  head: string;
  /** Statements inside `main`, ending in `vec2 offset`: the corner's move in pixels of the picture. */
  body: string;
  /** Cells per axis of the wearer's grid; 1 is the plain quad. */
  subdivisions: number;
}

export interface CompiledShader {
  fragment: string;
  /** Null without a Vertex output node: the wearer then draws on the quad every sprite shares. */
  vertex: ShaderVertex | null;
  uniforms: ShaderUniform[];
  /** A Time node feeds an output: the program takes `uTime`, and a sprite wearing it redraws every tick. */
  time: boolean;
  /**
   * The texture variant the sprite's mesh draws with – the Output's Edge where the graph samples
   * past the border, and the plain sheet where it cannot.
   */
  edge: EdgeMode;
  /** Pixels of the picture the wearer draws beyond it on every side (#69). */
  padding: number;
}

export interface ShaderBuild {
  /** Null while the graph has a problem – the sprite then draws without a filter. */
  compiled: CompiledShader | null;
  /** Per node id, what keeps the graph from compiling. */
  problems: Record<string, string>;
  /** Per node id, what the graph compiles with but its author will hardly have meant. */
  hints: Record<string, string>;
  /** One line for the tab: empty once compiled. */
  reason: string;
}

export const PADDING_CLAMP_HINT = 'Padding with Clamp smears the border';
export const perCornerProblem = (what: string) => `Not available per corner: ${what}`;

/**
 * A shader graph to the fragment of the sprite's own draw: every node reachable from the Output
 * becomes one local per output, in dependency order, and every Property a uniform. The texel
 * arrives premultiplied, so the prelude unpremultiplies and the epilogue multiplies back.
 */
export function compileShader(graph: Graph): ShaderBuild {
  const lines: string[] = [];
  /** What the vertex walk met that only a pixel has; the Vertex output node reports the first. */
  const perPixel: string[] = [];
  const walk = walker(graph, lines, 'n', null);
  const { problems } = walk;

  walk.checkNames(graph.nodes.filter(isPropertyNode), propertyName, (name) =>
    !isValidFieldName(name)
      ? 'Name must be one word: letters, digits, _'
      : isSpriteProp(name)
        ? `"${name}" is a sprite property`
        : '',
  );

  const output = graph.nodes.find((n) => n.type === SHADER_OUTPUT);
  if (!output) return { compiled: null, problems, hints: {}, reason: 'No Output node' };
  const mode = edgeOf(output.params);
  const padding = paddingOf(output.params);
  const hints: Record<string, string> =
    padding > 0 && mode === 'clamp' ? { [output.id]: PADDING_CLAMP_HINT } : {};
  const outputDef = nodeDef(SHADER_OUTPUT)!;
  const color = walk.input(output, outputDef.inputs[0]!);
  const alpha = walk.input(output, outputDef.inputs[1]!);

  const uniforms = shaderParamDefs(graph).map((row) => ({
    name: row.key,
    kind: row.kind,
    default: row.default,
  }));

  // The per-corner part is a walk of its own from the Vertex output: a node feeding both outputs is
  // emitted once for each stage.
  const corner = graph.nodes.find((n) => n.type === SHADER_VERTEX_OUTPUT);
  const cornerLines: string[] = [];
  const cornerWalk = walker(graph, cornerLines, 'c', perPixel);
  let vertex: ShaderVertex | null = null;
  if (corner) {
    const cornerDef = nodeDef(SHADER_VERTEX_OUTPUT)!;
    const x = cornerWalk.input(corner, cornerDef.inputs[0]!);
    const y = cornerWalk.input(corner, cornerDef.inputs[1]!);
    if (perPixel.length > 0) cornerWalk.problem(corner, perCornerProblem(perPixel[0]!));
    for (const [id, text] of Object.entries(cornerWalk.problems)) if (!(id in problems)) problems[id] = text;
    const reachesTime = graph.nodes.some((n) => n.type === SHADER_TIME && cornerWalk.reached(n));
    if (x && y)
      vertex = {
        // A uniform both stages declare has to have one precision; the fragment's default is mediump.
        head: [
          ...(reachesTime ? [`uniform mediump float ${TIME_UNIFORM};`] : []),
          ...uniforms.map((u) => `uniform mediump ${glslOf(u)} ${uniformName(u.name)};`),
          DIVIDE_FN,
        ].join('\n'),
        body: [...cornerLines, `vec2 offset = vec2(${x.code}, ${y.code});`].join('\n'),
        subdivisions: subdivisionsOf(corner.params),
      };
  }

  const failed = firstProblem(graph, problems);
  if (failed !== null || !color || !alpha || (corner && !vertex))
    return { compiled: null, problems, hints, reason: failed ?? 'Cannot compile' };

  const time = graph.nodes.some((n) => n.type === SHADER_TIME && (walk.reached(n) || cornerWalk.reached(n)));
  const fragmentTime = graph.nodes.some((n) => n.type === SHADER_TIME && walk.reached(n));
  const blurs = graph.nodes.some((n) => n.type === SHADER_BLUR && walk.reached(n));
  // The fragment runs in the sprite's own draw (#77). `vColor` is the tint, premultiplied like the
  // texel; it goes in before the graph, so `Texture color` is the sprite as it draws without a shader.
  // The version line is what makes Pixi compile the program as GLSL ES 3.00 (`GlProgram` looks for
  // it in the fragment); without it `texture()` is a `#define` and `textureSize` does not exist.
  const fragment = [
    '#version 300 es',
    'in vec2 vTextureCoord;',
    'in vec4 vColor;',
    'out vec4 finalColor;',
    'uniform sampler2D uTexture;',
    ...(fragmentTime ? [`uniform float ${TIME_UNIFORM};`] : []),
    ...uniforms.map((u) => `uniform ${glslOf(u)} ${uniformName(u.name)};`),
    STRAIGHT_FN,
    DIVIDE_FN,
    texelAtFn(mode),
    ...(blurs ? [BLUR_FN] : []),
    'void main()',
    '{',
    `    vec4 texel = ${TEXEL_AT}(vTextureCoord) * vColor;`,
    `    vec3 ${TEXEL_COLOR.code} = straight(texel);`,
    ...lines.map((line) => `    ${line}`),
    `    float alpha = clamp(${alpha.code}, 0.0, 1.0);`,
    `    finalColor = vec4(clamp(${color.code}, 0.0, 1.0) * alpha, alpha);`,
    '}',
    '',
  ].join('\n');
  // The edge only shows where a Sample reads somewhere other than its own pixel. A graph that never
  // shifts a UV stays on the plain sheet, so the common shader costs no second copy of the picture.
  // With padding every pixel of the margin reads past the border; Transparent is cut in the
  // fragment and takes the plain sheet as well. A Blur reads its neighbours, so it always shifts.
  const shifted =
    blurs ||
    graph.nodes.some(
      (n) =>
        n.type === SHADER_SAMPLE &&
        walk.reached(n) &&
        (edgesInto(graph, n.id, 'u').length > 0 || edgesInto(graph, n.id, 'v').length > 0),
    );
  const edge = (shifted || padding > 0) && mode !== 'transparent' ? mode : SHEET_EDGE;
  return { compiled: { fragment, vertex, uniforms, time, edge, padding }, problems, hints, reason: '' };
}

function firstProblem(graph: Graph, problems: Record<string, string>): string | null {
  const first = graph.nodes.find((n) => n.id in problems);
  return first ? problems[first.id]! : null;
}

/**
 * One walk of the graph into `lines`, a local per node output. `perPixel` is the list a per-corner
 * walk notes in what it met that only a pixel has; the fragment's walk hands in null.
 */
function walker(graph: Graph, lines: string[], prefix: string, perPixel: string[] | null) {
  let count = 0;
  const walk = graphEmitter<ShaderNodeDef, Expr>(graph, {
    isDef: isShaderNodeDef,
    noun: 'shader',
    emit: (def, inputs, node) => {
      if (def.perPixel) perPixel?.push(def.label);
      let prepared: string | undefined;
      if (def.prepare) {
        prepared = `${prefix}${count++}_prepared`;
        lines.push(`vec4 ${prepared} = ${def.prepare(inputs, node.params)};`);
      }
      const outputs: Record<string, Expr> = {};
      for (const [key, expr] of Object.entries(def.emit(inputs, node.params, prepared))) {
        const local = `${prefix}${count++}_${key}`;
        lines.push(`${expr.type} ${local} = ${expr.code};`);
        outputs[key] = { code: local, type: expr.type };
      }
      return outputs;
    },
    // With nothing to type and nothing the node says, an input is the texel: its colour for a
    // colour port, its alpha for a number.
    unfed: (node, port, def) => {
      const typed = node.params[port.key] ?? port.default;
      if (port.field?.kind === 'color')
        return { code: vec3Literal(colorOf(typed ?? 0xffffff, 0xffffff)), type: 'vec3' };
      if (port.field) return { code: floatLiteral(Number(typed ?? 0)), type: 'float' };
      const said = def?.unfed?.[port.key];
      if (said) return said;
      perPixel?.push('the picture');
      return port.type === 'number' ? TEXEL_ALPHA : TEXEL_COLOR;
    },
    conform: (node: GraphNode, port, expr) => {
      if (port.type === 'number' && expr.type === 'vec3') {
        walk.problem(node, `"${port.key}" needs a number, not a color`);
        return null;
      }
      if (port.type === 'vec3' && expr.type === 'float') return { code: asVec3(expr), type: 'vec3' };
      return expr;
    },
  });
  return walk;
}