From 113746f90e7ed4653436cbe03e977fb1d6562146 Mon Sep 17 00:00:00 2001
From: Admon Sasson <admon.sasson@snyk.io>
Date: Wed, 25 Nov 2020 18:03:34 +0200
Subject: [PATCH] feat: add memoization to when converting a graph into a tree

---
 src/legacy/index.ts        |  9 ++++++-
 test/legacy/stress.test.ts | 50 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 1 deletion(-)
 create mode 100644 test/legacy/stress.test.ts

diff --git a/src/legacy/index.ts b/src/legacy/index.ts
index 56f5ee9..b51d6a4 100644
--- a/src/legacy/index.ts
+++ b/src/legacy/index.ts
@@ -285,8 +285,12 @@ async function buildSubtree(
   depGraph: types.DepGraphInternal,
   nodeId: string,
   eventLoopSpinner: EventLoopSpinner,
-  maybeDeduplicationSet: Set<string> | null | false = null, // false = disabled; null = not in deduplication scope yet
+  maybeDeduplicationSet: Set<string> | null | false = false, // false = disabled; null = not in deduplication scope yet
+  memoizationMap: Map<string, DepTree> = new Map(),
 ): Promise<DepTree> {
+  if (!maybeDeduplicationSet && memoizationMap.has(nodeId)) {
+    return memoizationMap.get(nodeId)!;
+  }
   const isRoot = nodeId === depGraph.rootNodeId;
   const nodePkg = depGraph.getNodePkg(nodeId);
   const nodeInfo = depGraph.getNode(nodeId);
@@ -302,6 +306,7 @@ async function buildSubtree(
 
   const depInstanceIds = depGraph.getNodeDepsNodeIds(nodeId);
   if (!depInstanceIds || depInstanceIds.length === 0) {
+    memoizationMap.set(nodeId, depTree);
     return depTree;
   }
 
@@ -326,6 +331,7 @@ async function buildSubtree(
       depInstId,
       eventLoopSpinner,
       maybeDeduplicationSet,
+      memoizationMap,
     );
     if (!subtree) {
       continue;
@@ -341,6 +347,7 @@ async function buildSubtree(
   if (eventLoopSpinner.isStarving()) {
     await eventLoopSpinner.spin();
   }
+  memoizationMap.set(nodeId, depTree);
   return depTree;
 }
 
diff --git a/test/legacy/stress.test.ts b/test/legacy/stress.test.ts
new file mode 100644
index 0000000..8c1e743
--- /dev/null
+++ b/test/legacy/stress.test.ts
@@ -0,0 +1,50 @@
+import * as depGraphLib from '../../src';
+import { graphToDepTree } from '../../src/legacy';
+
+const dependencyName = 'needle';
+
+async function generateLargeGraph(width: number) {
+  const builder = new depGraphLib.DepGraphBuilder(
+    { name: 'npm' },
+    { name: 'root', version: '1.2.3' },
+  );
+  const rootNodeId = 'root-node';
+
+  const deepDependency = { name: dependencyName, version: '1.2.3' };
+  const deepDependency2 = { name: dependencyName + 2, version: '1.2.3' };
+
+  builder.addPkgNode(deepDependency, dependencyName);
+  builder.addPkgNode(deepDependency2, deepDependency2.name);
+  builder.connectDep(rootNodeId, dependencyName);
+  builder.connectDep(deepDependency.name, deepDependency2.name);
+
+  for (let j = 0; j < width / 2; j++) {
+    const shallowName = `id-${j}`;
+    const shallowDependency = { name: shallowName, version: '1.2.3' };
+
+    builder.addPkgNode(shallowDependency, shallowName);
+    builder.connectDep(rootNodeId, shallowName);
+    builder.connectDep(shallowName, dependencyName);
+  }
+
+  for (let j = 0; j < width / 2; j++) {
+    const shallowName = `second-${j}`;
+    const shallowDependency = { name: shallowName, version: '1.2.3' };
+
+    builder.addPkgNode(shallowDependency, shallowName);
+    builder.connectDep(deepDependency2.name, shallowName);
+  }
+
+  return builder.build();
+}
+
+describe('stress tests', () => {
+  test('graphToDepTree() with memoization (without deduplicateWithinTopLevelDeps) succeed for large dep-graphs', async () => {
+    const graph = await generateLargeGraph(125000);
+
+    const depTree = await graphToDepTree(graph, 'gomodules', {
+      deduplicateWithinTopLevelDeps: false,
+    });
+    expect(depTree).toBeDefined();
+  });
+});