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

…rformance

Memoizes additional calls when the typesystem has been locked.

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');
+  });
+});