feat(ivy): a generic visitor which allows prefixing nodes for ngtsc (#24230)

This adds ngtsc/util/src/visitor, a utility for visiting TS ASTs that
can add synthetic nodes immediately prior to certain types of nodes (e.g.
class declarations). It's useful to lift definitions that need to be
referenced repeatedly in generated code outside of the class that defines
them.

PR Close #24230

packages/compiler-cli/src/ngtsc/metadata/test/BUILD.bazel

@@ -12,6 +12,7 @@ ts_library(
     deps = [
         "//packages:types",
         "//packages/compiler-cli/src/ngtsc/metadata",
+        "//packages/compiler-cli/src/ngtsc/testing",
     ],
 )
 

packages/compiler-cli/src/ngtsc/metadata/test/reflector_spec.ts

@@ -8,14 +8,13 @@
 
 import * as ts from 'typescript';
 
+import {getDeclaration, makeProgram} from '../../testing/in_memory_typescript';
 import {Parameter, reflectConstructorParameters} from '../src/reflector';
 
-import {getDeclaration, makeProgram} from './in_memory_typescript';
-
 describe('reflector', () => {
   describe('ctor params', () => {
     it('should reflect a single argument', () => {
-      const program = makeProgram([{
+      const {program} = makeProgram([{
         name: 'entry.ts',
         contents: `
             class Bar {}
@@ -33,7 +32,7 @@ describe('reflector', () => {
     });
 
     it('should reflect a decorated argument', () => {
-      const program = makeProgram([
+      const {program} = makeProgram([
         {
           name: 'dec.ts',
           contents: `
@@ -61,7 +60,7 @@ describe('reflector', () => {
     });
 
     it('should reflect a decorated argument with a call', () => {
-      const program = makeProgram([
+      const {program} = makeProgram([
         {
           name: 'dec.ts',
           contents: `
@@ -89,7 +88,7 @@ describe('reflector', () => {
     });
 
     it('should reflect a decorated argument with an indirection', () => {
-      const program = makeProgram([
+      const {program} = makeProgram([
         {
           name: 'bar.ts',
           contents: `

packages/compiler-cli/src/ngtsc/metadata/test/resolver_spec.ts

@@ -8,17 +8,17 @@
 
 import * as ts from 'typescript';
 
+import {getDeclaration, makeProgram} from '../../testing/in_memory_typescript';
 import {ResolvedValue, staticallyResolve} from '../src/resolver';
 
-import {getDeclaration, makeProgram} from './in_memory_typescript';
-
 function makeSimpleProgram(contents: string): ts.Program {
-  return makeProgram([{name: 'entry.ts', contents}]);
+  return makeProgram([{name: 'entry.ts', contents}]).program;
 }
 
 function makeExpression(
     code: string, expr: string): {expression: ts.Expression, checker: ts.TypeChecker} {
-  const program = makeProgram([{name: 'entry.ts', contents: `${code}; const target$ = ${expr};`}]);
+  const {program} =
+      makeProgram([{name: 'entry.ts', contents: `${code}; const target$ = ${expr};`}]);
   const checker = program.getTypeChecker();
   const decl = getDeclaration(program, 'entry.ts', 'target$', ts.isVariableDeclaration);
   return {
@@ -34,7 +34,7 @@ function evaluate<T extends ResolvedValue>(code: string, expr: string): T {
 
 describe('ngtsc metadata', () => {
   it('reads a file correctly', () => {
-    const program = makeProgram([
+    const {program} = makeProgram([
       {
         name: 'entry.ts',
         contents: `
@@ -117,7 +117,7 @@ describe('ngtsc metadata', () => {
   });
 
   it('reads values from default exports', () => {
-    const program = makeProgram([
+    const {program} = makeProgram([
       {name: 'second.ts', contents: 'export default {property: "test"}'},
       {
         name: 'entry.ts',
@@ -135,7 +135,7 @@ describe('ngtsc metadata', () => {
   });
 
   it('reads values from named exports', () => {
-    const program = makeProgram([
+    const {program} = makeProgram([
       {name: 'second.ts', contents: 'export const a = {property: "test"};'},
       {
         name: 'entry.ts',
@@ -152,7 +152,7 @@ describe('ngtsc metadata', () => {
   });
 
   it('chain of re-exports works', () => {
-    const program = makeProgram([
+    const {program} = makeProgram([
       {name: 'const.ts', contents: 'export const value = {property: "test"};'},
       {name: 'def.ts', contents: `import {value} from './const'; export default value;`},
       {name: 'indirect-reexport.ts', contents: `import value from './def'; export {value};`},

packages/compiler-cli/src/ngtsc/testing/BUILD.bazel

@@ -0,0 +1,15 @@
+package(default_visibility = ["//visibility:public"])
+
+load("//tools:defaults.bzl", "ts_library")
+load("@build_bazel_rules_nodejs//:defs.bzl", "jasmine_node_test")
+
+ts_library(
+    name = "testing",
+    testonly = 1,
+    srcs = glob([
+        "**/*.ts",
+    ]),
+    deps = [
+        "//packages:types",
+    ],
+)

packages/compiler-cli/src/ngtsc/testing/in_memory_typescript.ts

@@ -9,7 +9,8 @@
 import * as path from 'path';
 import * as ts from 'typescript';
 
-export function makeProgram(files: {name: string, contents: string}[]): ts.Program {
+export function makeProgram(files: {name: string, contents: string}[]):
+    {program: ts.Program, host: ts.CompilerHost} {
   const host = new InMemoryHost();
   files.forEach(file => host.writeFile(file.name, file.contents));
 
@@ -17,10 +18,10 @@ export function makeProgram(files: {name: string, contents: string}[]): ts.Progr
   const program = ts.createProgram(rootNames, {noLib: true, experimentalDecorators: true}, host);
   const diags = [...program.getSyntacticDiagnostics(), ...program.getSemanticDiagnostics()];
   if (diags.length > 0) {
-    fail(diags.map(diag => diag.messageText).join(', '));
-    throw new Error(`Typescript diagnostics failed!`);
+    throw new Error(
+        `Typescript diagnostics failed! ${diags.map(diag => diag.messageText).join(', ')}`);
   }
-  return program;
+  return {program, host};
 }
 
 export class InMemoryHost implements ts.CompilerHost {

packages/compiler-cli/src/ngtsc/util/BUILD.bazel

@@ -0,0 +1,12 @@
+package(default_visibility = ["//visibility:public"])
+
+load("//tools:defaults.bzl", "ts_library")
+
+ts_library(
+    name = "util",
+    srcs = glob([
+        "index.ts",
+        "src/**/*.ts",
+    ]),
+    module_name = "@angular/compiler-cli/src/ngtsc/util",
+)

packages/compiler-cli/src/ngtsc/util/src/visitor.ts

@@ -0,0 +1,118 @@
+/**
+ * @license
+ * Copyright Google Inc. All Rights Reserved.
+ *
+ * Use of this source code is governed by an MIT-style license that can be
+ * found in the LICENSE file at https://angular.io/license
+ */
+
+import * as ts from 'typescript';
+
+/**
+ * Result type of visiting a node that's typically an entry in a list, which allows specifying that
+ * nodes should be added before the visited node in the output.
+ */
+export type VisitListEntryResult<B extends ts.Node, T extends B> = {
+  node: T,
+  before?: B[]
+};
+
+/**
+ * Visit a node with the given visitor and return a transformed copy.
+ */
+export function visit<T extends ts.Node>(
+    node: T, visitor: Visitor, context: ts.TransformationContext): T {
+  return visitor._visit(node, context);
+}
+
+/**
+ * Abstract base class for visitors, which processes certain nodes specially to allow insertion
+ * of other nodes before them.
+ */
+export abstract class Visitor {
+  /**
+   * Maps statements to an array of statements that should be inserted before them.
+   */
+  private _before = new Map<ts.Statement, ts.Statement[]>();
+
+  /**
+   * Visit a class declaration, returning at least the transformed declaration and optionally other
+   * nodes to insert before the declaration.
+   */
+  visitClassDeclaration(node: ts.ClassDeclaration):
+      VisitListEntryResult<ts.Statement, ts.ClassDeclaration> {
+    return {node};
+  }
+
+  private _visitClassDeclaration(node: ts.ClassDeclaration, context: ts.TransformationContext):
+      ts.ClassDeclaration {
+    const result = this.visitClassDeclaration(node);
+    const visited = ts.visitEachChild(result.node, child => this._visit(child, context), context);
+    if (result.before !== undefined) {
+      // Record that some nodes should be inserted before the given declaration. The declaration's
+      // parent's _visit call is responsible for performing this insertion.
+      this._before.set(visited, result.before);
+    }
+    return visited;
+  }
+
+  /**
+   * Visit types of nodes which don't have their own explicit visitor.
+   */
+  visitOtherNode<T extends ts.Node>(node: T): T { return node; }
+
+  private _visitOtherNode<T extends ts.Node>(node: T, context: ts.TransformationContext): T {
+    return ts.visitEachChild(
+        this.visitOtherNode(node), child => this._visit(child, context), context);
+  }
+
+  /**
+   * @internal
+   */
+  _visit<T extends ts.Node>(node: T, context: ts.TransformationContext): T {
+    // First, visit the node. visitedNode starts off as `null` but should be set after visiting
+    // is completed.
+    let visitedNode: T|null = null;
+    if (ts.isClassDeclaration(node)) {
+      visitedNode = this._visitClassDeclaration(node, context) as typeof node;
+    } else {
+      visitedNode = this._visitOtherNode(node, context);
+    }
+
+    // If the visited node has a `statements` array then process them, maybe replacing the visited
+    // node and adding additional statements.
+    if (hasStatements(visitedNode)) {
+      visitedNode = this._maybeProcessStatements(visitedNode);
+    }
+
+    return visitedNode;
+  }
+
+  private _maybeProcessStatements<T extends ts.Node&{statements: ts.NodeArray<ts.Statement>}>(
+      node: T): T {
+    // Shortcut - if every statement doesn't require nodes to be prepended, this is a no-op.
+    if (node.statements.every(stmt => !this._before.has(stmt))) {
+      return node;
+    }
+
+    // There are statements to prepend, so clone the original node.
+    const clone = ts.getMutableClone(node);
+
+    // Build a new list of statements and patch it onto the clone.
+    const newStatements: ts.Statement[] = [];
+    clone.statements.forEach(stmt => {
+      if (this._before.has(stmt)) {
+        newStatements.push(...(this._before.get(stmt) !as ts.Statement[]));
+        this._before.delete(stmt);
+      }
+      newStatements.push(stmt);
+    });
+    clone.statements = ts.createNodeArray(newStatements, node.statements.hasTrailingComma);
+    return clone;
+  }
+}
+
+function hasStatements(node: ts.Node): node is ts.Node&{statements: ts.NodeArray<ts.Statement>} {
+  const block = node as{statements?: any};
+  return block.statements !== undefined && Array.isArray(block.statements);
+}

packages/compiler-cli/src/ngtsc/util/test/BUILD.bazel

@@ -0,0 +1,26 @@
+package(default_visibility = ["//visibility:public"])
+
+load("//tools:defaults.bzl", "ts_library")
+load("@build_bazel_rules_nodejs//:defs.bzl", "jasmine_node_test")
+
+ts_library(
+    name = "test_lib",
+    testonly = 1,
+    srcs = glob([
+        "**/*.ts",
+    ]),
+    deps = [
+        "//packages:types",
+        "//packages/compiler-cli/src/ngtsc/testing",
+        "//packages/compiler-cli/src/ngtsc/util",
+    ],
+)
+
+jasmine_node_test(
+    name = "test",
+    bootstrap = ["angular/tools/testing/init_node_no_angular_spec.js"],
+    deps = [
+        ":test_lib",
+        "//tools/testing:node_no_angular",
+    ],
+)

packages/compiler-cli/src/ngtsc/util/test/visitor_spec.ts

@@ -0,0 +1,94 @@
+/**
+ * @license
+ * Copyright Google Inc. All Rights Reserved.
+ *
+ * Use of this source code is governed by an MIT-style license that can be
+ * found in the LICENSE file at https://angular.io/license
+ */
+
+import * as ts from 'typescript';
+
+import {makeProgram} from '../../testing/in_memory_typescript';
+import {VisitListEntryResult, Visitor, visit} from '../src/visitor';
+
+class TestAstVisitor extends Visitor {
+  visitClassDeclaration(node: ts.ClassDeclaration):
+      VisitListEntryResult<ts.Statement, ts.ClassDeclaration> {
+    const name = node.name !.text;
+    const statics =
+        node.members.filter(member => (member.modifiers as ReadonlyArray<ts.Modifier>|| [
+                                      ]).some(mod => mod.kind === ts.SyntaxKind.StaticKeyword));
+    const idStatic = statics
+                         .find(
+                             el => ts.isPropertyDeclaration(el) && ts.isIdentifier(el.name) &&
+                                 el.name.text === 'id') as ts.PropertyDeclaration |
+        undefined;
+    if (idStatic !== undefined) {
+      return {
+        node,
+        before: [
+          ts.createVariableStatement(
+              undefined,
+              [
+                ts.createVariableDeclaration(`${name}_id`, undefined, idStatic.initializer),
+              ]),
+        ],
+      };
+    }
+    return {node};
+  }
+}
+
+function testTransformerFactory(context: ts.TransformationContext): ts.Transformer<ts.SourceFile> {
+  return (file: ts.SourceFile) => visit(file, new TestAstVisitor(), context);
+}
+
+describe('AST Visitor', () => {
+  it('should add a statement before class in plain file', () => {
+    const {program, host} =
+        makeProgram([{name: 'main.ts', contents: `class A { static id = 3; }`}]);
+    const sf = program.getSourceFile('main.ts') !;
+    program.emit(sf, undefined, undefined, undefined, {before: [testTransformerFactory]});
+    const main = host.readFile('/main.js');
+    expect(main).toMatch(/^var A_id = 3;/);
+  });
+
+  it('should add a statement before class inside function definition', () => {
+    const {program, host} = makeProgram([{
+      name: 'main.ts',
+      contents: `
+      export function foo() {
+        var x = 3;
+        class A { static id = 2; }
+        return A;
+      }
+    `
+    }]);
+    const sf = program.getSourceFile('main.ts') !;
+    program.emit(sf, undefined, undefined, undefined, {before: [testTransformerFactory]});
+    const main = host.readFile('/main.js');
+    expect(main).toMatch(/var x = 3;\s+var A_id = 2;\s+var A =/);
+  });
+
+  it('handles nested statements', () => {
+    const {program, host} = makeProgram([{
+      name: 'main.ts',
+      contents: `
+      export class A {
+        static id = 3;
+
+        foo() {
+          class B {
+            static id = 4;
+          }
+          return B;
+        }
+    }`
+    }]);
+    const sf = program.getSourceFile('main.ts') !;
+    program.emit(sf, undefined, undefined, undefined, {before: [testTransformerFactory]});
+    const main = host.readFile('/main.js');
+    expect(main).toMatch(/var A_id = 3;\s+var A = /);
+    expect(main).toMatch(/var B_id = 4;\s+var B = /);
+  });
+});