From 9ba072a1fe2583a259393cbb7272bd72f84a020b Mon Sep 17 00:00:00 2001
From: Pedro Yamada <tacla.yamada@gmail.com>
Date: Tue, 3 Dec 2024 23:39:44 +0000
Subject: [PATCH] Optimize BundleGraph::getAssetWithDependency

Prevents creating temporary arrays when finding the asset linked to a
dependency. This is a relatively hot path in the bundler and can
improve its performance.
---
 packages/core/core/src/BundleGraph.js         | 21 +++++++++++++------
 packages/core/graph/src/AdjacencyList.js      |  6 ++++--
 packages/core/graph/src/Graph.js              |  2 +-
 .../core/graph/test/AdjacencyList.test.js     | 20 ++++++++++++++++++
 4 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/packages/core/core/src/BundleGraph.js b/packages/core/core/src/BundleGraph.js
index 8ab125629..912ac3dbd 100644
--- a/packages/core/core/src/BundleGraph.js
+++ b/packages/core/core/src/BundleGraph.js
@@ -1773,14 +1773,23 @@ export default class BundleGraph {
       return null;
     }
 
-    let res = this._graph.getNodeIdsConnectedTo(
+    let res = null;
+    let count = 0;
+    this._graph.forEachNodeIdConnectedTo(
       this._graph.getNodeIdByContentKey(dep.id),
+      (node) => {
+        res = node;
+        count += 1;
+        if (count > 1) {
+          throw new Error(
+            'Expected a single asset to be connected to a dependency',
+          );
+        }
+      },
     );
-    invariant(
-      res.length <= 1,
-      'Expected a single asset to be connected to a dependency',
-    );
-    let resNode = this._graph.getNode(res[0]);
+
+    invariant(res != null, 'Expected an asset to be connected to a dependency');
+    let resNode = this._graph.getNode(res);
     if (resNode?.type === 'asset') {
       return resNode.value;
     }
diff --git a/packages/core/graph/src/AdjacencyList.js b/packages/core/graph/src/AdjacencyList.js
index f95f2c41f..2fc7d8090 100644
--- a/packages/core/graph/src/AdjacencyList.js
+++ b/packages/core/graph/src/AdjacencyList.js
@@ -619,13 +619,15 @@ export default class AdjacencyList<TEdgeType: number = 1> {
     }
   }
 
-  forEachNodeIdConnectedTo(to: NodeId, fn: (nodeId: NodeId) => void) {
+  forEachNodeIdConnectedTo(to: NodeId, fn: (nodeId: NodeId) => boolean | void) {
     let node = this.#nodes.head(to);
     while (node !== null) {
       let edge = this.#nodes.firstIn(node);
       while (edge !== null) {
         let from = this.#edges.from(edge);
-        fn(from);
+        if (fn(from)) {
+          return;
+        }
         edge = this.#edges.nextIn(edge);
       }
       node = this.#nodes.next(node);
diff --git a/packages/core/graph/src/Graph.js b/packages/core/graph/src/Graph.js
index 40f0e8190..7b5d10d86 100644
--- a/packages/core/graph/src/Graph.js
+++ b/packages/core/graph/src/Graph.js
@@ -163,7 +163,7 @@ export default class Graph<TNode, TEdgeType: number = 1> {
     return this.adjacencyList.hasEdge(from, to, type);
   }
 
-  forEachNodeIdConnectedTo(to: NodeId, fn: (nodeId: NodeId) => void) {
+  forEachNodeIdConnectedTo(to: NodeId, fn: (nodeId: NodeId) => boolean | void) {
     this._assertHasNodeId(to);
 
     this.adjacencyList.forEachNodeIdConnectedTo(to, fn);
diff --git a/packages/core/graph/test/AdjacencyList.test.js b/packages/core/graph/test/AdjacencyList.test.js
index dc96e5572..4d3e6a2fb 100644
--- a/packages/core/graph/test/AdjacencyList.test.js
+++ b/packages/core/graph/test/AdjacencyList.test.js
@@ -306,6 +306,26 @@ describe('AdjacencyList', () => {
       assert.deepEqual(nodeIds, [a, c]);
     });
 
+    it('stops traversal if the return value is true', () => {
+      const graph = new AdjacencyList();
+      const a = graph.addNode();
+      const b = graph.addNode();
+      const c = graph.addNode();
+      const d = graph.addNode();
+
+      graph.addEdge(a, d);
+      graph.addEdge(b, d);
+      graph.addEdge(c, d);
+
+      const nodeIds = [];
+      graph.forEachNodeIdConnectedTo(d, (nodeId) => {
+        nodeIds.push(nodeId);
+        return true;
+      });
+
+      assert.deepEqual(nodeIds, [a]);
+    });
+
     it('terminates if the graph is cyclic', () => {
       const graph = new AdjacencyList();
       const a = graph.addNode();