feat: topologicalSort now accepts option shouldFollow to restrict graph traversal

Author: JesusTheHun

src/Graph.spec.ts

@@ -1,6 +1,6 @@
 import { describe, expect, it } from 'vitest';
 
-import { depthFirstSearch, Graph } from './index.js';
+import { Graph } from './index.js';
 import { indegree } from './utils/indegree.js';
 import { outdegree } from './utils/outdegree.js';
 

src/algorithms/topologicalSort/topologicalSort.spec.ts

@@ -72,14 +72,23 @@ describe('topologicalSort', () => {
   });
 
   it('Should exclude source nodes with a cycle.', function () {
-    const graph = new Graph().addEdge('a', 'b').addEdge('b', 'c').addEdge('c', 'a');
+    const graph = new Graph<string, { type: string }>();
+    graph
+      .addEdge('a', 'b', undefined, { type: 'foo' })
+      .addEdge('b', 'c', undefined, { type: 'foo' })
+      .addEdge('c', 'a', undefined, { type: 'bar' });
+
     const sorted = topologicalSort(graph, {
       sourceNodes: ['a'],
-      includeSourceNodes: false,
+      includeSourceNodes: true,
+      shouldFollow: (source, target) =>
+        graph.getEdgeProperties(source, target).type === 'foo',
     });
-    expect(sorted.length).toEqual(2);
-    expect(sorted[0]).toEqual('b');
-    expect(sorted[1]).toEqual('c');
+
+    expect(sorted.length).toEqual(3);
+    expect(sorted[0]).toEqual('a');
+    expect(sorted[1]).toEqual('b');
+    expect(sorted[2]).toEqual('c');
   });
 
   it('Should exclude source nodes with multiple cycles.', function () {

src/algorithms/topologicalSort/topologicalSort.ts

@@ -8,7 +8,7 @@ import { Graph } from '../../Graph.js';
 import { NoInfer } from '../../types.js';
 import { depthFirstSearch } from '../depthFirstSearch/index.js';
 
-export type TopologicalSortOptions<Node> = {
+export type TopologicalSortOptions<Node, LinkProps> = {
   /**
    * Run the first on those nodes.
    * @default all the nodes of the graph.
@@ -20,11 +20,24 @@ export type TopologicalSortOptions<Node> = {
    * @default true
    */
   includeSourceNodes?: boolean;
+
+  /**
+   * A function that is executed to determine if the branch should be visited or not.
+   * @param source the current node
+   * @param target the next node to explore
+   * @param graph the graph instance being explored
+   * @returns boolean
+   */
+  shouldFollow?: (
+    source: NoInfer<Node>,
+    target: NoInfer<Node>,
+    graph: Graph<Node, LinkProps>,
+  ) => boolean;
 };
 
-export function topologicalSort<Node>(
-  graph: Graph<Node>,
-  opts: TopologicalSortOptions<NoInfer<Node>> = {},
+export function topologicalSort<Node, LinkProps>(
+  graph: Graph<Node, LinkProps>,
+  opts: TopologicalSortOptions<NoInfer<Node>, NoInfer<LinkProps>> = {},
 ) {
   return depthFirstSearch(graph, {
     ...opts,

src/utils/findNodes.spec.ts

@@ -1,8 +1,6 @@
-import { describe, it, expect } from 'vitest';
+import { describe, it } from 'vitest';
 import { Graph } from '../Graph.js';
 import { findNodes } from './findNodes.js';
-import { getFirstNode } from './getFirstNode.js';
-import { getNode } from './getNode.js';
 
 describe('findNodes', () => {
   type Node = { id: string; type: string };

src/utils/getFirstNode.spec.ts

@@ -1,7 +1,6 @@
-import { describe, it, expect } from 'vitest';
+import { describe, it } from 'vitest';
 import { Graph } from '../Graph.js';
 import { getFirstNode } from './getFirstNode.js';
-import { getNode } from './getNode.js';
 
 describe('getFirstNode', () => {
   type Node = { id: string; type: string };

src/utils/getNode.spec.ts

@@ -1,4 +1,4 @@
-import { describe, it, expect } from 'vitest';
+import { describe, it } from 'vitest';
 import { Graph } from '../Graph.js';
 import { getNode } from './getNode.js';