refactor(compiler): separate compilation and transform phases (#38213)

This commit splits the transformation into 2 separate steps: Ivy compilation and actual transformation
of corresponding TS nodes. This is needed to have all `o.Expression`s generated before any TS transforms
happen. This allows `ConstantPool` to properly identify expressions that can be shared across multiple
components declared in the same file.

Resolves #38203.

PR Close #38213

packages/compiler-cli/src/ngtsc/transform/src/transform.ts

@@ -14,6 +14,7 @@ import {Decorator, ReflectionHost} from '../../reflection';
 import {ImportManager, translateExpression, translateStatement} from '../../translator';
 import {visit, VisitListEntryResult, Visitor} from '../../util/src/visitor';
 
+import {CompileResult} from './api';
 import {TraitCompiler} from './compilation';
 import {addImports} from './utils';
 
@@ -43,12 +44,15 @@ export function ivyTransformFactory(
   };
 }
 
-class IvyVisitor extends Visitor {
+/**
-  constructor(
+ * Visits all classes, performs Ivy compilation where Angular decorators are present and collects
-      private compilation: TraitCompiler, private reflector: ReflectionHost,
+ * result in a Map that associates a ts.ClassDeclaration with Ivy compilation results. This visitor
-      private importManager: ImportManager, private defaultImportRecorder: DefaultImportRecorder,
+ * does NOT perform any TS transformations.
-      private isClosureCompilerEnabled: boolean, private isCore: boolean,
+ */
-      private constantPool: ConstantPool) {
+class IvyCompilationVisitor extends Visitor {
+  public classCompilationMap = new Map<ts.ClassDeclaration, CompileResult[]>();
+
+  constructor(private compilation: TraitCompiler, private constantPool: ConstantPool) {
     super();
   }
 
@@ -56,14 +60,39 @@ class IvyVisitor extends Visitor {
       VisitListEntryResult<ts.Statement, ts.ClassDeclaration> {
     // Determine if this class has an Ivy field that needs to be added, and compile the field
     // to an expression if so.
-    const res = this.compilation.compile(node, this.constantPool);
+    const result = this.compilation.compile(node, this.constantPool);
+    if (result !== null) {
+      this.classCompilationMap.set(node, result);
+    }
+    return {node};
+  }
+}
+
+/**
+ * Visits all classes and performs transformation of corresponding TS nodes based on the Ivy
+ * compilation results (provided as an argument).
+ */
+class IvyTransformationVisitor extends Visitor {
+  constructor(
+      private compilation: TraitCompiler,
+      private classCompilationMap: Map<ts.ClassDeclaration, CompileResult[]>,
+      private reflector: ReflectionHost, private importManager: ImportManager,
+      private defaultImportRecorder: DefaultImportRecorder,
+      private isClosureCompilerEnabled: boolean, private isCore: boolean) {
+    super();
+  }
+
+  visitClassDeclaration(node: ts.ClassDeclaration):
+      VisitListEntryResult<ts.Statement, ts.ClassDeclaration> {
+    // If this class is not registered in the map, it means that it doesn't have Angular decorators,
+    // thus no further processing is required.
+    if (!this.classCompilationMap.has(node)) return {node};
 
-    if (res !== null) {
     // There is at least one field to add.
     const statements: ts.Statement[] = [];
     const members = [...node.members];
 
-      res.forEach(field => {
+    this.classCompilationMap.get(node)!.forEach(field => {
       // Translate the initializer for the field into TS nodes.
       const exprNode = translateExpression(
           field.initializer, this.importManager, this.defaultImportRecorder,
@@ -97,16 +126,13 @@ class IvyVisitor extends Visitor {
     node = ts.updateClassDeclaration(
         node,
         // Remove the decorator which triggered this compilation, leaving the others alone.
-          maybeFilterDecorator(node.decorators, this.compilation.decoratorsFor(node)),
+        maybeFilterDecorator(node.decorators, this.compilation.decoratorsFor(node)), node.modifiers,
-          node.modifiers, node.name, node.typeParameters, node.heritageClauses || [],
+        node.name, node.typeParameters, node.heritageClauses || [],
         // Map over the class members and remove any Angular decorators from them.
         members.map(member => this._stripAngularDecorators(member)));
     return {node, after: statements};
   }
 
-    return {node};
-  }
-
   /**
    * Return all decorators on a `Declaration` which are from @angular/core, or an empty set if none
    * are.
@@ -224,11 +250,26 @@ function transformIvySourceFile(
   const constantPool = new ConstantPool();
   const importManager = new ImportManager(importRewriter);
 
-  // Recursively scan through the AST and perform any updates requested by the IvyCompilation.
+  // The transformation process consists of 2 steps:
-  const visitor = new IvyVisitor(
+  //
-      compilation, reflector, importManager, defaultImportRecorder, isClosureCompilerEnabled,
+  //  1. Visit all classes, perform compilation and collect the results.
-      isCore, constantPool);
+  //  2. Perform actual transformation of required TS nodes using compilation results from the first
-  let sf = visit(file, visitor, context);
+  //     step.
+  //
+  // This is needed to have all `o.Expression`s generated before any TS transforms happen. This
+  // allows `ConstantPool` to properly identify expressions that can be shared across multiple
+  // components declared in the same file.
+
+  // Step 1. Go though all classes in AST, perform compilation and collect the results.
+  const compilationVisitor = new IvyCompilationVisitor(compilation, constantPool);
+  visit(file, compilationVisitor, context);
+
+  // Step 2. Scan through the AST again and perform transformations based on Ivy compilation
+  // results obtained at Step 1.
+  const transformationVisitor = new IvyTransformationVisitor(
+      compilation, compilationVisitor.classCompilationMap, reflector, importManager,
+      defaultImportRecorder, isClosureCompilerEnabled, isCore);
+  let sf = visit(file, transformationVisitor, context);
 
   // Generate the constant statements first, as they may involve adding additional imports
   // to the ImportManager.