feat(jsii-reflect): TypeSystem can be locked to improve reflection pe…

…rformance (#4318)

Memoizes additional calls that rely on the typesystem to retrieve type instances.
These calls can (theoretically) change when the typesystem is changed.
Therefore we cannot assume it's okay to always memoize the first call.

To workaround this limitation, we introduce a new mechanism to manually `lock` the typesystem once all assemblies are loaded.
We then can start memoizing the additional calls.

For example in `awslint` this reduces the runtime against `aws-cdk-lib` by ~20s.

---

By submitting this pull request, I confirm that my contribution is made under the terms of the [Apache 2.0 license].

[Apache 2.0 license]: https://www.apache.org/licenses/LICENSE-2.0

packages/jsii-reflect/lib/_memoized.ts

@@ -1,14 +1,41 @@
-/* eslint-disable @typescript-eslint/ban-types -- WeakMap<T, _> demands T extends object */
+import { TypeSystem } from './type-system';
 
+/* eslint-disable @typescript-eslint/ban-types -- WeakMap<T, _> demands T extends object */
 const CACHE = new WeakMap<object, Map<string, unknown>>();
 
+function memoizedGet(original: () => any, propertyKey: string): () => any {
+  return function (this: object): unknown {
+    let cache = CACHE.get(this);
+    if (cache == null) {
+      cache = new Map();
+      CACHE.set(this, cache);
+    }
+    if (cache.has(propertyKey)) {
+      const result = cache.get(propertyKey);
+      if (Array.isArray(result)) {
+        // Return a copy of arrays as a precaution
+        return Array.from(result);
+      }
+      return result;
+    }
+    const result = original.call(this);
+    // If the result is an array, memoize a copy for safety.
+    cache.set(propertyKey, Array.isArray(result) ? Array.from(result) : result);
+    return result;
+  };
+}
+
 /**
  * Decorates property readers for readonly properties so that their results are
  * memoized in a `WeakMap`-based cache. Those properties will consequently be
  * computed exactly once.
  *
  * This can only be applied to property accessors (`public get foo(): any`), and not to
  * property declarations (`public readonly foo: any`).
+ *
+ * This should not be applied to any computations relying on a typesystem.
+ * The typesystem can be changed and thus change the result of the call.
+ * Use `memoizedWhenLocked` instead.
  */
 export function memoized(
   _prototype: unknown,
@@ -23,23 +50,27 @@ export function memoized(
   }
 
   const original = descriptor.get;
-  descriptor.get = function memoizedGet(this: object): unknown {
-    let cache = CACHE.get(this);
-    if (cache == null) {
-      cache = new Map();
-      CACHE.set(this, cache);
-    }
-    if (cache.has(propertyKey)) {
-      const result = cache.get(propertyKey);
-      if (Array.isArray(result)) {
-        // Return a copy of arrays as a precaution
-        return Array.from(result);
-      }
-      return result;
+  descriptor.get = memoizedGet(original, propertyKey);
+}
+
+export function memoizedWhenLocked<T extends { system: TypeSystem }>(
+  _prototype: T,
+  propertyKey: string,
+  descriptor: PropertyDescriptor,
+): void {
+  if (!descriptor.get) {
+    throw new Error(`@memoized can only be applied to property getters!`);
+  }
+  if (descriptor.set) {
+    throw new Error(`@memoized can only be applied to readonly properties!`);
+  }
+
+  const original = descriptor.get;
+  descriptor.get = function (this: T): unknown {
+    if (this.system.isLocked) {
+      return memoizedGet(original, propertyKey).call(this);
     }
-    const result = original.call(this);
-    // If the result is an array, memoize a copy for safety.
-    cache.set(propertyKey, Array.isArray(result) ? Array.from(result) : result);
-    return result;
+
+    return original.call(this);
   };
 }

packages/jsii-reflect/lib/class.ts

@@ -1,5 +1,6 @@
 import * as jsii from '@jsii/spec';
 
+import { memoizedWhenLocked } from './_memoized';
 import { Assembly } from './assembly';
 import { Initializer } from './initializer';
 import { InterfaceType } from './interface';
@@ -20,6 +21,7 @@ export class ClassType extends ReferenceType {
   /**
    * Base class (optional).
    */
+  @memoizedWhenLocked
   public get base(): ClassType | undefined {
     if (!this.spec.base) {
       return undefined;
@@ -60,12 +62,22 @@ export class ClassType extends ReferenceType {
 
   /**
    * Returns list of all base classes (first is the direct base and last is the top-most).
+   *
+   * @deprecated use ClassType.ancestors instead
    */
   public getAncestors() {
+    return this.ancestors;
+  }
+
+  /**
+   * Returns list of all base classes (first is the direct base and last is the top-most).
+   */
+  @memoizedWhenLocked
+  public get ancestors() {
     const out = new Array<ClassType>();
     if (this.base) {
       out.push(this.base);
-      out.push(...this.base.getAncestors());
+      out.push(...this.base.ancestors);
     }
     return out;
   }

packages/jsii-reflect/lib/method.ts

@@ -1,5 +1,6 @@
 import * as jsii from '@jsii/spec';
 
+import { memoizedWhenLocked } from './_memoized';
 import { Assembly } from './assembly';
 import { Callable } from './callable';
 import { Documentable } from './docs';
@@ -42,6 +43,7 @@ export class Method
     return this.spec.name;
   }
 
+  @memoizedWhenLocked
   public get overrides(): Type | undefined {
     if (!this.spec.overrides) {
       return undefined;

packages/jsii-reflect/lib/type-system.ts

@@ -22,13 +22,28 @@ export class TypeSystem {
 
   private readonly _cachedClasses = new Map<Assembly, readonly ClassType[]>();
 
+  private _locked = false;
+  public get isLocked(): boolean {
+    return this._locked;
+  }
+
   /**
    * All assemblies in this type system.
    */
   public get assemblies(): readonly Assembly[] {
     return Array.from(this._assemblyLookup.values());
   }
 
+  /**
+   * Locks the TypeSystem from further changes
+   *
+   * Call this once all assemblies have been loaded.
+   * This allows the reflection to optimize and cache certain expensive calls.
+   */
+  public lock() {
+    this._locked = true;
+  }
+
   /**
    * Load all JSII dependencies of the given NPM package directory.
    *
@@ -164,6 +179,10 @@ export class TypeSystem {
   }
 
   public addAssembly(asm: Assembly, options: { isRoot?: boolean } = {}) {
+    if (this.isLocked) {
+      throw new Error('The typesystem has been locked from further changes');
+    }
+
     if (asm.system !== this) {
       throw new Error('Assembly has been created for different typesystem');
     }

packages/jsii-reflect/lib/type.ts

@@ -124,7 +124,7 @@ export abstract class Type implements Documentable, SourceLocatable {
       return this.getInterfaces(true).some((iface) => iface === base);
     }
     if (this.isClassType() && base.isClassType()) {
-      return this.getAncestors().some((clazz) => clazz === base);
+      return this.ancestors.some((clazz) => clazz === base);
     }
     return false;
   }

packages/jsii-reflect/test/_memoized.test.ts

@@ -0,0 +1,79 @@
+import { TypeSystem } from '../lib';
+import { memoized, memoizedWhenLocked } from '../lib/_memoized';
+
+const accessorSpy = jest.fn(() => 'foobar');
+
+class TestClass {
+  public constructor(public readonly system: TypeSystem) {}
+
+  public get uncached(): string {
+    return accessorSpy();
+  }
+
+  @memoized
+  public get cached(): string {
+    return accessorSpy();
+  }
+
+  @memoizedWhenLocked
+  public get cachedWhenLocked(): string {
+    return accessorSpy();
+  }
+}
+
+// eslint-disable-next-line @typescript-eslint/no-empty-function
+function noop(_val: unknown) {}
+
+describe('memoized', () => {
+  beforeEach(() => {
+    accessorSpy.mockClear();
+  });
+  const subject = new TestClass(new TypeSystem());
+
+  test('cached property is memoized', () => {
+    // Access the property twice
+    noop(subject.cached);
+    noop(subject.cached);
+
+    expect(accessorSpy).toHaveBeenCalledTimes(1);
+    expect(subject.cached).toBe('foobar');
+  });
+
+  test('uncached property is not memoized', () => {
+    // Access the property twice
+    noop(subject.uncached);
+    noop(subject.uncached);
+
+    expect(accessorSpy).toHaveBeenCalledTimes(2);
+    expect(subject.uncached).toBe('foobar');
+  });
+});
+
+describe('memoizedWhenLocked', () => {
+  let subject: TestClass;
+  beforeEach(() => {
+    accessorSpy.mockClear();
+    subject = new TestClass(new TypeSystem());
+  });
+
+  test('property is memoized when the typesystem is locked', () => {
+    // Lock the typesystem to enable memoizing
+    subject.system.lock();
+
+    // Access the property twice
+    noop(subject.cachedWhenLocked);
+    noop(subject.cachedWhenLocked);
+
+    expect(accessorSpy).toHaveBeenCalledTimes(1);
+    expect(subject.cachedWhenLocked).toBe('foobar');
+  });
+
+  test('property is not memoized when the typesystem is not locked', () => {
+    // Access the property twice
+    noop(subject.cachedWhenLocked);
+    noop(subject.cachedWhenLocked);
+
+    expect(accessorSpy).toHaveBeenCalledTimes(2);
+    expect(subject.cachedWhenLocked).toBe('foobar');
+  });
+});