From e11c0c42d2cbcdf8a5d75a4e24a6a5dbed33943e Mon Sep 17 00:00:00 2001
From: Paul Gschwendtner <paulgschwendtner@gmail.com>
Date: Wed, 31 Jul 2024 13:27:22 +0000
Subject: [PATCH] fix(compiler-cli): run JIT transforms on `@NgModule` classes
 with `jit: true` (#57212)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This commit is similar to 98ed5b609e76d3d2b464abfe49d70413c54d3eee, and
makes use of the preparation work implemented there.

Similar to directives and components marked via `jit: true`, we also
need to do the same for JIT marked `@NgModule` classes. This is mostly
important for downleveling of decorators to support dependency injection
of such classes.

Inside Google3, migrating from `ts_library` to `ng_module` turns of
decorator downleveling, so the `jit: true` for NgModule's is implicitly
requesting/reliant on this transform— as expected.

PR Close #57212
---
 .../annotations/ng_module/src/handler.ts      |  3 ++
 .../ng_module/test/ng_module_spec.ts          |  8 ++++-
 .../src/ngtsc/core/src/compiler.ts            |  1 +
 .../compiler-cli/test/ngtsc/ngtsc_spec.ts     | 31 +++++++++++++++++++
 4 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/packages/compiler-cli/src/ngtsc/annotations/ng_module/src/handler.ts b/packages/compiler-cli/src/ngtsc/annotations/ng_module/src/handler.ts
index 0370df99a5e..5853df12479 100644
--- a/packages/compiler-cli/src/ngtsc/annotations/ng_module/src/handler.ts
+++ b/packages/compiler-cli/src/ngtsc/annotations/ng_module/src/handler.ts
@@ -100,6 +100,7 @@ import {
   getValidConstructorDependencies,
   InjectableClassRegistry,
   isExpressionForwardReference,
+  JitDeclarationRegistry,
   ReferencesRegistry,
   resolveProvidersRequiringFactory,
   toR3Reference,
@@ -285,6 +286,7 @@ export class NgModuleDecoratorHandler
     private includeSelectorScope: boolean,
     private readonly compilationMode: CompilationMode,
     private readonly localCompilationExtraImportsTracker: LocalCompilationExtraImportsTracker | null,
+    private readonly jitDeclarationRegistry: JitDeclarationRegistry,
   ) {}
 
   readonly precedence = HandlerPrecedence.PRIMARY;
@@ -341,6 +343,7 @@ export class NgModuleDecoratorHandler
     const ngModule = reflectObjectLiteral(meta);
 
     if (ngModule.has('jit')) {
+      this.jitDeclarationRegistry.jitDeclarations.add(node);
       // The only allowed value is true, so there's no need to expand further.
       return {};
     }
diff --git a/packages/compiler-cli/src/ngtsc/annotations/ng_module/test/ng_module_spec.ts b/packages/compiler-cli/src/ngtsc/annotations/ng_module/test/ng_module_spec.ts
index 52246eeb971..8c2321329ca 100644
--- a/packages/compiler-cli/src/ngtsc/annotations/ng_module/test/ng_module_spec.ts
+++ b/packages/compiler-cli/src/ngtsc/annotations/ng_module/test/ng_module_spec.ts
@@ -28,7 +28,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 {NgModuleDecoratorHandler} from '../src/handler';
 
 function setup(program: ts.Program, compilationMode = CompilationMode.FULL) {
@@ -52,6 +56,7 @@ function setup(program: ts.Program, compilationMode = CompilationMode.FULL) {
   const refEmitter = new ReferenceEmitter([new LocalIdentifierStrategy()]);
   const injectableRegistry = new InjectableClassRegistry(reflectionHost, /* isCore */ false);
   const exportedProviderStatusResolver = new ExportedProviderStatusResolver(metaReader);
+  const jitDeclarationRegistry = new JitDeclarationRegistry();
 
   const handler = new NgModuleDecoratorHandler(
     reflectionHost,
@@ -72,6 +77,7 @@ function setup(program: ts.Program, compilationMode = CompilationMode.FULL) {
     true,
     compilationMode,
     /* localCompilationExtraImportsTracker */ null,
+    jitDeclarationRegistry,
   );
 
   return {handler, reflectionHost};
diff --git a/packages/compiler-cli/src/ngtsc/core/src/compiler.ts b/packages/compiler-cli/src/ngtsc/core/src/compiler.ts
index 0319f4f37c7..49a64d68b2a 100644
--- a/packages/compiler-cli/src/ngtsc/core/src/compiler.ts
+++ b/packages/compiler-cli/src/ngtsc/core/src/compiler.ts
@@ -1512,6 +1512,7 @@ export class NgCompiler {
         supportJitMode,
         compilationMode,
         localCompilationExtraImportsTracker,
+        jitDeclarationRegistry,
       ),
     ];
 
diff --git a/packages/compiler-cli/test/ngtsc/ngtsc_spec.ts b/packages/compiler-cli/test/ngtsc/ngtsc_spec.ts
index 3b9953c8bb8..0fe7b4b874a 100644
--- a/packages/compiler-cli/test/ngtsc/ngtsc_spec.ts
+++ b/packages/compiler-cli/test/ngtsc/ngtsc_spec.ts
@@ -645,6 +645,37 @@ runInEachFileSystem((os: string) => {
       expect(env.getContents('test.js')).not.toContain('NonJitComponent.propDecorators');
     });
 
+    it('should run JIT transform for `NgModule` marked as `jit: true`', () => {
+      env.write(
+        'test.ts',
+        `
+        import {NgModule, Injector} from '@angular/core';
+
+        @NgModule({
+          jit: true
+        })
+        class MyJitModule {
+          constructor(dep: Injector) {}
+        }
+
+        @NgModule({})
+        class NonJitModule {
+          constructor(dep: Injector) {}
+        }
+      `,
+      );
+
+      env.driveMain();
+
+      expect(trim(env.getContents('test.js'))).toContain(
+        trim(`
+        MyJitModule.ctorParameters = () => [
+          { type: Injector }
+        ];`),
+      );
+      expect(env.getContents('test.js')).not.toContain('NonJitModule.ctorParameters');
+    });
+
     // This test triggers the Tsickle compiler which asserts that the file-paths
     // are valid for the real OS. When on non-Windows systems it doesn't like paths
     // that start with `C:`.