feat(compiler-cli): run JIT transform on classes with jit: true opt-out (angular#56892)

Currently when compiling code with the Angular compiler, all classes
with Angular decorators are compiled with AOT. This includes type
checking, scope collection etc.

This may not be desirable for all components, e.g. dynamic components,
or test components w/ `TestBed.configureTestingModule` (if compiled with ngtsc).

Those components can opt out of AOT on a per component-basis via `jit:
true`. This is helpful as it allows incremental migrations/refactorings
to AOT. Whether we want to keep this capability long-term is something
to be discussed separately.

For now though, we should fix that components compiled with `jit: true`
actually work as expected. Currently this **not the case** as soon as
the new initializer APIs are used— as those do no longer declare class
metadata with decorators.

This commit runs the JIT transform on JIT-opted classes.

Related: https://docs.google.com/document/d/1ox4atCJldWWDXlaYgwM-hU8BNsTpKNW7gx8OfZ0HtRY/edit?resourcekey=0-G1haTNYtD-dN0vNRkQ8_OQ&tab=t.0

PR Close angular#56892

Author: devversion

packages/compiler-cli/src/ngtsc/annotations/common/index.ts

@@ -18,3 +18,4 @@ export * from './src/references_registry';
 export * from './src/schema';
 export * from './src/util';
 export * from './src/input_transforms';
+export * from './src/jit_declaration_registry';

packages/compiler-cli/src/ngtsc/annotations/common/src/jit_declaration_registry.ts

@@ -0,0 +1,17 @@
+/**
+ * @license
+ * Copyright Google LLC 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 {ClassDeclaration} from '../../../reflection';
+
+/**
+ * Registry that keeps track of Angular declarations that are explicitly
+ * marked for JIT compilation and are skipping compilation by trait handlers.
+ */
+export class JitDeclarationRegistry {
+  jitDeclarations = new Set<ClassDeclaration>();
+}

packages/compiler-cli/src/ngtsc/annotations/component/src/handler.ts

@@ -177,6 +177,7 @@ import {
   validateAndFlattenComponentImports,
 } from './util';
 import {getTemplateDiagnostics} from '../../../typecheck';
+import {JitDeclarationRegistry} from '../../common/src/jit_declaration_registry';
 
 const EMPTY_ARRAY: any[] = [];
 
@@ -249,6 +250,7 @@ export class ComponentDecoratorHandler
     private readonly enableBlockSyntax: boolean,
     private readonly enableLetSyntax: boolean,
     private readonly localCompilationExtraImportsTracker: LocalCompilationExtraImportsTracker | null,
+    private readonly jitDeclarationRegistry: JitDeclarationRegistry,
   ) {
     this.extractTemplateOptions = {
       enableI18nLegacyMessageIdFormat: this.enableI18nLegacyMessageIdFormat,
@@ -410,10 +412,11 @@ export class ComponentDecoratorHandler
       this.compilationMode,
       this.elementSchemaRegistry.getDefaultComponentElementName(),
     );
-    if (directiveResult === undefined) {
-      // `extractDirectiveMetadata` returns undefined when the @Directive has `jit: true`. In this
-      // case, compilation of the decorator is skipped. Returning an empty object signifies
-      // that no analysis was produced.
+    // `extractDirectiveMetadata` returns `jitForced = true` when the `@Component` has
+    // set `jit: true`. In this case, compilation of the decorator is skipped. Returning
+    // an empty object signifies that no analysis was produced.
+    if (directiveResult.jitForced) {
+      this.jitDeclarationRegistry.jitDeclarations.add(node);
       return {};
     }
 

packages/compiler-cli/src/ngtsc/annotations/component/test/component_spec.ts

@@ -39,6 +39,7 @@ import {getDeclaration, makeProgram} from '../../../testing';
 import {CompilationMode} from '../../../transform';
 import {
   InjectableClassRegistry,
+  JitDeclarationRegistry,
   NoopReferencesRegistry,
   ResourceLoader,
   ResourceLoaderContext,
@@ -105,6 +106,7 @@ function setup(
   );
   const resourceLoader = new StubResourceLoader();
   const importTracker = new ImportedSymbolsTracker();
+  const jitDeclarationRegistry = new JitDeclarationRegistry();
 
   const handler = new ComponentDecoratorHandler(
     reflectionHost,
@@ -144,6 +146,7 @@ function setup(
     /* enableBlockSyntax */ true,
     /* enableLetSyntax */ true,
     /* localCompilationExtraImportsTracker */ null,
+    jitDeclarationRegistry,
   );
   return {reflectionHost, handler, resourceLoader, metaRegistry};
 }

packages/compiler-cli/src/ngtsc/annotations/directive/src/handler.ts

@@ -76,6 +76,7 @@ import {
 
 import {extractDirectiveMetadata} from './shared';
 import {DirectiveSymbol} from './symbol';
+import {JitDeclarationRegistry} from '../../common/src/jit_declaration_registry';
 
 const FIELD_DECORATORS = [
   'Input',
@@ -133,7 +134,7 @@ export class DirectiveDecoratorHandler
     private importTracker: ImportedSymbolsTracker,
     private includeClassMetadata: boolean,
     private readonly compilationMode: CompilationMode,
-    private readonly generateExtraImportsInLocalMode: boolean,
+    private readonly jitDeclarationRegistry: JitDeclarationRegistry,
   ) {}
 
   readonly precedence = HandlerPrecedence.PRIMARY;
@@ -189,9 +190,14 @@ export class DirectiveDecoratorHandler
       this.compilationMode,
       /* defaultSelector */ null,
     );
-    if (directiveResult === undefined) {
+    // `extractDirectiveMetadata` returns `jitForced = true` when the `@Directive` has
+    // set `jit: true`. In this case, compilation of the decorator is skipped. Returning
+    // an empty object signifies that no analysis was produced.
+    if (directiveResult.jitForced) {
+      this.jitDeclarationRegistry.jitDeclarations.add(node);
       return {};
     }
+
     const analysis = directiveResult.metadata;
 
     let providersRequiringFactory: Set<Reference<ClassDeclaration>> | null = null;

packages/compiler-cli/src/ngtsc/annotations/directive/src/shared.ts

@@ -117,6 +117,7 @@ export function extractDirectiveMetadata(
   defaultSelector: string | null,
 ):
   | {
+      jitForced: false;
       decorator: Map<string, ts.Expression>;
       metadata: R3DirectiveMetadata;
       inputs: ClassPropertyMapping<InputMapping>;
@@ -125,7 +126,7 @@ export function extractDirectiveMetadata(
       hostDirectives: HostDirectiveMeta[] | null;
       rawHostDirectives: ts.Expression | null;
     }
-  | undefined {
+  | {jitForced: true} {
   let directive: Map<string, ts.Expression>;
   if (decorator.args === null || decorator.args.length === 0) {
     directive = new Map<string, ts.Expression>();
@@ -149,7 +150,7 @@ export function extractDirectiveMetadata(
 
   if (directive.has('jit')) {
     // The only allowed value is true, so there's no need to expand further.
-    return undefined;
+    return {jitForced: true};
   }
 
   const members = reflector.getMembersOfClass(clazz);
@@ -407,6 +408,7 @@ export function extractDirectiveMetadata(
       null,
   };
   return {
+    jitForced: false,
     decorator: directive,
     metadata,
     inputs,

packages/compiler-cli/src/ngtsc/annotations/directive/test/directive_spec.ts

@@ -29,7 +29,11 @@ import {
 import {LocalModuleScopeRegistry, MetadataDtsModuleScopeResolver} from '../../../scope';
 import {getDeclaration, makeProgram} from '../../../testing';
 import {CompilationMode} from '../../../transform';
-import {InjectableClassRegistry, NoopReferencesRegistry} from '../../common';
+import {
+  InjectableClassRegistry,
+  JitDeclarationRegistry,
+  NoopReferencesRegistry,
+} from '../../common';
 import {DirectiveDecoratorHandler} from '../index';
 
 runInEachFileSystem(() => {
@@ -190,6 +194,8 @@ runInEachFileSystem(() => {
     );
     const injectableRegistry = new InjectableClassRegistry(reflectionHost, /* isCore */ false);
     const importTracker = new ImportedSymbolsTracker();
+    const jitDeclarationRegistry = new JitDeclarationRegistry();
+
     const handler = new DirectiveDecoratorHandler(
       reflectionHost,
       evaluator,
@@ -207,7 +213,7 @@ runInEachFileSystem(() => {
       importTracker,
       /*includeClassMetadata*/ true,
       /*compilationMode */ CompilationMode.FULL,
-      /*generateExtraImportsInLocalMode*/ false,
+      jitDeclarationRegistry,
     );
 
     const DirNode = getDeclaration(program, _('/entry.ts'), dirName, isNamedClassDeclaration);

packages/compiler-cli/src/ngtsc/annotations/index.ts

@@ -16,6 +16,7 @@ export {
   ReferencesRegistry,
   ResourceLoader,
   ResourceLoaderContext,
+  JitDeclarationRegistry,
 } from './common';
 export {ComponentDecoratorHandler} from './component';
 export {

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

@@ -44,6 +44,7 @@ ts_library(
         "//packages/compiler-cli/src/ngtsc/util",
         "//packages/compiler-cli/src/ngtsc/validation",
         "//packages/compiler-cli/src/ngtsc/xi18n",
+        "//packages/compiler-cli/src/transformers/jit_transforms",
         "@npm//@types/semver",
         "@npm//semver",
         "@npm//typescript",

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

@@ -18,7 +18,7 @@ import {
   PipeDecoratorHandler,
   ReferencesRegistry,
 } from '../../annotations';
-import {InjectableClassRegistry} from '../../annotations/common';
+import {InjectableClassRegistry, JitDeclarationRegistry} from '../../annotations/common';
 import {CycleAnalyzer, CycleHandlingStrategy, ImportGraph} from '../../cycles';
 import {
   COMPILER_ERRORS_WITH_GUIDES,
@@ -124,6 +124,7 @@ import {DiagnosticCategoryLabel, NgCompilerAdapter, NgCompilerOptions} from '../
 
 import {coreHasSymbol} from './core_version';
 import {coreVersionSupportsFeature} from './feature_detection';
+import {angularJitApplicationTransform} from '@angular/compiler-cli/src/transformers/jit_transforms';
 
 /**
  * State information about a compilation which is only generated once some data is requested from
@@ -145,6 +146,8 @@ interface LazyCompilationState {
   extendedTemplateChecker: ExtendedTemplateChecker | null;
   templateSemanticsChecker: TemplateSemanticsChecker | null;
   sourceFileValidator: SourceFileValidator | null;
+  jitDeclarationRegistry: JitDeclarationRegistry;
+  supportJitMode: boolean;
 
   /**
    * Only available in local compilation mode when option `generateExtraImportsInLocalMode` is set.
@@ -814,6 +817,39 @@ export class NgCompiler {
       defaultImportTracker.importPreservingTransformer(),
     ];
 
+    // If there are JIT declarations, wire up the JIT transform and efficiently
+    // run it against the target declarations.
+    if (compilation.supportJitMode && compilation.jitDeclarationRegistry.jitDeclarations.size > 0) {
+      const {jitDeclarations} = compilation.jitDeclarationRegistry;
+      const jitDeclarationsArray = Array.from(jitDeclarations);
+      const jitDeclarationOriginalNodes = new Set(
+        jitDeclarationsArray.map((d) => ts.getOriginalNode(d)),
+      );
+      const sourceFilesWithJit = new Set(
+        jitDeclarationsArray.map((d) => d.getSourceFile().fileName),
+      );
+
+      before.push((ctx) => {
+        const reflectionHost = new TypeScriptReflectionHost(this.inputProgram.getTypeChecker());
+        const jitTransform = angularJitApplicationTransform(
+          this.inputProgram,
+          compilation.isCore,
+          (node) => {
+            // Class may be synthetic at this point due to Ivy transform.
+            node = ts.getOriginalNode(node, ts.isClassDeclaration);
+            return reflectionHost.isClass(node) && jitDeclarationOriginalNodes.has(node);
+          },
+        )(ctx);
+
+        return (sourceFile) => {
+          if (!sourceFilesWithJit.has(sourceFile.fileName)) {
+            return sourceFile;
+          }
+          return jitTransform(sourceFile);
+        };
+      });
+    }
+
     const afterDeclarations: ts.TransformerFactory<ts.SourceFile>[] = [];
 
     // In local compilation mode we don't make use of .d.ts files for Angular compilation, so their
@@ -1361,6 +1397,8 @@ export class NgCompiler {
       );
     }
 
+    const jitDeclarationRegistry = new JitDeclarationRegistry();
+
     // Set up the IvyCompilation, which manages state for the Ivy transformer.
     const handlers: DecoratorHandler<unknown, unknown, SemanticSymbol | null, unknown>[] = [
       new ComponentDecoratorHandler(
@@ -1401,6 +1439,7 @@ export class NgCompiler {
         this.enableBlockSyntax,
         this.enableLetSyntax,
         localCompilationExtraImportsTracker,
+        jitDeclarationRegistry,
       ),
 
       // TODO(alxhub): understand why the cast here is necessary (something to do with `null`
@@ -1422,7 +1461,7 @@ export class NgCompiler {
         importTracker,
         supportTestBed,
         compilationMode,
-        !!this.options.generateExtraImportsInLocalMode,
+        jitDeclarationRegistry,
       ) as Readonly<DecoratorHandler<unknown, unknown, SemanticSymbol | null, unknown>>,
       // Pipe handler must be before injectable handler in list so pipe factories are printed
       // before injectable factories (so injectable factories can delegate to them)
@@ -1545,8 +1584,10 @@ export class NgCompiler {
       resourceRegistry,
       extendedTemplateChecker,
       localCompilationExtraImportsTracker,
+      jitDeclarationRegistry,
       templateSemanticsChecker,
       sourceFileValidator,
+      supportJitMode,
     };
   }
 }