Make 'removeOrphans' optional on replaceNodeIdsConnectedTo

The graph clean-up is a traversal, so there's a large linear performance hit if we clean-up on every modification to the graph. This makes it configurable to disable clean-up on replacement. Test Plan: `yarn mocha src/packages/core/graph` Pull Request: #330
atlassian-labs · Feb 5, 2025 · 8d1af44 · 8d1af44
1 parent c30a222
commit 8d1af44
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 26 deletions.
diff --git a/packages/core/graph/src/Graph.js b/packages/core/graph/src/Graph.js
@@ -330,6 +330,7 @@ export default class Graph<TNode, TEdgeType: number = NullEdgeType> {
     toNodeIds: $ReadOnlyArray<NodeId>,
     replaceFilter?: null | ((NodeId) => boolean),
     type?: TEdgeType | NullEdgeType = NULL_EDGE_TYPE,
+    removeOrphans: boolean = true,
   ): void {
     this._assertHasNodeId(fromNodeId);
 
@@ -348,7 +349,7 @@ export default class Graph<TNode, TEdgeType: number = NullEdgeType> {
     }
 
     for (let child of childrenToRemove) {
-      this._removeEdge(fromNodeId, child, type);
+      this._removeEdge(fromNodeId, child, type, removeOrphans);
     }
   }
 

diff --git a/packages/core/graph/test/Graph.test.js b/packages/core/graph/test/Graph.test.js
@@ -37,13 +37,6 @@ describe('Graph', () => {
     }, /Does not have node/);
   });
 
-  it("errors if replaceNodeIdsConnectedTo is called with a node that doesn't belong", () => {
-    let graph = new Graph();
-    assert.throws(() => {
-      graph.replaceNodeIdsConnectedTo(toNodeId(-1), []);
-    }, /Does not have node/);
-  });
-
   it("errors when adding an edge to a node that doesn't exist", () => {
     let graph = new Graph();
     let node = graph.addNode({});
@@ -274,26 +267,63 @@ describe('Graph', () => {
     }
   });
 
-  it("replaceNodeIdsConnectedTo should update a node's downstream nodes", () => {
-    let graph = new Graph();
-    let nodeA = graph.addNode('a');
-    graph.setRootNodeId(nodeA);
-    let nodeB = graph.addNode('b');
-    let nodeC = graph.addNode('c');
-    graph.addEdge(nodeA, nodeB);
-    graph.addEdge(nodeA, nodeC);
+  describe('replaceNodeIdsConnectedTo', () => {
+    it("errors if replaceNodeIdsConnectedTo is called with a node that doesn't belong", () => {
+      let graph = new Graph();
+      assert.throws(() => {
+        graph.replaceNodeIdsConnectedTo(toNodeId(-1), []);
+      }, /Does not have node/);
+    });
 
-    let nodeD = graph.addNode('d');
-    graph.replaceNodeIdsConnectedTo(nodeA, [nodeB, nodeD]);
+    it("replaceNodeIdsConnectedTo should update a node's downstream nodes", () => {
+      let graph = new Graph();
+      let nodeA = graph.addNode('a');
+      graph.setRootNodeId(nodeA);
+      let nodeB = graph.addNode('b');
+      let nodeC = graph.addNode('c');
+      graph.addEdge(nodeA, nodeB);
+      graph.addEdge(nodeA, nodeC);
+
+      let nodeD = graph.addNode('d');
+      graph.replaceNodeIdsConnectedTo(nodeA, [nodeB, nodeD]);
+
+      assert(graph.hasNode(nodeA));
+      assert(graph.hasNode(nodeB));
+      assert(!graph.hasNode(nodeC)); // orphan removed
+      assert(graph.hasNode(nodeD));
+      assert.deepEqual(Array.from(graph.getAllEdges()), [
+        {from: nodeA, to: nodeB, type: 1},
+        {from: nodeA, to: nodeD, type: 1},
+      ]);
+    });
 
-    assert(graph.hasNode(nodeA));
-    assert(graph.hasNode(nodeB));
-    assert(!graph.hasNode(nodeC));
-    assert(graph.hasNode(nodeD));
-    assert.deepEqual(Array.from(graph.getAllEdges()), [
-      {from: nodeA, to: nodeB, type: 1},
-      {from: nodeA, to: nodeD, type: 1},
-    ]);
+    it('replaceNodeIdsConnectedTo does not remove orphan nodes if flag is false', () => {
+      let graph = new Graph();
+      let nodeA = graph.addNode('a');
+      graph.setRootNodeId(nodeA);
+      let nodeB = graph.addNode('b');
+      let nodeC = graph.addNode('c');
+      graph.addEdge(nodeA, nodeB);
+      graph.addEdge(nodeA, nodeC);
+
+      let nodeD = graph.addNode('d');
+      graph.replaceNodeIdsConnectedTo(
+        nodeA,
+        [nodeB, nodeD],
+        () => true,
+        1,
+        false,
+      );
+
+      assert(graph.hasNode(nodeA));
+      assert(graph.hasNode(nodeB));
+      assert(graph.hasNode(nodeC)); // orphan kept
+      assert(graph.hasNode(nodeD));
+      assert.deepEqual(Array.from(graph.getAllEdges()), [
+        {from: nodeA, to: nodeB, type: 1},
+        {from: nodeA, to: nodeD, type: 1},
+      ]);
+    });
   });
 
   it('traverses along edge types if a filter is given', () => {