refactor(ngcc): use a special map for memoizing expensive-to-compute values (#34512)

Previously, in cases were values were expensive to compute and would be
used multiple times, a combination of a regular `Map` and a helper
function (`getOrDefault()`) was used to ensure values were only computed
once.

This commit uses a special `Map`-like structure to compute and memoize
such expensive values without the need to a helper function.

PR Close #34512

packages/compiler-cli/ngcc/src/host/commonjs_host.ts

@@ -11,15 +11,20 @@ import {absoluteFrom} from '../../../src/ngtsc/file_system';
 import {Declaration, Import} from '../../../src/ngtsc/reflection';
 import {Logger} from '../logging/logger';
 import {BundleProgram} from '../packages/bundle_program';
-import {getOrDefault, isDefined, stripExtension} from '../utils';
+import {FactoryMap, isDefined, stripExtension} from '../utils';
 
 import {ExportDeclaration, ExportStatement, ReexportStatement, RequireCall, findNamespaceOfIdentifier, findRequireCallReference, isExportStatement, isReexportStatement, isRequireCall} from './commonjs_umd_utils';
 import {Esm5ReflectionHost} from './esm5_host';
 import {NgccClassSymbol} from './ngcc_host';
 
 export class CommonJsReflectionHost extends Esm5ReflectionHost {
-  protected commonJsExports = new Map<ts.SourceFile, Map<string, Declaration>|null>();
-  protected topLevelHelperCalls = new Map<string, Map<ts.SourceFile, ts.CallExpression[]>>();
+  protected commonJsExports = new FactoryMap<ts.SourceFile, Map<string, Declaration>|null>(
+      sf => this.computeExportsOfCommonJsModule(sf));
+  protected topLevelHelperCalls =
+      new FactoryMap<string, FactoryMap<ts.SourceFile, ts.CallExpression[]>>(
+          helperName => new FactoryMap<ts.SourceFile, ts.CallExpression[]>(
+              sf => sf.statements.map(stmt => this.getHelperCall(stmt, [helperName]))
+                        .filter(isDefined)));
   protected program: ts.Program;
   protected compilerHost: ts.CompilerHost;
   constructor(logger: Logger, isCore: boolean, src: BundleProgram, dts: BundleProgram|null = null) {
@@ -46,12 +51,7 @@ export class CommonJsReflectionHost extends Esm5ReflectionHost {
   }
 
   getExportsOfModule(module: ts.Node): Map<string, Declaration>|null {
-    return super.getExportsOfModule(module) || this.getCommonJsExports(module.getSourceFile());
-  }
-
-  getCommonJsExports(sourceFile: ts.SourceFile): Map<string, Declaration>|null {
-    return getOrDefault(
-        this.commonJsExports, sourceFile, () => this.computeExportsOfCommonJsModule(sourceFile));
+    return super.getExportsOfModule(module) || this.commonJsExports.get(module.getSourceFile());
   }
 
   /**
@@ -63,7 +63,8 @@ export class CommonJsReflectionHost extends Esm5ReflectionHost {
    * no helpers are found.
    *
    * @param classSymbol the class whose helper calls we are interested in.
-   * @param helperName the name of the helper (e.g. `__decorate`) whose calls we are interested in.
+   * @param helperNames the names of the helpers (e.g. `__decorate`) whose calls we are interested
+   * in.
    * @returns an array of nodes of calls to the helper with the given name.
    */
   protected getHelperCallsForClass(classSymbol: NgccClassSymbol, helperNames: string[]):
@@ -92,11 +93,8 @@ export class CommonJsReflectionHost extends Esm5ReflectionHost {
       ts.CallExpression[] {
     const calls: ts.CallExpression[] = [];
     helperNames.forEach(helperName => {
-      const helperCallsMap = getOrDefault(this.topLevelHelperCalls, helperName, () => new Map());
-      calls.push(...getOrDefault(
-          helperCallsMap, sourceFile,
-          () => sourceFile.statements.map(statement => this.getHelperCall(statement, [helperName]))
-                    .filter(isDefined)));
+      const helperCallsMap = this.topLevelHelperCalls.get(helperName);
+      calls.push(...helperCallsMap.get(sourceFile));
     });
     return calls;
   }

packages/compiler-cli/ngcc/src/host/umd_host.ts

@@ -12,16 +12,20 @@ import {absoluteFrom} from '../../../src/ngtsc/file_system';
 import {Declaration, Import} from '../../../src/ngtsc/reflection';
 import {Logger} from '../logging/logger';
 import {BundleProgram} from '../packages/bundle_program';
-import {getOrDefault, stripExtension} from '../utils';
+import {FactoryMap, stripExtension} from '../utils';
 import {ExportDeclaration, ExportStatement, ReexportStatement, findNamespaceOfIdentifier, findRequireCallReference, isExportStatement, isReexportStatement, isRequireCall} from './commonjs_umd_utils';
 import {Esm5ReflectionHost, stripParentheses} from './esm5_host';
 
 export class UmdReflectionHost extends Esm5ReflectionHost {
-  protected umdModules = new Map<ts.SourceFile, UmdModule|null>();
-  protected umdExports = new Map<ts.SourceFile, Map<string, Declaration>|null>();
-  protected umdImportPaths = new Map<ts.ParameterDeclaration, string|null>();
+  protected umdModules =
+      new FactoryMap<ts.SourceFile, UmdModule|null>(sf => this.computeUmdModule(sf));
+  protected umdExports = new FactoryMap<ts.SourceFile, Map<string, Declaration>|null>(
+      sf => this.computeExportsOfUmdModule(sf));
+  protected umdImportPaths =
+      new FactoryMap<ts.ParameterDeclaration, string|null>(param => this.computeImportPath(param));
   protected program: ts.Program;
   protected compilerHost: ts.CompilerHost;
+
   constructor(logger: Logger, isCore: boolean, src: BundleProgram, dts: BundleProgram|null = null) {
     super(logger, isCore, src, dts);
     this.program = src.program;
@@ -48,7 +52,7 @@ export class UmdReflectionHost extends Esm5ReflectionHost {
   }
 
   getExportsOfModule(module: ts.Node): Map<string, Declaration>|null {
-    return super.getExportsOfModule(module) || this.getUmdExports(module.getSourceFile());
+    return super.getExportsOfModule(module) || this.umdExports.get(module.getSourceFile());
   }
 
   getUmdModule(sourceFile: ts.SourceFile): UmdModule|null {
@@ -56,44 +60,11 @@ export class UmdReflectionHost extends Esm5ReflectionHost {
       return null;
     }
 
-    return getOrDefault(this.umdModules, sourceFile, (sf: ts.SourceFile) => {
-      if (sf.statements.length !== 1) {
-        throw new Error(
-            `Expected UMD module file (${sf.fileName}) to contain exactly one statement, but found ${sf.statements.length}.`);
-      }
-      return parseStatementForUmdModule(sf.statements[0]);
-    });
+    return this.umdModules.get(sourceFile);
   }
 
   getUmdImportPath(importParameter: ts.ParameterDeclaration): string|null {
-    return getOrDefault(this.umdImportPaths, importParameter, (param: ts.ParameterDeclaration) => {
-      const umdModule = this.getUmdModule(param.getSourceFile());
-      if (umdModule === null) {
-        return null;
-      }
-
-      const imports = getImportsOfUmdModule(umdModule);
-      if (imports === null) {
-        return null;
-      }
-
-      let importPath: string|null = null;
-
-      for (const i of imports) {
-        // Add all imports to the map to speed up future look ups.
-        this.umdImportPaths.set(i.parameter, i.path);
-        if (i.parameter === importParameter) {
-          importPath = i.path;
-        }
-      }
-
-      return importPath;
-    });
-  }
-
-  getUmdExports(sourceFile: ts.SourceFile): Map<string, Declaration>|null {
-    return getOrDefault(
-        this.umdExports, sourceFile, (sf: ts.SourceFile) => this.computeExportsOfUmdModule(sf));
+    return this.umdImportPaths.get(importParameter);
   }
 
   /** Get the top level statements for a module.
@@ -108,6 +79,16 @@ export class UmdReflectionHost extends Esm5ReflectionHost {
     return umdModule !== null ? Array.from(umdModule.factoryFn.body.statements) : [];
   }
 
+  private computeUmdModule(sourceFile: ts.SourceFile): UmdModule|null {
+    if (sourceFile.statements.length !== 1) {
+      throw new Error(
+          `Expected UMD module file (${sourceFile.fileName}) to contain exactly one statement, ` +
+          `but found ${sourceFile.statements.length}.`);
+    }
+
+    return parseStatementForUmdModule(sourceFile.statements[0]);
+  }
+
   private computeExportsOfUmdModule(sourceFile: ts.SourceFile): Map<string, Declaration>|null {
     const moduleMap = new Map<string, Declaration>();
     for (const statement of this.getModuleStatements(sourceFile)) {
@@ -124,6 +105,30 @@ export class UmdReflectionHost extends Esm5ReflectionHost {
     return moduleMap;
   }
 
+  private computeImportPath(param: ts.ParameterDeclaration): string|null {
+    const umdModule = this.getUmdModule(param.getSourceFile());
+    if (umdModule === null) {
+      return null;
+    }
+
+    const imports = getImportsOfUmdModule(umdModule);
+    if (imports === null) {
+      return null;
+    }
+
+    let importPath: string|null = null;
+
+    for (const i of imports) {
+      // Add all imports to the map to speed up future look ups.
+      this.umdImportPaths.set(i.parameter, i.path);
+      if (i.parameter === param) {
+        importPath = i.path;
+      }
+    }
+
+    return importPath;
+  }
+
   private extractUmdExportDeclaration(statement: ExportStatement): ExportDeclaration {
     const exportExpression = statement.expression.right;
     const declaration = this.getDeclarationOfExpression(exportExpression);

packages/compiler-cli/ngcc/src/utils.ts

@@ -66,13 +66,6 @@ export function findAll<T>(node: ts.Node, test: (node: ts.Node) => node is ts.No
   }
 }
 
-export function getOrDefault<K, V>(map: Map<K, V>, key: K, factory: (key: K) => V): V {
-  if (!map.has(key)) {
-    map.set(key, factory(key));
-  }
-  return map.get(key) !;
-}
-
 /**
  * Does the given declaration have a name which is an identifier?
  * @param declaration The declaration to test.
@@ -98,6 +91,33 @@ export function isRelativePath(path: string): boolean {
   return /^\/|^\.\.?($|\/)/.test(path);
 }
 
+/**
+ * A `Map`-like object that can compute and memoize a missing value for any key.
+ *
+ * The computed values are memoized, so the factory function is not called more than once per key.
+ * This is useful for storing values that are expensive to compute and may be used multiple times.
+ */
+// NOTE:
+// Ideally, this class should extend `Map`, but that causes errors in ES5 transpiled code:
+// `TypeError: Constructor Map requires 'new'`
+export class FactoryMap<K, V> {
+  private internalMap: Map<K, V>;
+
+  constructor(private factory: (key: K) => V, entries?: readonly(readonly[K, V])[]|null) {
+    this.internalMap = new Map(entries);
+  }
+
+  get(key: K): V {
+    if (!this.internalMap.has(key)) {
+      this.internalMap.set(key, this.factory(key));
+    }
+
+    return this.internalMap.get(key) !;
+  }
+
+  set(key: K, value: V): void { this.internalMap.set(key, value); }
+}
+
 /**
  * Attempt to resolve a `path` to a file by appending the provided `postFixes`
  * to the `path` and checking if the file exists on disk.

packages/compiler-cli/ngcc/test/utils_spec.ts

@@ -6,29 +6,47 @@
  * found in the LICENSE file at https://angular.io/license
  */
 
-import {getOrDefault, isRelativePath, stripExtension} from '../src/utils';
+import {FactoryMap, isRelativePath, stripExtension} from '../src/utils';
 
-describe('getOrDefault()', () => {
+describe('FactoryMap', () => {
   it('should return an existing value', () => {
-    const map = new Map([['k1', 'v1'], ['k2', 'v2']]);
-    const factorySpy = jasmine.createSpy('factory');
+    const factoryFnSpy = jasmine.createSpy('factory');
+    const factoryMap = new FactoryMap<string, string>(factoryFnSpy, [['k1', 'v1'], ['k2', 'v2']]);
 
-    expect(getOrDefault(map, 'k1', factorySpy)).toBe('v1');
-    expect(getOrDefault(map, 'k2', factorySpy)).toBe('v2');
-    expect(factorySpy).not.toHaveBeenCalled();
+    expect(factoryMap.get('k1')).toBe('v1');
+    expect(factoryMap.get('k2')).toBe('v2');
+    expect(factoryFnSpy).not.toHaveBeenCalled();
+  });
+
+  it('should not treat falsy values as missing', () => {
+    const factoryFnSpy = jasmine.createSpy('factory').and.returnValue('never gonna happen');
+    const factoryMap = new FactoryMap<string, any>(factoryFnSpy, [
+      ['k1', ''],
+      ['k2', 0],
+      ['k3', false],
+      ['k4', null],
+      ['k5', undefined],
+    ]);
+
+    expect(factoryMap.get('k1')).toBe('');
+    expect(factoryMap.get('k2')).toBe(0);
+    expect(factoryMap.get('k3')).toBe(false);
+    expect(factoryMap.get('k4')).toBe(null);
+    expect(factoryMap.get('k5')).toBe(undefined);
+    expect(factoryFnSpy).not.toHaveBeenCalled();
   });
 
   it('should create, store and return the value if it does not exist', () => {
-    const map = new Map([['k1', 'v1'], ['k2', 'v2']]);
-    const factorySpy = jasmine.createSpy('factory').and.returnValues('v3', 'never gonna happen');
+    const factoryFnSpy = jasmine.createSpy('factory').and.returnValues('v3', 'never gonna happen');
+    const factoryMap = new FactoryMap(factoryFnSpy, [['k1', 'v1'], ['k2', 'v2']]);
 
-    expect(getOrDefault(map, 'k3', factorySpy)).toBe('v3');
-    expect(factorySpy).toHaveBeenCalledTimes(1);
+    expect(factoryMap.get('k3')).toBe('v3');
+    expect(factoryFnSpy).toHaveBeenCalledTimes(1);
 
-    factorySpy.calls.reset();
+    factoryFnSpy.calls.reset();
 
-    expect(getOrDefault(map, 'k3', factorySpy)).toBe('v3');
-    expect(factorySpy).not.toHaveBeenCalled();
+    expect(factoryMap.get('k3')).toBe('v3');
+    expect(factoryFnSpy).not.toHaveBeenCalled();
   });
 });