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 { isPropertyNode, isValidPropertyName, propertyName, shaderParamDefs } from '../properties/shaders';
import {
  asVec3,
  type Expr,
  floatLiteral,
  isShaderNodeDef,
  SHADER_OUTPUT,
  TEXEL_COLOR,
  uniformName,
  vec3Literal,
} from './nodes';

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

export interface CompiledShader {
  fragment: string;
  uniforms: ShaderUniform[];
}

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 a Pixi filter fragment: every node reachable from the Output becomes one
 * local per output, in dependency order, and every Property a float uniform. Pixi hands a filter
 * premultiplied pixels, so the prelude unpremultiplies and the epilogue multiplies back.
 */
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 (!isValidPropertyName(name)) problem(node, 'Name must be one word: letters, digits, _');
    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 – the texel. */
  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' };
    return 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 first = graph.nodes.find((n) => n.id in problems);
  if (first || !color)
    return { compiled: null, problems, reason: first ? problems[first.id]! : 'Cannot compile' };

  const uniforms = shaderParamDefs(graph).map((row) => ({ name: row.key, default: row.default }));
  const fragment = [
    'in vec2 vTextureCoord;',
    'out vec4 finalColor;',
    'uniform sampler2D uTexture;',
    ...uniforms.map((u) => `uniform float ${uniformName(u.name)};`),
    'void main()',
    '{',
    '    vec4 texel = texture(uTexture, vTextureCoord);',
    `    vec3 ${TEXEL_COLOR.code} = texel.a > 0.0 ? texel.rgb / texel.a : texel.rgb;`,
    ...lines.map((line) => `    ${line}`),
    `    finalColor = vec4(clamp(${color.code}, 0.0, 1.0) * texel.a, texel.a);`,
    '}',
    '',
  ].join('\n');
  return { compiled: { fragment, uniforms }, problems, reason: '' };
}