|
1 | 1 | import type { LGraph, LGraphNode, Subgraph } from '@comfyorg/litegraph' |
2 | 2 |
|
3 | | -import type { NodeLocatorId } from '@/types/nodeIdentification' |
| 3 | +import type { NodeExecutionId, NodeLocatorId } from '@/types/nodeIdentification' |
4 | 4 | import { parseNodeLocatorId } from '@/types/nodeIdentification' |
5 | 5 |
|
6 | 6 | import { isSubgraphIoNode } from './typeGuardUtil' |
@@ -351,3 +351,106 @@ export function mapSubgraphNodes<T>( |
351 | 351 | export function getAllNonIoNodesInSubgraph(subgraph: Subgraph): LGraphNode[] { |
352 | 352 | return subgraph.nodes.filter((node) => !isSubgraphIoNode(node)) |
353 | 353 | } |
| 354 | + |
| 355 | +/** |
| 356 | + * Performs depth-first traversal of nodes and their subgraphs. |
| 357 | + * Generic visitor pattern that can be used for various node processing tasks. |
| 358 | + * |
| 359 | + * @param nodes - Starting nodes for traversal |
| 360 | + * @param visitor - Function called for each node with its context |
| 361 | + * @param expandSubgraphs - Whether to traverse into subgraph nodes (default: true) |
| 362 | + */ |
| 363 | +export function traverseNodesDepthFirst<T>( |
| 364 | + nodes: LGraphNode[], |
| 365 | + visitor: (node: LGraphNode, context: T) => T, |
| 366 | + initialContext: T, |
| 367 | + expandSubgraphs: boolean = true |
| 368 | +): void { |
| 369 | + type StackItem = { node: LGraphNode; context: T } |
| 370 | + const stack: StackItem[] = [] |
| 371 | + |
| 372 | + // Initialize stack with starting nodes |
| 373 | + for (const node of nodes) { |
| 374 | + stack.push({ node, context: initialContext }) |
| 375 | + } |
| 376 | + |
| 377 | + // Process stack iteratively (DFS) |
| 378 | + while (stack.length > 0) { |
| 379 | + const { node, context } = stack.pop()! |
| 380 | + |
| 381 | + // Visit node and get updated context for children |
| 382 | + const childContext = visitor(node, context) |
| 383 | + |
| 384 | + // If it's a subgraph and we should expand, add children to stack |
| 385 | + if (expandSubgraphs && node.isSubgraphNode?.() && node.subgraph) { |
| 386 | + // Process children in reverse order to maintain left-to-right DFS processing |
| 387 | + // when popping from stack (LIFO). Iterate backwards to avoid array reversal. |
| 388 | + const children = node.subgraph.nodes |
| 389 | + for (let i = children.length - 1; i >= 0; i--) { |
| 390 | + stack.push({ node: children[i], context: childContext }) |
| 391 | + } |
| 392 | + } |
| 393 | + } |
| 394 | +} |
| 395 | + |
| 396 | +/** |
| 397 | + * Collects nodes with custom data during depth-first traversal. |
| 398 | + * Generic collector that can gather any type of data per node. |
| 399 | + * |
| 400 | + * @param nodes - Starting nodes for traversal |
| 401 | + * @param collector - Function that returns data to collect for each node |
| 402 | + * @param contextBuilder - Function that builds context for child nodes |
| 403 | + * @param expandSubgraphs - Whether to traverse into subgraph nodes |
| 404 | + * @returns Array of collected data |
| 405 | + */ |
| 406 | +export function collectFromNodes<T, C>( |
| 407 | + nodes: LGraphNode[], |
| 408 | + collector: (node: LGraphNode, context: C) => T | null, |
| 409 | + contextBuilder: (node: LGraphNode, parentContext: C) => C, |
| 410 | + initialContext: C, |
| 411 | + expandSubgraphs: boolean = true |
| 412 | +): T[] { |
| 413 | + const results: T[] = [] |
| 414 | + |
| 415 | + traverseNodesDepthFirst( |
| 416 | + nodes, |
| 417 | + (node, context) => { |
| 418 | + const data = collector(node, context) |
| 419 | + if (data !== null) { |
| 420 | + results.push(data) |
| 421 | + } |
| 422 | + return contextBuilder(node, context) |
| 423 | + }, |
| 424 | + initialContext, |
| 425 | + expandSubgraphs |
| 426 | + ) |
| 427 | + |
| 428 | + return results |
| 429 | +} |
| 430 | + |
| 431 | +/** |
| 432 | + * Collects execution IDs for selected nodes and all their descendants. |
| 433 | + * Uses the generic DFS traversal with optimized string building. |
| 434 | + * |
| 435 | + * @param selectedNodes - The selected nodes to process |
| 436 | + * @returns Array of execution IDs for selected nodes and all nodes within selected subgraphs |
| 437 | + */ |
| 438 | +export function getExecutionIdsForSelectedNodes( |
| 439 | + selectedNodes: LGraphNode[] |
| 440 | +): NodeExecutionId[] { |
| 441 | + return collectFromNodes( |
| 442 | + selectedNodes, |
| 443 | + // Collector: build execution ID for each node |
| 444 | + (node, parentExecutionId: string): NodeExecutionId => { |
| 445 | + const nodeId = String(node.id) |
| 446 | + return parentExecutionId ? `${parentExecutionId}:${nodeId}` : nodeId |
| 447 | + }, |
| 448 | + // Context builder: pass execution ID to children |
| 449 | + (node, parentExecutionId: string) => { |
| 450 | + const nodeId = String(node.id) |
| 451 | + return parentExecutionId ? `${parentExecutionId}:${nodeId}` : nodeId |
| 452 | + }, |
| 453 | + '', // Initial context: empty parent execution ID |
| 454 | + true // Expand subgraphs |
| 455 | + ) |
| 456 | +} |
0 commit comments