fix(jsii): slow with deep dependency tree (#3294)

The `_loadDependencies` calls in jsii were taking a very long time with
a deep dependency tree (like in `cdk-watchful`). The reason for this is
another accidentally quadratic algorithm: every dependency would be
processed for every time it transitively occurred in any of the
dependency graph's dependencies.

For a reason I don't completely understand yet, the timings got a whole
lot worse recently in 1.49.0.

Refactor to add some short-circuiting: if a directory has already been
processed, no need to recurse into it yet again (the first recursion
pass on that tree will have gathered all the information that is
necessary already).

I stared long and hard at the `semver-intersect` call... but as far as I
could tell we would never be in a situation where `packageVersions`
would/could already contain a previous version, and the intersection
would actually do anything. I validated with a print statement.  As far
as I can tell this was dead code and I removed it.

Timings on the `cdk-watchful` project for reference:

```
$ time npx tsc
    5.88 real         9.62 user         0.38 sys

$ time npx jsii@1.47.0
    35.64 real        39.38 user         0.81 sys

$ time npx jsii@1.49.0
    334.26 real       239.60 user       141.71 sys

$ time /path/to/my/jsii
    6.11 real         9.82 user         0.42 sys
```

Diff a `diff` between a jsii assembly produced with the old and the
new version. No change.

packages/jsii/lib/project-info.ts

@@ -3,7 +3,6 @@ import * as fs from 'fs-extra';
 import * as log4js from 'log4js';
 import * as path from 'path';
 import * as semver from 'semver';
-import { intersect } from 'semver-intersect';
 import * as ts from 'typescript';
 
 import { JsiiDiagnostic } from './jsii-diagnostic';
@@ -120,23 +119,19 @@ export async function loadProjectInfo(
     }
   }
 
-  const transitiveAssemblies: { [name: string]: spec.Assembly } = {};
-  const assemblyCache = new Map<string, spec.Assembly>();
-  const dependencies = await _loadDependencies(
-    pkg.dependencies,
-    projectRoot,
-    transitiveAssemblies,
-    assemblyCache,
-    new Set<string>(Object.keys(bundleDependencies ?? {})),
-  );
-  const peerDependencies = await _loadDependencies(
-    pkg.peerDependencies,
-    projectRoot,
-    transitiveAssemblies,
-    assemblyCache,
+  const bundled = new Set(Object.keys(bundleDependencies ?? {}));
+  const dependencies: Record<string, string> = filterDictByKey(
+    pkg.dependencies ?? {},
+    (depName) => !bundled.has(depName),
   );
+  const peerDependencies: Record<string, string> = pkg.peerDependencies ?? {};
 
-  const transitiveDependencies = Object.values(transitiveAssemblies);
+  const resolver = new DependencyResolver();
+  const resolved = await resolver.discoverDependencyTree(projectRoot, {
+    ...dependencies,
+    ...peerDependencies,
+  });
+  const transitiveDependencies = resolver.assemblyClosure(resolved);
 
   const metadata = mergeMetadata(
     {
@@ -235,75 +230,123 @@ function _guessRepositoryType(url: string): string {
   );
 }
 
-async function _loadDependencies(
-  dependencies: { [name: string]: string } | undefined,
-  searchPath: string,
-  transitiveAssemblies: { [name: string]: spec.Assembly },
-  assemblyCache: Map<string, spec.Assembly>,
-  bundled = new Set<string>(),
-): Promise<{ [name: string]: string }> {
-  if (!dependencies) {
-    return {};
+interface DependencyInfo {
+  readonly assembly: spec.Assembly;
+  readonly resolvedDependencies: Record<string, string>;
+}
+
+class DependencyResolver {
+  private readonly cache = new Map<string, DependencyInfo>();
+
+  /**
+   * Discover the dependency tree starting at 'root', validating versions as we go along
+   *
+   * This primes the data structures in this class and should be called first.
+   *
+   * Return the resolved jsii dependency paths
+   */
+  public async discoverDependencyTree(
+    root: string,
+    dependencies: Record<string, string>,
+  ): Promise<Record<string, string>> {
+    const ret: Record<string, string> = {};
+    for (const [name, declaration] of Object.entries(dependencies)) {
+      // eslint-disable-next-line no-await-in-loop
+      const resolved = await this.resolveDependency(root, name, declaration);
+
+      const actualVersion = resolved.dependencyInfo.assembly.version;
+      if (!semver.satisfies(actualVersion, declaration)) {
+        throw new Error(
+          `Declared dependency on version ${declaration} of ${name}, but version ${actualVersion} was found`,
+        );
+      }
+
+      ret[name] = resolved.resolvedFile;
+    }
+    return ret;
   }
-  const packageVersions: { [name: string]: string } = {};
-  for (const name of Object.keys(dependencies)) {
-    if (bundled.has(name)) {
-      continue;
+
+  /**
+   * From a set of resolved paths, recursively return all assemblies
+   */
+  public assemblyClosure(resolved: Record<string, string>): spec.Assembly[] {
+    const closure = new Map<string, spec.Assembly>();
+    const queue = Array.from(Object.values(resolved));
+    while (queue.length > 0) {
+      const next = queue.shift()!;
+      const resolved = this.cache.get(next);
+      if (!resolved) {
+        throw new Error(`Path ${next} not seen before`);
+      }
+      if (closure.has(next)) {
+        continue;
+      }
+
+      closure.set(next, resolved.assembly);
+      queue.push(...Object.values(resolved.resolvedDependencies));
     }
+    return Array.from(closure.values());
+  }
+
+  private async resolveDependency(
+    root: string,
+    name: string,
+    declaration: string,
+  ) {
     const { version: versionString, localPackage } = _resolveVersion(
-      dependencies[name],
-      searchPath,
+      declaration,
+      root,
     );
     const version = new semver.Range(versionString);
     if (!version) {
       throw new Error(
         `Invalid semver expression for ${name}: ${versionString}`,
       );
     }
-    // eslint-disable-next-line no-await-in-loop
-    const pkg = await _tryResolveAssembly(name, localPackage, searchPath);
-    LOG.debug(`Resolved dependency ${name} to ${pkg}`);
-    // eslint-disable-next-line no-await-in-loop
-    const assm = await loadAndValidateAssembly(pkg, assemblyCache);
-    if (!semver.satisfies(assm.version, version)) {
-      throw new Error(
-        `Declared dependency on version ${versionString} of ${name}, but version ${assm.version} was found`,
-      );
-    }
-    packageVersions[assm.name] =
-      packageVersions[assm.name] != null
-        ? intersect(versionString, packageVersions[assm.name])
-        : versionString;
-    transitiveAssemblies[assm.name] = assm;
-    const pkgDir = path.dirname(pkg);
-    if (assm.dependencies) {
-      // eslint-disable-next-line no-await-in-loop
-      await _loadDependencies(
-        assm.dependencies,
-        pkgDir,
-        transitiveAssemblies,
-        assemblyCache,
-      );
+    const jsiiFile = await _tryResolveAssembly(name, localPackage, root);
+    LOG.debug(`Resolved dependency ${name} to ${jsiiFile}`);
+    return {
+      resolvedVersion: versionString,
+      resolvedFile: jsiiFile,
+      dependencyInfo: await this.loadAssemblyAndRecurse(jsiiFile),
+    };
+  }
+
+  private async loadAssemblyAndRecurse(jsiiFile: string) {
+    // Only recurse if we haven't seen this assembly yet
+    if (this.cache.has(jsiiFile)) {
+      return this.cache.get(jsiiFile)!;
     }
+
+    // eslint-disable-next-line no-await-in-loop
+    const assembly = await this.loadAssembly(jsiiFile);
+    // Continue loading any dependencies declared in the asm
+
+    const resolvedDependencies = assembly.dependencies
+      ? await this.discoverDependencyTree(
+          path.dirname(jsiiFile),
+          assembly.dependencies,
+        )
+      : {};
+
+    const depInfo: DependencyInfo = {
+      assembly,
+      resolvedDependencies,
+    };
+    this.cache.set(jsiiFile, depInfo);
+    return depInfo;
   }
-  return packageVersions;
-}
 
-/**
- * Load a JSII filename and validate it; cached to avoid redundant loads of the same JSII assembly
- */
-async function loadAndValidateAssembly(
-  jsiiFileName: string,
-  cache: Map<string, spec.Assembly>,
-): Promise<spec.Assembly> {
-  if (!cache.has(jsiiFileName)) {
+  /**
+   * Load a JSII filename and validate it; cached to avoid redundant loads of the same JSII assembly
+   */
+  private async loadAssembly(jsiiFileName: string): Promise<spec.Assembly> {
     try {
-      cache.set(jsiiFileName, await fs.readJson(jsiiFileName));
+      return await fs.readJson(jsiiFileName);
     } catch (e) {
       throw new Error(`Error loading ${jsiiFileName}: ${e}`);
     }
   }
-  return cache.get(jsiiFileName)!;
 }
 
 function _required<T>(value: T, message: string): T {
@@ -518,3 +561,16 @@ function _loadDiagnostics(entries?: { [key: string]: string }):
   }
   return result;
 }
+
+function filterDictByKey<A>(
+  xs: Record<string, A>,
+  predicate: (key: string) => boolean,
+): Record<string, A> {
+  const ret: Record<string, A> = {};
+  for (const [key, value] of Object.entries(xs)) {
+    if (predicate(key)) {
+      ret[key] = value;
+    }
+  }
+  return ret;
+}