feuerwerk/src/core/shader/functions.ts

import {
  type DeclaredPort,
  edgesInto,
  type Graph,
  type GraphEdge,
  graphFingerprint,
  type GraphNode,
} from '../graph/Graph';
import type { GraphLookup } from '../graph/lookup';
import {
  calledFunctions,
  FUNCTION_PARAM,
  SHADER_CALL,
  SHADER_FUNCTION_INPUT,
  SHADER_FUNCTION_OUTPUT,
  terminalNode,
} from '../graph/nested';
import { nodeDef, type PortDef } from '../graph/NodeRegistry';
import { textOf } from '../model/Properties';
import { SHADER_COLOR_PROPERTY, SHADER_PROPERTY } from '../properties/shaders';
import {
  NUMBER_STEP,
  SHADER_COLOR,
  SHADER_NUMBER,
  SHADER_OUTPUT,
  SHADER_VERTEX_OUTPUT,
  type ShaderNodeDef,
} from './nodes';

/** The shader functions a compile may set in, by asset id (#303). */
export type FunctionLookup = GraphLookup['shaderFunction'];
export const NO_FUNCTIONS: FunctionLookup = () => null;

/** A function takes and gives values only: a picture stays the caller's. */
const FUNCTION_PORT_KINDS = ['number', 'color'] as const;
const DEFAULT_COLOR = 0xffffff;

/** A declared port of a function as the nodes on either side of it see it. */
export function functionPortOf(port: DeclaredPort): PortDef {
  const { key } = port;
  if (port.kind === 'color')
    return {
      key,
      label: key,
      type: 'vec3',
      default: DEFAULT_COLOR,
      field: { key, label: key, kind: 'color', default: DEFAULT_COLOR },
    };
  return {
    key,
    label: key,
    type: 'number',
    default: 0,
    field: { key, label: key, kind: 'number', default: 0, step: NUMBER_STEP },
  };
}

function functionPorts(graph: Graph | undefined, type: string): PortDef[] {
  const node = graph && terminalNode(graph, type);
  return (node?.ports ?? []).map(functionPortOf);
}

export const ShaderFunctionInput: ShaderNodeDef = {
  type: SHADER_FUNCTION_INPUT,
  label: 'Function input',
  category: 'shader',
  description:
    'The values a shader calling this function hands in. Every input declared here, a number or a color, becomes an input port on the node that calls the function.',
  inputs: [],
  outputs: [],
  params: [],
  declaresPorts: 'outputs',
  portKinds: FUNCTION_PORT_KINDS,
  declaredPort: (port) => functionPortOf(port),
  emit: () => ({}),
};

export const ShaderFunctionOutput: ShaderNodeDef = {
  type: SHADER_FUNCTION_OUTPUT,
  label: 'Function output',
  category: 'shader',
  description:
    'What this function hands back. Every output declared here becomes an output port on the node that calls the function; unconnected, it is the value typed in.',
  inputs: [],
  outputs: [],
  params: [],
  declaresPorts: 'inputs',
  portKinds: FUNCTION_PORT_KINDS,
  declaredPort: (port) => functionPortOf(port),
  emit: () => ({}),
};

/** Never emitted in place: the compiler sets the function's graph in where it stands. */
export const ShaderCall: ShaderNodeDef = {
  type: SHADER_CALL,
  label: 'Shader function',
  category: 'shader',
  description:
    "A shader function's graph, set in here. Its ports are the ones that function declares at its Function input and Function output, and the picture it reads is this shader's. Double-click opens the function.",
  inputs: [],
  outputs: [],
  params: [
    { key: FUNCTION_PARAM, label: 'Function', kind: 'asset', default: '', assetTypes: ['shaderfunction'] },
  ],
  derivePorts: ({ node, lookup }) => {
    const called = lookup.shaderFunction(textOf(node.params[FUNCTION_PARAM] ?? ''));
    return {
      label: called?.name,
      inputs: functionPorts(called?.graph, SHADER_FUNCTION_INPUT),
      outputs: functionPorts(called?.graph, SHADER_FUNCTION_OUTPUT),
    };
  },
  emit: () => ({}),
};

export const SHADER_FUNCTION_NODES: readonly ShaderNodeDef[] = [
  ShaderFunctionInput,
  ShaderFunctionOutput,
  ShaderCall,
];

/**
 * A function hands values back and nothing else, so what ends a shader or gives the wearer a row
 * has no place in it; the ends of a function have none in a shader.
 */
const ONLY_IN_SHADER = new Set([SHADER_OUTPUT, SHADER_VERTEX_OUTPUT, SHADER_PROPERTY, SHADER_COLOR_PROPERTY]);
const ONLY_IN_FUNCTION = new Set([SHADER_FUNCTION_INPUT, SHADER_FUNCTION_OUTPUT]);

/** Why a node of that type cannot stand in a shader (or, `inFunction`, a function); null if it can. */
export function misplacedProblem(type: string, inFunction: boolean): string | null {
  if (inFunction && ONLY_IN_SHADER.has(type)) return `A function cannot hold ${nodeDef(type)?.label ?? type}`;
  if (!inFunction && ONLY_IN_FUNCTION.has(type)) return 'Only in a shader function';
  return null;
}

/** Whether a node of that type belongs in a function's palette – or, with false, in a shader's. */
export function offeredIn(type: string, inFunction: boolean): boolean {
  return misplacedProblem(type, inFunction) === null;
}

export interface Inlined {
  /** The graph with every call that could be set in replaced by the function's nodes. */
  graph: Graph;
  /** Per call set in, the function's name: a problem inside it is reported at the call. */
  names: Record<string, string>;
  /** Per call left in place, why it could not be set in. */
  unresolved: Record<string, string>;
}

/**
 * Every call in `graph` replaced by a copy of the function it names (#303), its nodes renamed
 * `<call id>/<node id>`: what fed a port of the call feeds what the Function input's port fed,
 * and what read an output of the call reads what fed the Function output's port. A port nothing
 * feeds becomes a constant with its typed-in value. `chain` is the functions already being set in
 * – meeting one again is a function calling itself.
 */
export function inlineFunctions(
  graph: Graph,
  functions: FunctionLookup,
  chain: readonly string[] = [],
): Inlined {
  const names: Record<string, string> = {};
  const unresolved: Record<string, string> = {};
  let current = graph;
  for (const call of graph.nodes) {
    if (call.type !== SHADER_CALL) continue;
    const id = textOf(call.params[FUNCTION_PARAM] ?? '');
    const called = id ? functions(id) : null;
    const set = !id
      ? 'No function picked'
      : !called
        ? 'That function is gone'
        : chain.includes(id)
          ? 'A function cannot call itself'
          : setIn(current, call, called.name, called.graph, functions, [...chain, id]);
    if (typeof set === 'string') unresolved[call.id] = set;
    else {
      current = set;
      names[call.id] = called!.name;
    }
  }
  return { graph: current, names, unresolved };
}

type End = { node: string; port: string };

function setIn(
  graph: Graph,
  call: GraphNode,
  name: string,
  body: Graph,
  functions: FunctionLookup,
  chain: readonly string[],
): Graph | string {
  for (const node of body.nodes) {
    const misplaced = misplacedProblem(node.type, true);
    if (misplaced) return `In ${name}: ${misplaced}`;
  }
  const inner = inlineFunctions(body, functions, chain);
  const stuck = Object.values(inner.unresolved)[0];
  if (stuck) return `In ${name}: ${stuck}`;
  const fi = terminalNode(inner.graph, SHADER_FUNCTION_INPUT);
  const fo = terminalNode(inner.graph, SHADER_FUNCTION_OUTPUT);
  const outs = new Map((fo?.ports ?? []).map((p) => [p.key, p]));
  const gone = graph.edges.find((e) => e.from.node === call.id && !outs.has(e.from.port));
  if (gone) return `"${gone.from.port}" is not an output of ${name}`;

  const prefix = `${call.id}/`;
  const constants: GraphNode[] = [];
  const constant = (key: string, port: DeclaredPort, typed: unknown): End => {
    const color = functionPortOf(port).type === 'vec3';
    const value = typed ?? functionPortOf(port).default;
    const id = `${prefix}${key}`;
    constants.push({
      id,
      type: color ? SHADER_COLOR : SHADER_NUMBER,
      x: 0,
      y: 0,
      params: { value: value as number },
    });
    return { node: id, port: color ? 'color' : 'value' };
  };
  const fed: Record<string, End> = {};
  for (const port of fi?.ports ?? []) {
    const edge = edgesInto(graph, call.id, port.key)[0];
    fed[port.key] = edge ? { ...edge.from } : constant(`in:${port.key}`, port, call.params[port.key]);
  }
  /** Where an end inside the function comes from once it is set in; null for one that is gone. */
  const outside = (end: End): End | null =>
    end.node === fi?.id
      ? (fed[end.port] ?? null)
      : end.node === fo?.id
        ? null
        : { node: prefix + end.node, port: end.port };

  const edges: GraphEdge[] = [];
  const link = (from: End, to: End) => edges.push({ id: `${prefix}e${edges.length}`, from, to });
  for (const edge of inner.graph.edges) {
    if (edge.to.node === fi?.id || edge.to.node === fo?.id) continue;
    const from = outside(edge.from);
    if (from) link(from, { node: prefix + edge.to.node, port: edge.to.port });
  }
  const results: Record<string, End> = {};
  for (const [key, port] of outs) {
    const edge = edgesInto(inner.graph, fo!.id, key)[0];
    results[key] = (edge && outside(edge.from)) ?? constant(`out:${key}`, port, fo!.params[key]);
  }
  const kept: GraphEdge[] = [];
  for (const edge of graph.edges) {
    if (edge.to.node === call.id) continue;
    if (edge.from.node === call.id) link(results[edge.from.port]!, { ...edge.to });
    else kept.push(edge);
  }
  const copies = inner.graph.nodes
    .filter((n) => n !== fi && n !== fo)
    .map((n) => ({ ...n, id: prefix + n.id }));
  return {
    ...graph,
    nodes: [...graph.nodes.filter((n) => n !== call), ...copies, ...constants],
    edges: [...kept, ...edges],
  };
}

/**
 * Problems of the graph a compile saw, moved onto the graph its author sees: one inside a function
 * goes to the call it was set in at, in that function's words.
 */
export function problemsAtCalls(problems: Record<string, string>, inlined: Inlined): Record<string, string> {
  const lifted: Record<string, string> = {};
  for (const [id, text] of Object.entries(problems)) {
    const cut = id.indexOf('/');
    const call = cut < 0 ? id : id.slice(0, cut);
    const said = cut < 0 ? text : `In ${inlined.names[call] ?? 'the function'}: ${text}`;
    if (!(call in lifted)) lifted[call] = said;
  }
  return lifted;
}

/**
 * What a compile of `graph` depends on beyond its own nodes: every function it calls, however
 * deep, by content. A shader recompiles when one of them changes.
 */
export function functionsFingerprint(graph: Graph, functions: FunctionLookup): string {
  const parts: string[] = [];
  const seen = new Set<string>();
  const visit = (g: Graph) => {
    for (const id of calledFunctions(g)) {
      if (seen.has(id)) continue;
      seen.add(id);
      const called = functions(id);
      parts.push(`${id}:${called ? `${called.name}:${graphFingerprint(called.graph)}` : 'gone'}`);
      if (called) visit(called.graph);
    }
  };
  visit(graph);
  return parts.join('\n');
}