feuerwerk/src/core/commands/portCommands.ts

import {
  BEHAVIOR_PARAM,
  declaresPropertyPorts,
  declaresVariablePorts,
  type DeclaredPort,
  findNode,
  fixedPortKeys,
  FUNCTION_PARAM,
  type Graph,
  type GraphEdge,
  type GraphNode,
  isEntryTerminal,
  isFlowPort,
  isTerminalNode,
  isValidPortName,
  namesBehavior,
  OBJECT_BEHAVIOR_PARAM,
  propertyTargetOf,
  SHADER_CALL,
  variablesSourceOf,
} from '../graph';
import { isGraphAsset } from '../model/Asset';
import type { PropertyValue } from '../model/Properties';
import { kindLabel } from '../properties/kinds';
import type { Project } from '../model/Project';
import { type Removed, removeFirst, removeWhere, restoreAt } from '../util/list';
import { type Command, CommandError, RenameCommand } from './Command';
import { nextFieldName } from './naming';
import { requireGraph, requireNode, requireNodeDef } from './require';

interface PortSite {
  graph: Graph;
  nodeId: string;
  /** Whether the port is an input there – the terminal and the call node face opposite ways. */
  asInput: boolean;
}

/**
 * Where a change to `node`'s declaration has to be mirrored. For most nodes that is the node
 * itself. A terminal is the exception: its ports are also ports of every node running or hearing the
 * behavior, or calling the shader function – facing the other way, because what a start node hands
 * out the caller passes in.
 * A node that has no such port (an `On behavior` and a flow port of the start) has no edges on it either.
 */
function portSites(project: Project, assetId: string, node: GraphNode): PortSite[] {
  const sites = ownSites(requireGraph(project, assetId), node);
  if (!isTerminalNode(node)) return sites;
  // A shader function's ends are the ports of its calls in shaders and in other functions (#303).
  const callsIt = (caller: GraphNode) =>
    caller.type === SHADER_CALL
      ? caller.params[FUNCTION_PARAM] === assetId
      : namesBehavior(caller.type) && caller.params[BEHAVIOR_PARAM] === assetId;
  for (const other of project.assets) {
    if (!isGraphAsset(other)) continue;
    for (const caller of other.graph.nodes) {
      if (!callsIt(caller)) continue;
      sites.push({ graph: other.graph, nodeId: caller.id, asInput: isEntryTerminal(node.type) });
    }
  }
  return sites;
}

function ownSites(graph: Graph, node: GraphNode): PortSite[] {
  const side = requireNodeDef(node.type).declaresPorts;
  const site = (asInput: boolean): PortSite => ({ graph, nodeId: node.id, asInput });
  return side === 'both' ? [site(false), site(true)] : [site(side === 'inputs')];
}

type EdgeTest = { graph: Graph; test: (edge: GraphEdge) => boolean };
type DroppedEdges = { graph: Graph; entries: Removed<GraphEdge>[] };

/** One pass per edge list: a second pass over the same list would be restored at shifted indices. */
function dropEdges(tests: readonly EdgeTest[]): DroppedEdges[] {
  const byGraph = new Map<Graph, EdgeTest['test'][]>();
  for (const { graph, test } of tests) byGraph.set(graph, [...(byGraph.get(graph) ?? []), test]);
  return [...byGraph].map(([graph, all]) => ({
    graph,
    entries: removeWhere(graph.edges, (edge) => all.some((test) => test(edge))),
  }));
}

function restoreEdges(dropped: readonly DroppedEdges[]): void {
  for (const { graph, entries } of dropped) restoreAt(graph.edges, entries);
}

/** A node declaring both sides has its flow pins only while one of its ports writes. */
function flowEdgesOnceSilent(graph: Graph, node: GraphNode): EdgeTest[] {
  if (requireNodeDef(node.type).declaresPorts !== 'both' || (node.ports ?? []).some((p) => p.write))
    return [];
  const onFlow = (end: { node: string; port: string }) => end.node === node.id && end.port === 'exec';
  return [{ graph, test: (edge) => onFlow(edge.to) || onFlow(edge.from) }];
}

/** An output's edges include those on its field ports, while it is split. */
function edgesOnPort(site: PortSite, key: string): (edge: GraphEdge) => boolean {
  return (edge) =>
    site.asInput
      ? edge.to.node === site.nodeId && edge.to.port === key
      : edge.from.node === site.nodeId && (edge.from.port === key || edge.from.port.startsWith(`${key}.`));
}

function nodeAt(site: PortSite): GraphNode | undefined {
  return findNode(site.graph, site.nodeId);
}

/** The port's edges follow its new key, and so does an input's typed-in value, which lives under it. */
function carryPort(site: PortSite, from: string, to: string): void {
  for (const edge of site.graph.edges.filter(edgesOnPort(site, from))) {
    if (site.asInput) edge.to.port = to;
    else edge.from.port = `${to}${edge.from.port.slice(from.length)}`;
  }
  const split = site.asInput ? undefined : nodeAt(site)?.splits?.find((s) => s.port === from);
  if (split) split.port = to;
  const node = site.asInput ? nodeAt(site) : undefined;
  const value = node?.params[from];
  if (!node || value === undefined) return;
  node.params[to] = value;
  delete node.params[from];
}

// A struct field may be called `exec`, a port may not – so this message has to name that reason.
const INVALID_NAME =
  'Port names must start with a letter, contain only letters, digits or _, and "exec" is taken by the flow port';

function checkPortName(node: GraphNode, key: string): void {
  if (!isValidPortName(key)) throw new CommandError(INVALID_NAME);
  if (fixedPortKeys(requireNodeDef(node.type)).includes(key))
    throw new CommandError(`"${key}" is taken by a port of the node itself`);
}

/** Several ports as one step, so seeding a node from a struct is one undo rather than one per field. */
export class AddPortsCommand implements Command {
  readonly label: string;
  protected readonly ports: readonly DeclaredPort[];

  constructor(
    private readonly assetId: string,
    private readonly nodeId: string,
    ports: readonly DeclaredPort[],
  ) {
    this.ports = ports.map((p) => ({ ...p }));
    this.label = this.ports.length === 1 ? `Add ${this.ports[0]?.key}` : `Add ${this.ports.length} ports`;
  }

  execute(project: Project): void {
    const node = requireNode(project, this.assetId, this.nodeId);
    const keys = new Set((node.ports ?? []).map((p) => p.key));
    const kinds = requireNodeDef(node.type).portKinds;
    for (const port of this.ports) {
      checkPortName(node, port.key);
      if (kinds && !kinds.includes(port.kind))
        throw new CommandError(`A ${kindLabel(port.kind)} port has no place on this node`);
      if (isFlowPort(port) && !isTerminalNode(node))
        throw new CommandError('Only a start or an end node declares a flow port');
      if (keys.has(port.key)) throw new CommandError(`"${port.key}" already exists`);
      keys.add(port.key);
    }
    node.ports = [...(node.ports ?? []), ...this.ports.map((p) => ({ ...p }))];
  }

  undo(project: Project): void {
    const node = requireNode(project, this.assetId, this.nodeId);
    const added = new Set(this.ports.map((p) => p.key));
    node.ports = (node.ports ?? []).filter((p) => !added.has(p.key));
  }
}

export class AddPortCommand extends AddPortsCommand {
  constructor(assetId: string, nodeId: string, port: DeclaredPort) {
    super(assetId, nodeId, [port]);
  }
}

/** Takes the edges hanging on the port with it – in this graph and in every calling one. */
export class RemovePortCommand implements Command {
  readonly label: string;
  private removed?: Removed<DeclaredPort>;
  private edges: DroppedEdges[] = [];
  private values: { site: PortSite; value: PropertyValue }[] = [];

  constructor(
    private readonly assetId: string,
    private readonly nodeId: string,
    private readonly key: string,
  ) {
    this.label = `Remove ${key}`;
  }

  execute(project: Project): void {
    const node = requireNode(project, this.assetId, this.nodeId);
    const sites = portSites(project, this.assetId, node);
    const ports = [...(node.ports ?? [])];
    const removed = removeFirst(ports, (p) => p.key === this.key);
    if (!removed) throw new CommandError(`Port "${this.key}" not found`);
    this.removed = removed;
    node.ports = ports;
    const graph = requireGraph(project, this.assetId);
    this.edges = dropEdges([
      ...sites.map((site) => ({ graph: site.graph, test: edgesOnPort(site, this.key) })),
      ...flowEdgesOnceSilent(graph, node),
    ]);
    this.values = [];
    for (const site of sites) {
      const target = site.asInput ? nodeAt(site) : undefined;
      const value = target?.params[this.key];
      if (!target || value === undefined) continue;
      this.values.push({ site, value });
      delete target.params[this.key];
    }
  }

  undo(project: Project): void {
    if (!this.removed) return;
    const node = requireNode(project, this.assetId, this.nodeId);
    const ports = [...(node.ports ?? [])];
    restoreAt(ports, [this.removed]);
    node.ports = ports;
    restoreEdges(this.edges);
    for (const { site, value } of this.values) {
      const target = nodeAt(site);
      if (target) target.params[this.key] = value;
    }
  }
}

/**
 * Switches whether a port of a node declaring both sides is written. Switched off, its input goes
 * with its edge and typed-in value, and the last one off takes the flow pins and their edges along.
 */
export class SetPortWriteCommand implements Command {
  readonly label: string;
  private edges: DroppedEdges[] = [];
  private value: PropertyValue | undefined;
  private was = false;

  constructor(
    private readonly assetId: string,
    private readonly nodeId: string,
    private readonly key: string,
    private readonly write: boolean,
  ) {
    this.label = write ? `Set ${key}` : `Stop setting ${key}`;
  }

  execute(project: Project): void {
    const node = requireNode(project, this.assetId, this.nodeId);
    if (requireNodeDef(node.type).declaresPorts !== 'both')
      throw new CommandError('This node does not write its ports');
    const port = (node.ports ?? []).find((p) => p.key === this.key);
    if (!port) throw new CommandError(`Port "${this.key}" not found`);
    this.was = port.write === true;
    this.mark(node, this.write);
    this.edges = [];
    this.value = undefined;
    if (this.write) return;
    const graph = requireGraph(project, this.assetId);
    const input: PortSite = { graph, nodeId: node.id, asInput: true };
    this.edges = dropEdges([
      { graph, test: edgesOnPort(input, this.key) },
      ...flowEdgesOnceSilent(graph, node),
    ]);
    this.value = node.params[this.key];
    delete node.params[this.key];
  }

  undo(project: Project): void {
    const node = requireNode(project, this.assetId, this.nodeId);
    this.mark(node, this.was);
    restoreEdges(this.edges);
    if (this.value !== undefined) node.params[this.key] = this.value;
  }

  private mark(node: GraphNode, write: boolean): void {
    node.ports = (node.ports ?? []).map((port) => {
      if (port.key !== this.key) return port;
      const { write: _, ...plain } = port;
      return write ? { ...plain, write: true } : plain;
    });
  }
}

/** Rewrites the edges and inline values pointing at the port too, in every graph that has any. */
export class RenamePortCommand extends RenameCommand {
  constructor(
    private readonly assetId: string,
    private readonly nodeId: string,
    from: string,
    to: string,
  ) {
    super(`Rename ${from}`, from, to);
  }

  protected rename(project: Project, from: string, to: string): void {
    if (from === to) return;
    const node = requireNode(project, this.assetId, this.nodeId);
    const sites = portSites(project, this.assetId, node);
    checkPortName(node, to);
    const ports = node.ports ?? [];
    if (!ports.some((p) => p.key === from)) throw new CommandError(`Port "${from}" not found`);
    if (ports.some((p) => p.key === to)) throw new CommandError(`"${to}" already exists`);
    node.ports = ports.map((p) => (p.key === from ? { ...p, key: to } : p));
    for (const site of sites) carryPort(site, from, to);
  }
}

/** Renames one declared port on one node, taking its edges and its typed-in value along. */
function renamePortOn(graph: Graph, node: GraphNode, from: string, to: string): void {
  const ports = node.ports ?? [];
  if (from === to || !ports.some((p) => p.key === from) || ports.some((p) => p.key === to)) return;
  node.ports = ports.map((p) => (p.key === from ? { ...p, key: to } : p));
  for (const site of ownSites(graph, node)) carryPort(site, from, to);
}

/**
 * Points one declared port at another variable or property. Key and kind change together, because
 * there the kind is the variable's (the property's) rather than a choice of its own; the edges
 * hanging on the port come along, so correcting a pick is not a rewiring.
 */
export class RepointPortCommand implements Command {
  readonly label: string;
  private previous?: DeclaredPort;

  constructor(
    private readonly assetId: string,
    private readonly nodeId: string,
    private readonly from: string,
    private readonly port: DeclaredPort,
  ) {
    this.label = `Point ${from} at ${port.key}`;
  }

  execute(project: Project): void {
    const node = requireNode(project, this.assetId, this.nodeId);
    const ports = node.ports ?? [];
    const current = ports.find((p) => p.key === this.from);
    if (!current) throw new CommandError(`Port "${this.from}" not found`);
    checkPortName(node, this.port.key);
    if (this.from !== this.port.key && ports.some((p) => p.key === this.port.key))
      throw new CommandError(`"${this.port.key}" already exists`);
    this.previous = { ...current };
    this.swap(project, node, this.from, this.port);
  }

  undo(project: Project): void {
    if (!this.previous) return;
    this.swap(project, requireNode(project, this.assetId, this.nodeId), this.port.key, this.previous);
  }

  private swap(project: Project, node: GraphNode, from: string, to: DeclaredPort): void {
    renamePortOn(requireGraph(project, this.assetId), node, from, to.key);
    node.ports = (node.ports ?? []).map((p) => (p.key === to.key ? { ...p, kind: to.kind } : p));
  }
}

/**
 * Follows a renamed variable into the nodes that name it as a port: the Variables nodes of its own
 * behavior, and Properties nodes anywhere that point at that behavior. Their edges and typed-in
 * values move with it – the port key *is* the variable name there.
 */
export function renameVariablePorts(project: Project, behaviorId: string, from: string, to: string): void {
  for (const asset of project.assets) {
    if (asset.type !== 'behavior') continue;
    for (const node of asset.graph.nodes)
      if (namesVariablesOf(node, asset.id, behaviorId)) renamePortOn(asset.graph, node, from, to);
  }
}

function namesVariablesOf(node: GraphNode, graphOwner: string, behaviorId: string): boolean {
  if (declaresVariablePorts(node.type)) {
    // A node naming a Variables asset follows that asset wherever it sits, not the graph around it.
    const source = variablesSourceOf(node.params);
    return source ? source === behaviorId : graphOwner === behaviorId;
  }
  return (
    declaresPropertyPorts(node.type) &&
    propertyTargetOf(node.params) === 'behavior' &&
    node.params[OBJECT_BEHAVIOR_PARAM] === behaviorId
  );
}

/** "in1", "in2", … – the first port name the terminal node does not use yet. */
export function nextPortName(node: GraphNode, base: string): string {
  return nextFieldName(node.ports ?? [], base);
}