From 46a486c6943d57e611225dd43e16f56a6fa9802f Mon Sep 17 00:00:00 2001 From: Christian Byrne Date: Fri, 1 Aug 2025 21:40:29 -0700 Subject: [PATCH] [fix] Optimize subgraph serialization to exclude unused definitions (#1185) --- src/LGraph.ts | 11 +- src/subgraph/subgraphUtils.ts | 50 +++++++ test/subgraph/subgraphSerialization.test.ts | 145 +++++++++++++++++++ test/subgraph/subgraphUtils.test.ts | 147 ++++++++++++++++++++ 4 files changed, 351 insertions(+), 2 deletions(-) create mode 100644 test/subgraph/subgraphSerialization.test.ts create mode 100644 test/subgraph/subgraphUtils.test.ts diff --git a/src/LGraph.ts b/src/LGraph.ts index 37f59d887..10eaf43e8 100644 --- a/src/LGraph.ts +++ b/src/LGraph.ts @@ -36,7 +36,7 @@ import { stringOrEmpty } from "./strings" import { type GraphOrSubgraph, Subgraph } from "./subgraph/Subgraph" import { SubgraphInput } from "./subgraph/SubgraphInput" import { SubgraphOutput } from "./subgraph/SubgraphOutput" -import { getBoundaryLinks, groupResolvedByOutput, mapSubgraphInputsAndLinks, mapSubgraphOutputsAndLinks, multiClone, splitPositionables } from "./subgraph/subgraphUtils" +import { findUsedSubgraphIds, getBoundaryLinks, groupResolvedByOutput, mapSubgraphInputsAndLinks, mapSubgraphOutputsAndLinks, multiClone, splitPositionables } from "./subgraph/subgraphUtils" import { Alignment, LGraphEventMode } from "./types/globalEnums" import { getAllNestedItems } from "./utils/collections" @@ -1674,7 +1674,14 @@ export class LGraph implements LinkNetwork, BaseLGraph, Serialisable x.asSerialisable()) } + const usedSubgraphIds = findUsedSubgraphIds(this, this._subgraphs) + const usedSubgraphs = [...this._subgraphs.values()] + .filter(subgraph => usedSubgraphIds.has(subgraph.id)) + .map(x => x.asSerialisable()) + + if (usedSubgraphs.length > 0) { + data.definitions = { subgraphs: usedSubgraphs } + } } this.onSerialize?.(data) diff --git a/src/subgraph/subgraphUtils.ts b/src/subgraph/subgraphUtils.ts index 078c9167e..d8036b9a8 100644 --- a/src/subgraph/subgraphUtils.ts +++ b/src/subgraph/subgraphUtils.ts @@ -1,8 +1,10 @@ +import type { GraphOrSubgraph } from "./Subgraph" import type { SubgraphInput } from "./SubgraphInput" import type { SubgraphOutput } from "./SubgraphOutput" import type { INodeInputSlot, INodeOutputSlot, Positionable } from "@/interfaces" import type { LGraph } from "@/LGraph" import type { ISerialisedNode, SerialisableLLink, SubgraphIO } from "@/types/serialisation" +import type { UUID } from "@/utils/uuid" import { SUBGRAPH_INPUT_ID, SUBGRAPH_OUTPUT_ID } from "@/constants" import { LGraphGroup } from "@/LGraphGroup" @@ -339,6 +341,54 @@ export function mapSubgraphOutputsAndLinks(resolvedOutputLinks: ResolvedConnecti return outputs } +/** + * Collects all subgraph IDs used directly in a single graph (non-recursive). + * @param graph The graph to check for subgraph nodes + * @returns Set of subgraph IDs used in this graph + */ +export function getDirectSubgraphIds(graph: GraphOrSubgraph): Set { + const subgraphIds = new Set() + + for (const node of graph._nodes) { + if (node.isSubgraphNode()) { + subgraphIds.add(node.type) + } + } + + return subgraphIds +} + +/** + * Collects all subgraph IDs referenced in a graph hierarchy using BFS. + * @param rootGraph The graph to start from + * @param subgraphRegistry Map of all available subgraphs + * @returns Set of all subgraph IDs found + */ +export function findUsedSubgraphIds( + rootGraph: GraphOrSubgraph, + subgraphRegistry: Map, +): Set { + const usedSubgraphIds = new Set() + const toVisit: GraphOrSubgraph[] = [rootGraph] + + while (toVisit.length > 0) { + const graph = toVisit.shift()! + const directIds = getDirectSubgraphIds(graph) + + for (const id of directIds) { + if (!usedSubgraphIds.has(id)) { + usedSubgraphIds.add(id) + const subgraph = subgraphRegistry.get(id) + if (subgraph) { + toVisit.push(subgraph) + } + } + } + } + + return usedSubgraphIds +} + /** * Type guard to check if a slot is a SubgraphInput. * @param slot The slot to check diff --git a/test/subgraph/subgraphSerialization.test.ts b/test/subgraph/subgraphSerialization.test.ts new file mode 100644 index 000000000..4f75f0a62 --- /dev/null +++ b/test/subgraph/subgraphSerialization.test.ts @@ -0,0 +1,145 @@ +import { describe, expect, it } from "vitest" + +import { LGraph } from "@/litegraph" + +import { createTestSubgraph, createTestSubgraphNode } from "./fixtures/subgraphHelpers" + +describe("Subgraph Serialization", () => { + describe("LGraph.asSerialisable", () => { + it("should not include unused subgraph definitions", () => { + const rootGraph = new LGraph() + + // Create subgraphs + const usedSubgraph = createTestSubgraph({ name: "Used Subgraph" }) + const unusedSubgraph = createTestSubgraph({ name: "Unused Subgraph" }) + + // Add both to registry + rootGraph._subgraphs.set(usedSubgraph.id, usedSubgraph) + rootGraph._subgraphs.set(unusedSubgraph.id, unusedSubgraph) + + // Only add node for used subgraph + const node = createTestSubgraphNode(usedSubgraph) + rootGraph.add(node) + + // Serialize + const serialized = rootGraph.asSerialisable() + + // Check that only used subgraph is included + expect(serialized.definitions?.subgraphs).toBeDefined() + expect(serialized.definitions!.subgraphs!.length).toBe(1) + expect(serialized.definitions!.subgraphs![0].id).toBe(usedSubgraph.id) + expect(serialized.definitions!.subgraphs![0].name).toBe("Used Subgraph") + }) + + it("should include nested subgraphs", () => { + const rootGraph = new LGraph() + + // Create nested subgraphs + const level1Subgraph = createTestSubgraph({ name: "Level 1" }) + const level2Subgraph = createTestSubgraph({ name: "Level 2" }) + + // Add to registry + rootGraph._subgraphs.set(level1Subgraph.id, level1Subgraph) + rootGraph._subgraphs.set(level2Subgraph.id, level2Subgraph) + + // Add level1 to root + const level1Node = createTestSubgraphNode(level1Subgraph) + rootGraph.add(level1Node) + + // Add level2 to level1 + const level2Node = createTestSubgraphNode(level2Subgraph) + level1Subgraph.add(level2Node) + + // Serialize + const serialized = rootGraph.asSerialisable() + + // Both subgraphs should be included + expect(serialized.definitions?.subgraphs).toBeDefined() + expect(serialized.definitions!.subgraphs!.length).toBe(2) + + const ids = serialized.definitions!.subgraphs!.map(s => s.id) + expect(ids).toContain(level1Subgraph.id) + expect(ids).toContain(level2Subgraph.id) + }) + + it("should handle circular subgraph references", () => { + const rootGraph = new LGraph() + + // Create two subgraphs that reference each other + const subgraph1 = createTestSubgraph({ name: "Subgraph 1" }) + const subgraph2 = createTestSubgraph({ name: "Subgraph 2" }) + + // Add to registry + rootGraph._subgraphs.set(subgraph1.id, subgraph1) + rootGraph._subgraphs.set(subgraph2.id, subgraph2) + + // Add subgraph1 to root + const node1 = createTestSubgraphNode(subgraph1) + rootGraph.add(node1) + + // Add subgraph2 to subgraph1 + const node2 = createTestSubgraphNode(subgraph2) + subgraph1.add(node2) + + // Add subgraph1 to subgraph2 (circular) + const node3 = createTestSubgraphNode(subgraph1, { id: 3 }) + subgraph2.add(node3) + + // Serialize - should not hang + const serialized = rootGraph.asSerialisable() + + // Both should be included + expect(serialized.definitions?.subgraphs).toBeDefined() + expect(serialized.definitions!.subgraphs!.length).toBe(2) + }) + + it("should handle empty subgraph registry", () => { + const rootGraph = new LGraph() + + // Serialize with no subgraphs + const serialized = rootGraph.asSerialisable() + + // Should not include definitions + expect(serialized.definitions).toBeUndefined() + }) + + it("should only serialize from root graph", () => { + const rootGraph = new LGraph() + const subgraph = createTestSubgraph({ name: "Parent Subgraph" }) + + // Add subgraph to root registry + rootGraph._subgraphs.set(subgraph.id, subgraph) + + // Try to serialize from subgraph (not root) + const serialized = subgraph.asSerialisable() + + // Should not include definitions since it's not the root + expect(serialized.definitions).toBeUndefined() + }) + + it("should handle multiple instances of same subgraph", () => { + const rootGraph = new LGraph() + const subgraph = createTestSubgraph({ name: "Reused Subgraph" }) + + // Add to registry + rootGraph._subgraphs.set(subgraph.id, subgraph) + + // Add multiple instances + const node1 = createTestSubgraphNode(subgraph, { id: 1 }) + const node2 = createTestSubgraphNode(subgraph, { id: 2 }) + const node3 = createTestSubgraphNode(subgraph, { id: 3 }) + + rootGraph.add(node1) + rootGraph.add(node2) + rootGraph.add(node3) + + // Serialize + const serialized = rootGraph.asSerialisable() + + // Should only include one definition + expect(serialized.definitions?.subgraphs).toBeDefined() + expect(serialized.definitions!.subgraphs!.length).toBe(1) + expect(serialized.definitions!.subgraphs![0].id).toBe(subgraph.id) + }) + }) +}) diff --git a/test/subgraph/subgraphUtils.test.ts b/test/subgraph/subgraphUtils.test.ts new file mode 100644 index 000000000..e8175faec --- /dev/null +++ b/test/subgraph/subgraphUtils.test.ts @@ -0,0 +1,147 @@ +import type { UUID } from "@/utils/uuid" + +import { describe, expect, it } from "vitest" + +import { LGraph } from "@/litegraph" +import { + findUsedSubgraphIds, + getDirectSubgraphIds, +} from "@/subgraph/subgraphUtils" + +import { createTestSubgraph, createTestSubgraphNode } from "./fixtures/subgraphHelpers" + +describe("subgraphUtils", () => { + describe("getDirectSubgraphIds", () => { + it("should return empty set for graph with no subgraph nodes", () => { + const graph = new LGraph() + const result = getDirectSubgraphIds(graph) + expect(result.size).toBe(0) + }) + + it("should find single subgraph node", () => { + const graph = new LGraph() + const subgraph = createTestSubgraph() + const subgraphNode = createTestSubgraphNode(subgraph) + graph.add(subgraphNode) + + const result = getDirectSubgraphIds(graph) + expect(result.size).toBe(1) + expect(result.has(subgraph.id)).toBe(true) + }) + + it("should find multiple unique subgraph nodes", () => { + const graph = new LGraph() + const subgraph1 = createTestSubgraph({ name: "Subgraph 1" }) + const subgraph2 = createTestSubgraph({ name: "Subgraph 2" }) + + const node1 = createTestSubgraphNode(subgraph1) + const node2 = createTestSubgraphNode(subgraph2) + + graph.add(node1) + graph.add(node2) + + const result = getDirectSubgraphIds(graph) + expect(result.size).toBe(2) + expect(result.has(subgraph1.id)).toBe(true) + expect(result.has(subgraph2.id)).toBe(true) + }) + + it("should return unique IDs when same subgraph is used multiple times", () => { + const graph = new LGraph() + const subgraph = createTestSubgraph() + + const node1 = createTestSubgraphNode(subgraph, { id: 1 }) + const node2 = createTestSubgraphNode(subgraph, { id: 2 }) + + graph.add(node1) + graph.add(node2) + + const result = getDirectSubgraphIds(graph) + expect(result.size).toBe(1) + expect(result.has(subgraph.id)).toBe(true) + }) + }) + + describe("findUsedSubgraphIds", () => { + it("should handle graph with no subgraphs", () => { + const graph = new LGraph() + const registry = new Map() + + const result = findUsedSubgraphIds(graph, registry) + expect(result.size).toBe(0) + }) + + it("should find nested subgraphs", () => { + const rootGraph = new LGraph() + const subgraph1 = createTestSubgraph({ name: "Level 1" }) + const subgraph2 = createTestSubgraph({ name: "Level 2" }) + + // Add subgraph1 node to root + const node1 = createTestSubgraphNode(subgraph1) + rootGraph.add(node1) + + // Add subgraph2 node inside subgraph1 + const node2 = createTestSubgraphNode(subgraph2) + subgraph1.add(node2) + + const registry = new Map([ + [subgraph1.id, subgraph1], + [subgraph2.id, subgraph2], + ]) + + const result = findUsedSubgraphIds(rootGraph, registry) + expect(result.size).toBe(2) + expect(result.has(subgraph1.id)).toBe(true) + expect(result.has(subgraph2.id)).toBe(true) + }) + + it("should handle circular references without infinite loop", () => { + const rootGraph = new LGraph() + const subgraph1 = createTestSubgraph({ name: "Subgraph 1" }) + const subgraph2 = createTestSubgraph({ name: "Subgraph 2" }) + + // Add subgraph1 to root + const node1 = createTestSubgraphNode(subgraph1) + rootGraph.add(node1) + + // Add subgraph2 to subgraph1 + const node2 = createTestSubgraphNode(subgraph2) + subgraph1.add(node2) + + // Add subgraph1 to subgraph2 (circular reference) + const node3 = createTestSubgraphNode(subgraph1, { id: 3 }) + subgraph2.add(node3) + + const registry = new Map([ + [subgraph1.id, subgraph1], + [subgraph2.id, subgraph2], + ]) + + const result = findUsedSubgraphIds(rootGraph, registry) + expect(result.size).toBe(2) + expect(result.has(subgraph1.id)).toBe(true) + expect(result.has(subgraph2.id)).toBe(true) + }) + + it("should handle missing subgraphs in registry gracefully", () => { + const rootGraph = new LGraph() + const subgraph1 = createTestSubgraph({ name: "Subgraph 1" }) + const subgraph2 = createTestSubgraph({ name: "Subgraph 2" }) + + // Add both subgraph nodes + const node1 = createTestSubgraphNode(subgraph1) + const node2 = createTestSubgraphNode(subgraph2) + + rootGraph.add(node1) + rootGraph.add(node2) + + // Only register subgraph1 + const registry = new Map([[subgraph1.id, subgraph1]]) + + const result = findUsedSubgraphIds(rootGraph, registry) + expect(result.size).toBe(2) + expect(result.has(subgraph1.id)).toBe(true) + expect(result.has(subgraph2.id)).toBe(true) // Still found, just can't recurse into it + }) + }) +})