diff --git a/index.ts b/index.ts index 886a0b4..27f50ac 100644 --- a/index.ts +++ b/index.ts @@ -2,6 +2,8 @@ interface Identifiable { readonly id: string; } +type DAGVisitor = (node: T, parent: T | null, depth: number, index: number, total: number, context: C) => void; + class Node { constructor(readonly data: T, readonly dependencies = new Set()) {} @@ -112,6 +114,52 @@ class DAGraph { return reverseGraph; } + /** + * Traverses the graph in depth-first order and calls the visitor function for each node. + * + * @param visitor the visitor function to call for each node. The visitor signature is: + * (node: T, parent: T | null, depth: number, index: number, total: number, context: C) => void + * - node: the current node data + * - parent: the parent node data (null for roots) + * - depth: the depth of the current node (0 for roots) + * - index: the index of the current node among its siblings + * - total: the total number of siblings + * - context: the context object passed to traverse + * @param context the context object to pass to the visitor. + */ + traverse(visitor: DAGVisitor, context: C): void { + const outgoing = new Map(); + for (const node of this.nodesById.values()) { + for (const depId of node.dependencies) { + let children = outgoing.get(depId); + if (!children) { + children = []; + outgoing.set(depId, children); + } + children.push(node.id); + } + } + + const visitNode = (nodeId: string, parent: T | null, depth: number, index: number, total: number) => { + const node = this.nodesById.get(nodeId); + if (!node) { + return; + } + + visitor(node.data, parent, depth, index, total, context); + + const children = outgoing.get(nodeId) || []; + children.forEach((childId, i) => { + visitNode(childId, node.data, depth + 1, i, children.length); + }); + }; + + const roots = [...this.roots()]; + roots.forEach((root, i) => { + visitNode(root.id, null, 0, i, roots.length); + }); + } + private ensureNode(data: T): Node { let node = this.nodesById.get(data.id); if (node) { @@ -161,6 +209,6 @@ function createDAG(): DAGraph { return new DAGraph(); } -export type { DAGraph, Identifiable }; +export type { DAGraph, Identifiable, DAGVisitor }; export default createDAG; export { createDAG }; diff --git a/test/DAGraph.spec.ts b/test/DAGraph.spec.ts index 8ee64e6..3d4d024 100644 --- a/test/DAGraph.spec.ts +++ b/test/DAGraph.spec.ts @@ -157,33 +157,108 @@ describe('DAGraph', () => { expect([...dag.reverse().topologicalSort()]).toEqual(expected); - const myDAG = createDAG(); - myDAG.addEdge( - { - id: 'a', - name: 'my thing', - doSomething: () => { - console.log('A'); - } - }, - { - id: 'b', - name: 'my other thing', - doSomething: () => { - console.log('B'); - } - } - ); + // Additional verification for simple case + const smallDag = createDAG(); + const a = { id: 'A' }; + const b = { id: 'B' }; + smallDag.addEdge(a, b); + expect([...smallDag.reverse().topologicalSort()]).toEqual([b, a]); }); }); -}); -type MyThing = { - id: string; - name: string; - doSomething(): void; -}; + describe('traverse', () => { + test('should visit nodes in depth-first order', () => { + const dag = createDAG(); + const nodeA = aNode('A'); + const nodeB = aNode('B'); + const nodeC = aNode('C'); + const nodeD = aNode('D'); + + dag.addEdge(nodeA, nodeB); + dag.addEdge(nodeB, nodeC); + dag.addEdge(nodeA, nodeD); + + // Structure: + // A -> B -> C + // A -> D + + const visited: string[] = []; + dag.traverse((node, parent, depth) => { + visited.push(`${node.id}:${depth}`); + }, {}); + + // Since traverse uses roots() which iterates over Map.values(), the order of branches is insertion order related but implementation specific for Map. + // We just check that we visited everyone and depths are correct. + expect(visited).toIncludeAllMembers(['A:0', 'B:1', 'C:2', 'D:1']); + }); + + test('should pass context to visitor', () => { + const dag = createDAG(); + const nodeA = aNode('A'); + dag.addNode(nodeA); + + const context = { count: 0 }; + dag.traverse((_node, _parent, _depth, _index, _total, ctx) => { + ctx.count++; + }, context); + + expect(context.count).toBe(1); + }); + + test('should provide parent node', () => { + const dag = createDAG(); + const nodeA = aNode('A'); + const nodeB = aNode('B'); + dag.addEdge(nodeA, nodeB); + + const parents: Record = {}; + dag.traverse((node, parent) => { + parents[node.id] = parent ? parent.id : null; + }, {}); + + expect(parents['A']).toBeNull(); + expect(parents['B']).toBe('A'); + }); + + test('should handle multiple roots', () => { + const dag = createDAG(); + const nodeA = aNode('A'); + const nodeB = aNode('B'); + const nodeC = aNode('C'); + + dag.addNode(nodeA); + dag.addNode(nodeB); + dag.addEdge(nodeB, nodeC); + + const visited: string[] = []; + dag.traverse(node => visited.push(node.id), {}); + + expect(visited).toIncludeSameMembers(['A', 'B', 'C']); + }); + + test('should provide index and total siblings count', () => { + const dag = createDAG(); + const a = { id: 'A' }; + const b = { id: 'B' }; + const c = { id: 'C' }; + + // A -> B + // A -> C + dag.addEdge(a, b); + dag.addEdge(a, c); + + const visits: string[] = []; + dag.traverse((node, _parent, _depth, index, total) => { + visits.push(`${node.id}(${index}/${total})`); + }, {}); + + // Roots: A (0/1) + // Children of A: B (0/2), C (1/2) - deterministic order due to insertion + expect(visits).toIncludeSameMembers(['A(0/1)', 'B(0/2)', 'C(1/2)']); + }); + }); +}); -function aNode(): Identifiable { - return { id: uuid() }; +function aNode(id?: string): Identifiable { + return { id: id || uuid() }; }