nodes/frontend/src/lib/graph-manager.ts

312 lines
8.8 KiB
TypeScript
Raw Normal View History

import { writable, type Writable } from "svelte/store";
import { type Graph, type Node, type Edge, type Socket, type NodeRegistry, type RuntimeExecutor } from "./types";
2024-03-13 16:18:48 +01:00
import { HistoryManager } from "./history-manager";
import * as templates from "./graphs";
2024-03-19 16:56:42 +01:00
import EventEmitter from "./helpers/EventEmitter";
2024-03-06 18:31:06 +01:00
2024-03-19 16:56:42 +01:00
export class GraphManager extends EventEmitter<{ "save": Graph }> {
2024-03-06 18:31:06 +01:00
status: Writable<"loading" | "idle" | "error"> = writable("loading");
graph: Graph = { nodes: [], edges: [] };
private _nodes: Map<number, Node> = new Map();
nodes: Writable<Map<number, Node>> = writable(new Map());
2024-03-11 22:00:16 +01:00
private _edges: Edge[] = [];
edges: Writable<Edge[]> = writable([]);
2024-03-06 18:31:06 +01:00
2024-03-14 16:55:11 +01:00
inputSockets: Writable<Set<string>> = writable(new Set());
2024-03-13 16:18:48 +01:00
history: HistoryManager = new HistoryManager(this);
constructor(private nodeRegistry: NodeRegistry, private runtime: RuntimeExecutor) {
2024-03-19 16:56:42 +01:00
super();
2024-03-11 22:00:16 +01:00
this.nodes.subscribe((nodes) => {
this._nodes = nodes;
});
this.edges.subscribe((edges) => {
this._edges = edges;
2024-03-14 16:55:11 +01:00
const s = new Set<string>();
for (const edge of edges) {
s.add(`${edge[2].id}-${edge[3]}`);
}
this.inputSockets.set(s);
2024-03-11 22:00:16 +01:00
});
2024-03-13 14:30:30 +01:00
}
2024-03-13 16:18:48 +01:00
serialize(): Graph {
2024-03-13 14:30:30 +01:00
const nodes = Array.from(this._nodes.values()).map(node => ({
id: node.id,
2024-03-13 16:18:48 +01:00
position: { x: node.position.x, y: node.position.y },
2024-03-13 14:30:30 +01:00
type: node.type,
props: node.props,
}));
2024-03-13 16:18:48 +01:00
const edges = this._edges.map(edge => [edge[0].id, edge[1], edge[2].id, edge[3]]) as Graph["edges"];
2024-03-13 14:30:30 +01:00
return { nodes, edges };
2024-03-06 18:31:06 +01:00
}
execute() {
if (!this.runtime["loaded"]) return;
const start = performance.now();
const result = this.runtime.execute(this.serialize());
const end = performance.now();
console.log(`Execution took ${end - start}ms -> ${result}`);
}
2024-03-13 16:18:48 +01:00
private _init(graph: Graph) {
const nodes = new Map(graph.nodes.map(node => {
const nodeType = this.nodeRegistry.getNode(node.type);
2024-03-13 16:18:48 +01:00
if (nodeType) {
node.tmp = node.tmp || {};
node.tmp.type = nodeType;
2024-03-06 18:31:06 +01:00
}
2024-03-13 16:18:48 +01:00
return [node.id, node]
}));
2024-03-13 16:18:48 +01:00
this.edges.set(graph.edges.map((edge) => {
const from = nodes.get(edge[0]);
const to = nodes.get(edge[2]);
if (!from || !to) {
console.error("Edge references non-existing node");
return;
};
from.tmp = from.tmp || {};
from.tmp.children = from.tmp.children || [];
from.tmp.children.push(to);
to.tmp = to.tmp || {};
to.tmp.parents = to.tmp.parents || [];
to.tmp.parents.push(from);
2024-03-11 22:00:16 +01:00
return [from, edge[1], to, edge[3]] as const;
})
.filter(Boolean) as unknown as [Node, number, Node, string][]
);
this.nodes.set(nodes);
2024-03-11 22:00:16 +01:00
2024-03-13 16:18:48 +01:00
}
async load(graph: Graph) {
this.graph = graph;
this.status.set("loading");
2024-03-13 16:18:48 +01:00
for (const node of this.graph.nodes) {
const nodeType = this.nodeRegistry.getNode(node.type);
if (!nodeType) {
console.error(`Node type not found: ${node.type}`);
this.status.set("error");
return;
}
node.tmp = node.tmp || {};
node.tmp.type = nodeType;
}
this._init(this.graph);
setTimeout(() => {
this.status.set("idle");
this.save();
}, 100)
2024-03-06 18:31:06 +01:00
}
2024-03-11 22:00:16 +01:00
getAllNodes() {
return Array.from(this._nodes.values());
}
2024-03-11 19:37:58 +01:00
getNode(id: number) {
return this._nodes.get(id);
2024-03-06 18:31:06 +01:00
}
2024-03-13 14:30:30 +01:00
getNodeType(id: string) {
return this.nodeRegistry.getNode(id);
}
getChildrenOfNode(node: Node) {
const children = [];
const stack = node.tmp?.children?.slice(0);
while (stack?.length) {
const child = stack.pop();
if (!child) continue;
children.push(child);
stack.push(...child.tmp?.children || []);
}
return children;
}
2024-03-11 19:37:58 +01:00
2024-03-13 14:30:30 +01:00
getNodesBetween(from: Node, to: Node): Node[] | undefined {
// < - - - - from
const toParents = this.getParentsOfNode(to);
// < - - - - from - - - - to
const fromParents = this.getParentsOfNode(from);
if (toParents.includes(from)) {
return toParents.splice(toParents.indexOf(from));
} else if (fromParents.includes(to)) {
return [...fromParents.splice(fromParents.indexOf(to)), from];
2024-03-13 14:30:30 +01:00
} else {
// these two nodes are not connected
return;
}
}
removeNode(node: Node) {
const edges = this._edges.filter((edge) => edge[0].id !== node.id && edge[2].id !== node.id);
this.edges.set(edges);
2024-03-13 14:30:30 +01:00
this.nodes.update((nodes) => {
nodes.delete(node.id);
return nodes;
});
this.save();
2024-03-13 14:30:30 +01:00
}
createEdge(from: Node, fromSocket: number, to: Node, toSocket: string) {
const existingEdges = this.getEdgesToNode(to);
// check if this exact edge already exists
const existingEdge = existingEdges.find(e => e[0].id === from.id && e[1] === fromSocket && e[3] === toSocket);
if (existingEdge) {
console.log("Edge already exists");
console.log(existingEdge)
return;
};
// check if socket types match
const fromSocketType = from.tmp?.type?.outputs?.[fromSocket];
const toSocketType = to.tmp?.type?.inputs?.[toSocket]?.type;
if (fromSocketType !== toSocketType) {
console.error(`Socket types do not match: ${fromSocketType} !== ${toSocketType}`);
return;
}
this.edges.update((edges) => {
return [...edges.filter(e => e[2].id !== to.id || e[3] !== toSocket), [from, fromSocket, to, toSocket]];
});
2024-03-13 16:18:48 +01:00
this.save();
}
save() {
2024-03-19 16:56:42 +01:00
this.emit("save", this.serialize());
2024-03-13 16:18:48 +01:00
this.history.save();
}
getParentsOfNode(node: Node) {
const parents = [];
const stack = node.tmp?.parents?.slice(0);
while (stack?.length) {
const parent = stack.pop();
if (!parent) continue;
parents.push(parent);
stack.push(...parent.tmp?.parents || []);
}
2024-03-13 14:30:30 +01:00
return parents.reverse();
}
getPossibleSockets({ node, index }: Socket): [Node, string | number][] {
const nodeType = node?.tmp?.type;
2024-03-11 19:37:58 +01:00
if (!nodeType) return [];
const sockets: [Node, string | number][] = []
// if index is a string, we are an input looking for outputs
if (typeof index === "string") {
// filter out self and child nodes
const children = new Set(this.getChildrenOfNode(node).map(n => n.id));
const nodes = this.getAllNodes().filter(n => n.id !== node.id && !children.has(n.id));
2024-03-11 19:37:58 +01:00
const ownType = nodeType?.inputs?.[index].type;
2024-03-11 19:37:58 +01:00
for (const node of nodes) {
const nodeType = node?.tmp?.type;
2024-03-11 19:37:58 +01:00
const inputs = nodeType?.outputs;
if (!inputs) continue;
for (let index = 0; index < inputs.length; index++) {
if (inputs[index] === ownType) {
sockets.push([node, index]);
}
}
}
} else if (typeof index === "number") {
// if index is a number, we are an output looking for inputs
2024-03-11 19:37:58 +01:00
// filter out self and parent nodes
const parents = new Set(this.getParentsOfNode(node).map(n => n.id));
const nodes = this.getAllNodes().filter(n => n.id !== node.id && !parents.has(n.id));
// get edges from this socket
const edges = new Map(this.getEdgesFromNode(node).filter(e => e[1] === index).map(e => [e[2].id, e[3]]));
const ownType = nodeType.outputs?.[index];
2024-03-11 19:37:58 +01:00
for (const node of nodes) {
const inputs = node?.tmp?.type?.inputs;
if (!inputs) continue;
for (const key in inputs) {
if (inputs[key].type === ownType && edges.get(node.id) !== key) {
sockets.push([node, key]);
2024-03-11 19:37:58 +01:00
}
}
2024-03-11 19:37:58 +01:00
}
}
return sockets;
}
2024-03-11 22:00:16 +01:00
removeEdge(edge: Edge) {
const id0 = edge[0].id;
const sid0 = edge[1];
const id2 = edge[2].id;
const sid2 = edge[3];
this.edges.update((edges) => {
return edges.filter((e) => e[0].id !== id0 || e[1] !== sid0 || e[2].id !== id2 || e[3] !== sid2);
});
this.save();
2024-03-11 22:00:16 +01:00
}
getEdgesToNode(node: Node) {
return this._edges
.filter((edge) => edge[2].id === node.id)
2024-03-06 18:31:06 +01:00
.map((edge) => {
const from = this.getNode(edge[0].id);
const to = this.getNode(edge[2].id);
2024-03-06 18:31:06 +01:00
if (!from || !to) return;
2024-03-11 22:00:16 +01:00
return [from, edge[1], to, edge[3]] as const;
2024-03-06 18:31:06 +01:00
})
2024-03-11 22:00:16 +01:00
.filter(Boolean) as unknown as [Node, number, Node, string][];
2024-03-06 18:31:06 +01:00
}
2024-03-11 22:00:16 +01:00
getEdgesFromNode(node: Node) {
return this._edges
.filter((edge) => edge[0].id === node.id)
2024-03-11 22:00:16 +01:00
.map((edge) => {
const from = this.getNode(edge[0].id);
const to = this.getNode(edge[2].id);
2024-03-11 22:00:16 +01:00
if (!from || !to) return;
return [from, edge[1], to, edge[3]] as const;
})
.filter(Boolean) as unknown as [Node, number, Node, string][];
}
2024-03-06 18:31:06 +01:00
createTemplate<T extends keyof typeof templates>(template: T, ...args: Parameters<typeof templates[T]>) {
switch (template) {
case "grid":
return templates.grid(args?.[0] || 5, args?.[1] || 5);
case "tree":
return templates.tree(args?.[0] || 4);
default:
throw new Error(`Template not found: ${template}`);
2024-03-06 18:31:06 +01:00
}
2024-03-11 22:00:16 +01:00
2024-03-06 18:31:06 +01:00
}
}