fix(compiler-cli): run JIT transforms on @NgModule classes with jit: true (#57212)

This commit is similar to 98ed5b609e, 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

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 {};
     }

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};

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

@@ -1512,6 +1512,7 @@ export class NgCompiler {
         supportJitMode,
         compilationMode,
         localCompilationExtraImportsTracker,
+        jitDeclarationRegistry,
       ),
     ];
 

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:`.