feat: add runtime type info (fqn and version) to compiled sources

Adds a static member to each exported class with the fully-qualified name and
current version. This is the first, preparatory step for a richer runtime type
information module for jsii, as well as enhanced metadata for CDKv2.

See https://github.com/aws/aws-cdk-rfcs/blob/master/text/0253-cdk-metadata-v2.md#appendix-1-jsii-run-time-type-information
for RFC and specification I'm following here.

packages/jsii/lib/assembler.ts

@@ -21,6 +21,7 @@ import * as bindings from './node-bindings';
 import { ProjectInfo } from './project-info';
 import { isReservedName } from './reserved-words';
 import { DeprecatedRemover } from './transforms/deprecated-remover';
+import { RuntimeTypeInfoInjector } from './transforms/runtime-info';
 import { TsCommentReplacer } from './transforms/ts-comment-replacer';
 import { combinedTransformers } from './transforms/utils';
 import { Validator } from './validator';
@@ -37,6 +38,7 @@ const LOG = log4js.getLogger('jsii/assembler');
  */
 export class Assembler implements Emitter {
   private readonly commentReplacer = new TsCommentReplacer();
+  private readonly runtimeTypeInfoInjector: RuntimeTypeInfoInjector;
   private readonly deprecatedRemover?: DeprecatedRemover;
 
   private readonly mainFile: string;
@@ -92,11 +94,15 @@ export class Assembler implements Emitter {
     }
 
     this.mainFile = path.resolve(projectInfo.projectRoot, mainFile);
+    this.runtimeTypeInfoInjector = new RuntimeTypeInfoInjector(
+      projectInfo.version,
+    );
   }
 
   public get customTransformers(): ts.CustomTransformers {
     return combinedTransformers(
       this.deprecatedRemover?.customTransformers ?? {},
+      this.runtimeTypeInfoInjector.makeTransformers(),
       this.commentReplacer.makeTransformers(),
     );
   }
@@ -891,6 +897,9 @@ export class Assembler implements Emitter {
         this._typeChecker.getTypeAtLocation(node),
         context,
       );
+      if (jsiiType) {
+        this.registerExportedClassFqn(node, jsiiType.fqn);
+      }
     } else if (ts.isInterfaceDeclaration(node) && _isExported(node)) {
       // export interface Name { ... }
       this._validateHeritageClauses(node.heritageClauses);
@@ -2574,6 +2583,14 @@ export class Assembler implements Emitter {
     }
   }
 
+  /**
+   * Updates the runtime type info with the fully-qualified name for the current class definition.
+   * Used by the runtime type info injector to add this information to the compiled file.
+   */
+  private registerExportedClassFqn(clazz: ts.ClassDeclaration, fqn: string) {
+    this.runtimeTypeInfoInjector.registerClassFqn(clazz, fqn);
+  }
+
   /**
    * From the given JSIIDocs, re-render the TSDoc comment for the Node
    *

packages/jsii/lib/transforms/runtime-info.ts

@@ -0,0 +1,144 @@
+import * as ts from 'typescript';
+
+/**
+ * Provides a TransformerFactory to annotate classes with runtime information
+ * (e.g., fully-qualified name, version).
+ *
+ * It does this by first inserting this definition at the top of each source file:
+ * ```
+ * var JSII_RTTI_SYMBOL_1 = Symbol.for("jsii.rtti");
+ * ```
+ *
+ * Then, for each class that has registered runtime information during assembly,
+ * insert a static member to the class with its fqn and version:
+ * ```
+ * private static readonly [JSII_RTTI_SYMBOL_1] = { fqn: "ModuleName.ClassName", version: "1.2.3" }
+ * ```
+ */
+export class RuntimeTypeInfoInjector {
+  private readonly fqnsByClass = new Map<ts.ClassDeclaration, string>();
+
+  public constructor(private readonly version: string) {}
+
+  /**
+   * Register the fully-qualified name (fqn) of a class with its ClassDeclaration.
+   * Only ClassDeclarations with registered fqns will be annotated.
+   */
+  public registerClassFqn(clazz: ts.ClassDeclaration, fqn: string) {
+    this.fqnsByClass.set(clazz, fqn);
+  }
+
+  /**
+   * Return the set of Transformers to be used in TSC's program.emit()
+   */
+  public makeTransformers(): ts.CustomTransformers {
+    return {
+      before: [this.runtimeTypeTransformer()],
+    };
+  }
+
+  public runtimeTypeTransformer(): ts.TransformerFactory<ts.SourceFile> {
+    return (context) => {
+      return (sourceFile) => {
+        const rttiSymbolIdentifier = ts.createUniqueName('JSII_RTTI_SYMBOL');
+
+        let classesAnnotated = false;
+        const visitor = (node: ts.Node): ts.Node => {
+          if (ts.isClassDeclaration(node)) {
+            const fqn = this.getClassFqn(node);
+            if (fqn) {
+              classesAnnotated = true;
+              return this.addRuntimeInfoToClass(
+                node,
+                fqn,
+                rttiSymbolIdentifier,
+              );
+            }
+          }
+          return ts.visitEachChild(node, visitor, context);
+        };
+
+        // Visit the source file, annotating the classes.
+        let annotatedSourceFile = ts.visitNode(sourceFile, visitor);
+
+        // Only add the symbol definition if it's actually used.
+        if (classesAnnotated) {
+          const rttiSymbol = ts.createCall(
+            ts.createPropertyAccess(
+              ts.createIdentifier('Symbol'),
+              ts.createIdentifier('for'),
+            ),
+            undefined,
+            [ts.createStringLiteral('jsii.rtti')],
+          );
+          const rttiSymbolDeclaration = ts.createVariableDeclaration(
+            rttiSymbolIdentifier,
+            undefined,
+            rttiSymbol,
+          );
+          const variableDeclaration = ts.createVariableStatement(
+            [],
+            ts.createVariableDeclarationList(
+              [rttiSymbolDeclaration],
+              ts.NodeFlags.Const,
+            ),
+          );
+
+          annotatedSourceFile = ts.updateSourceFileNode(annotatedSourceFile, [
+            variableDeclaration,
+            ...annotatedSourceFile.statements,
+          ]);
+        }
+
+        return annotatedSourceFile;
+      };
+    };
+  }
+
+  /** Used instead of direct access to the map to faciliate testing. */
+  protected getClassFqn(clazz: ts.ClassDeclaration): string | undefined {
+    return this.fqnsByClass.get(clazz);
+  }
+
+  /**
+   * If the ClassDeclaration has an associated fully-qualified name registered,
+   * will append a static property to the class with the fqn and version.
+   */
+  private addRuntimeInfoToClass(
+    node: ts.ClassDeclaration,
+    fqn: string,
+    rttiSymbol: ts.Identifier,
+  ): ts.ClassDeclaration {
+    const runtimeInfo = ts.createObjectLiteral([
+      ts.createPropertyAssignment(
+        ts.createIdentifier('fqn'),
+        ts.createStringLiteral(fqn),
+      ),
+      ts.createPropertyAssignment(
+        ts.createIdentifier('version'),
+        ts.createStringLiteral(this.version),
+      ),
+    ]);
+    const runtimeProperty = ts.createProperty(
+      undefined,
+      ts.createModifiersFromModifierFlags(
+        ts.ModifierFlags.Private |
+          ts.ModifierFlags.Static |
+          ts.ModifierFlags.Readonly,
+      ),
+      ts.createComputedPropertyName(rttiSymbol),
+      undefined,
+      undefined,
+      runtimeInfo,
+    );
+    return ts.updateClassDeclaration(
+      node,
+      node.decorators,
+      node.modifiers,
+      node.name,
+      node.typeParameters,
+      node.heritageClauses,
+      [runtimeProperty, ...node.members],
+    );
+  }
+}

packages/jsii/test/transforms/runtime-info.test.ts

@@ -0,0 +1,154 @@
+import * as ts from 'typescript';
+
+import { RuntimeTypeInfoInjector } from '../../lib/transforms/runtime-info';
+
+test('leaves files without classes unaltered', () => {
+  expect(transformedSource(EXAMPLE_NO_CLASS, 'Foo')).not.toContain(
+    'JSII_RTTI_SYMBOL',
+  );
+});
+
+test('leaves files without classes with metadata unaltered', () => {
+  expect(transformedSource(EXAMPLE_SINGLE_CLASS)).not.toContain(
+    'JSII_RTTI_SYMBOL',
+  );
+});
+
+test('adds jsii.rtti symbol at the top of each file when classes are present', () => {
+  expect(transformedSource(EXAMPLE_SINGLE_CLASS, 'Foo')).toContain(
+    'const JSII_RTTI_SYMBOL_1 = Symbol.for("jsii.rtti");',
+  );
+});
+
+test('adds runtime info for a class', () => {
+  expect(transformedSource(EXAMPLE_SINGLE_CLASS, 'Foo')).toContain(
+    'private static readonly [JSII_RTTI_SYMBOL_1] = { fqn: "RuntimeInfoTest.Foo", version: "1.2.3" }',
+  );
+});
+
+test('adds runtime info for each class', () => {
+  const transformed = transformedSource(EXAMPLE_MULTIPLE_CLASSES, 'Foo', 'Bar');
+  expect(transformed).toContain(
+    'private static readonly [JSII_RTTI_SYMBOL_1] = { fqn: "RuntimeInfoTest.Foo", version: "1.2.3" }',
+  );
+  expect(transformed).toContain(
+    'private static readonly [JSII_RTTI_SYMBOL_1] = { fqn: "RuntimeInfoTest.Bar", version: "1.2.3" }',
+  );
+});
+
+test('skips runtime info if not available', () => {
+  const transformed = transformedSource(EXAMPLE_MULTIPLE_CLASSES, 'Foo');
+  expect(transformed).toContain(
+    'private static readonly [JSII_RTTI_SYMBOL_1] = { fqn: "RuntimeInfoTest.Foo", version: "1.2.3" }',
+  );
+  expect(transformed).not.toContain(
+    'private static readonly [JSII_RTTI_SYMBOL_1] = { fqn: "RuntimeInfoTest.Bar", version: "1.2.3" }',
+  );
+});
+
+test('creates a unique name if the default is taken', () => {
+  // Conflicting example has existing variable for JSII_RTTI_SYMBOL_1, so transformation should use _2.
+  const transformed = transformedSource(EXAMPLE_CONFLICTING_NAME, 'Foo');
+  expect(transformed).toContain(
+    'const JSII_RTTI_SYMBOL_2 = Symbol.for("jsii.rtti");',
+  );
+  expect(transformed).toContain(
+    'private static readonly [JSII_RTTI_SYMBOL_2] = { fqn: "RuntimeInfoTest.Foo", version: "1.2.3" }',
+  );
+});
+
+function transformedSource(source: string, ...classNames: string[]) {
+  const mockedTypeInfo = mockedTypeInfoForClasses(...classNames);
+  const injector = new TestRuntimeTypeInfoInjector(mockedTypeInfo);
+  const transformed = ts.transform(
+    ts.createSourceFile('source.ts', source, ts.ScriptTarget.Latest),
+    [injector.runtimeTypeTransformer()],
+  );
+  return ts
+    .createPrinter()
+    .printBundle(ts.createBundle(transformed.transformed));
+}
+
+/** Test subclass of RuntimeTypeInfoInjector that accepts overrides for type info */
+class TestRuntimeTypeInfoInjector extends RuntimeTypeInfoInjector {
+  public constructor(private readonly typeInfo: Map<string, string>) {
+    super('1.2.3');
+  }
+
+  protected getClassFqn(clazz: ts.ClassDeclaration): string | undefined {
+    return clazz.name ? this.typeInfo.get(clazz.name.text) : undefined;
+  }
+}
+
+/**
+ * Mock the Map<ts.ClassDefinition, string> of classes to fqns.
+ * This assumes each class name only appears once in the source,
+ * which is a reasonable assumption for these tests.
+ */
+function mockedTypeInfoForClasses(
+  ...classNames: string[]
+): Map<string, string> {
+  const typeInfoMap = new Map<string, string>();
+  classNames.forEach((clazz) =>
+    typeInfoMap.set(clazz, `RuntimeInfoTest.${clazz}`),
+  );
+  return typeInfoMap;
+}
+
+/**
+ * ===============================
+ * =    EXAMPLE SOURCE FILES     =
+ * ===============================
+ */
+
+const EXAMPLE_NO_CLASS = `
+import * as ts from 'typescript';
+
+interface Foo {
+  readonly foobar: string;
+}
+`;
+
+const EXAMPLE_SINGLE_CLASS = `
+import * as ts from 'typescript';
+
+class Foo {
+  constructor(public readonly bar: string) {}
+}
+`;
+
+const EXAMPLE_MULTIPLE_CLASSES = `
+class Foo {
+  constructor(public readonly bar: string) {}
+  public doStuff() { return 42; }
+}
+
+interface FooBar {
+  readonly answer: number;
+}
+
+/**
+ * A bar.
+ */
+class Bar {
+  public doStuffToo() {
+    return new class implements FooBar {
+      public readonly answer = 21;
+    }();
+  }
+}
+
+export default class {
+  constructor() {}
+}
+`;
+
+const EXAMPLE_CONFLICTING_NAME = `
+import * as ts from 'typescript';
+
+const JSII_RTTI_SYMBOL_1 = 42;
+
+class Foo {
+  constructor(public readonly bar: string) {}
+}
+`;