feuerwerk/src/core/shader/compile.ts

import { edgesInto, findNode, type Graph, type GraphNode } from '../graph/Graph';
import { nodeDef, type PortDef } from '../graph/NodeRegistry';
import { colorOf } from '../model/Properties';
import { isSpriteProp, isValidFieldName } from '../properties/PropertySchema';
import { isPropertyNode, propertyName, shaderParamDefs } from '../properties/shaders';
import {
  asVec3,
  DIVIDE_FN,
  type EdgeMode,
  edgeOf,
  type Expr,
  floatLiteral,
  isShaderNodeDef,
  SHADER_OUTPUT,
  SHADER_SAMPLE,
  SHADER_TIME,
  SHEET_EDGE,
  STRAIGHT_FN,
  TEXEL_ALPHA,
  TEXEL_COLOR,
  TIME_UNIFORM,
  uniformName,
  vec3Literal,
} from './nodes';

export interface ShaderUniform {
  name: string;
  default: number;
}

export interface CompiledShader {
  fragment: string;
  uniforms: ShaderUniform[];
  /** A Time node feeds the Output: the fragment 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;
}

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>;
  /** One line for the tab: empty once compiled. */
  reason: string;
}

/**
 * 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 float uniform. The texel
 * arrives premultiplied, so the prelude unpremultiplies and the epilogue multiplies back.
 */
/**
 * What the fragment depends on and nothing else: node ids, types, params and declared ports, and
 * which port feeds which. Positions, collapsing and folding are the editor's; they stay out so a
 * moved node does not compile again (#65).
 */
export function shaderFingerprint(graph: Graph): string {
  return JSON.stringify({
    nodes: graph.nodes.map((n) => [n.id, n.type, n.params, n.ports ?? null]),
    edges: graph.edges.map((e) => [e.from.node, e.from.port, e.to.node, e.to.port]),
  });
}

export function compileShader(graph: Graph): ShaderBuild {
  const problems: Record<string, string> = {};
  const problem = (node: GraphNode, text: string) => {
    if (!(node.id in problems)) problems[node.id] = text;
  };

  const names = new Set<string>();
  for (const node of graph.nodes) {
    if (!isPropertyNode(node)) continue;
    const name = propertyName(node);
    if (!isValidFieldName(name)) problem(node, 'Name must be one word: letters, digits, _');
    else if (isSpriteProp(name)) problem(node, `"${name}" is a sprite property`);
    else if (names.has(name)) problem(node, `Another property is named "${name}"`);
    names.add(name);
  }

  const lines: string[] = [];
  const emitted = new Map<string, Record<string, Expr>>();
  const visiting = new Set<string>();
  let count = 0;

  const emitNode = (node: GraphNode): Record<string, Expr> | null => {
    const done = emitted.get(node.id);
    if (done) return done;
    if (visiting.has(node.id)) {
      problem(node, 'Feeds itself');
      return null;
    }
    const def = nodeDef(node.type);
    if (!def || !isShaderNodeDef(def)) {
      problem(node, 'Not a shader node');
      return null;
    }
    visiting.add(node.id);
    const inputs: Record<string, Expr> = {};
    let complete = true;
    for (const port of def.inputs) {
      const expr = input(node, port);
      if (expr) inputs[port.key] = expr;
      else complete = false;
    }
    visiting.delete(node.id);
    if (!complete) return null;
    const outputs: Record<string, Expr> = {};
    for (const [key, expr] of Object.entries(def.emit(inputs, node.params))) {
      const local = `n${count++}_${key}`;
      lines.push(`${expr.type} ${local} = ${expr.code};`);
      outputs[key] = { code: local, type: expr.type };
    }
    emitted.set(node.id, outputs);
    return outputs;
  };

  /**
   * What feeds an input: the edge's source, the typed-in value, or – with nothing to type – what
   * the node says (`unfed`), else the texel: its colour for a colour port, its alpha for a number.
   */
  const input = (node: GraphNode, port: PortDef): Expr | null => {
    const edge = edgesInto(graph, node.id, port.key)[0];
    if (edge) {
      const source = findNode(graph, edge.from.node);
      const outputs = source && emitNode(source);
      const expr = outputs?.[edge.from.port];
      if (!expr) return null;
      return conform(node, port, expr);
    }
    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 def = nodeDef(node.type);
    const unfed = def && isShaderNodeDef(def) ? def.unfed?.[port.key] : undefined;
    return unfed ?? (port.type === 'number' ? TEXEL_ALPHA : TEXEL_COLOR);
  };

  const conform = (node: GraphNode, port: PortDef, expr: Expr): Expr | null => {
    if (port.type === 'number' && expr.type === 'vec3') {
      problem(node, `"${port.key}" needs a number, not a colour`);
      return null;
    }
    if (port.type === 'vec3' && expr.type === 'float') return { code: asVec3(expr), type: 'vec3' };
    return expr;
  };

  const output = graph.nodes.find((n) => n.type === SHADER_OUTPUT);
  if (!output) return { compiled: null, problems, reason: 'No Output node' };
  const outputDef = nodeDef(SHADER_OUTPUT)!;
  const color = input(output, outputDef.inputs[0]!);
  const alpha = input(output, outputDef.inputs[1]!);

  const first = graph.nodes.find((n) => n.id in problems);
  if (first || !color || !alpha)
    return { compiled: null, problems, reason: first ? problems[first.id]! : 'Cannot compile' };

  const uniforms = shaderParamDefs(graph).map((row) => ({ name: row.key, default: row.default }));
  const time = graph.nodes.some((n) => n.type === SHADER_TIME && emitted.has(n.id));
  // 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;',
    ...(time ? [`uniform float ${TIME_UNIFORM};`] : []),
    ...uniforms.map((u) => `uniform float ${uniformName(u.name)};`),
    STRAIGHT_FN,
    DIVIDE_FN,
    'void main()',
    '{',
    '    vec4 texel = texture(uTexture, 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.
  const shifted = graph.nodes.some(
    (n) =>
      n.type === SHADER_SAMPLE &&
      emitted.has(n.id) &&
      (edgesInto(graph, n.id, 'u').length > 0 || edgesInto(graph, n.id, 'v').length > 0),
  );
  const edge = shifted ? edgeOf(output.params) : SHEET_EDGE;
  return { compiled: { fragment, uniforms, time, edge }, problems, reason: '' };
}