perf(core): avoid recursive scope recalculation when TestBed.overrideModule is used (#35454)

Currently if TestBed detects that TestBed.overrideModule was used for module X, transitive scopes are recalculated recursively for all modules that X imports and previously calculated data (stored in cache) is ignored. This behavior was introduced in https://github.com/angular/angular/pull/33787 to fix stale transitive scopes issue (cache was not updated if module overrides are present).

The perf issue comes from a "diamond" problem, where module X is overridden which imports modules A and B, which both import module C. Under previous logic, module C gets its transitive deps recomputed multiple times, during the recompute for both A and B. For deep graphs and big common/shared modules this can be super costly.

This commit updates the logic to recalculate ransitive scopes for the overridden module, while keeping previously calculated scopes of other modules untouched.

PR Close #35454

packages/core/src/render3/jit/module.ts

@@ -427,25 +427,19 @@ export function patchComponentDefWithScope<C>(
 /**
  * Compute the pair of transitive scopes (compilation scope and exported scope) for a given module.
  *
- * By default this operation is memoized and the result is cached on the module's definition. You
+ * This operation is memoized and the result is cached on the module's definition. This function can
- * can avoid memoization and previously stored results (if available) by providing the second
+ * be called on modules with components that have not fully compiled yet, but the result should not
- * argument with the `true` value (forcing transitive scopes recalculation).
+ * be used until they have.
- *
- * This function can be called on modules with components that have not fully compiled yet, but the
- * result should not be used until they have.
  *
  * @param moduleType module that transitive scope should be calculated for.
- * @param forceRecalc flag that indicates whether previously calculated and memoized values should
- * be ignored and transitive scope to be fully recalculated.
  */
-export function transitiveScopesFor<T>(
+export function transitiveScopesFor<T>(moduleType: Type<T>): NgModuleTransitiveScopes {
-    moduleType: Type<T>, forceRecalc: boolean = false): NgModuleTransitiveScopes {
   if (!isNgModule(moduleType)) {
     throw new Error(`${moduleType.name} does not have a module def (ɵmod property)`);
   }
   const def = getNgModuleDef(moduleType) !;
 
-  if (!forceRecalc && def.transitiveCompileScopes !== null) {
+  if (def.transitiveCompileScopes !== null) {
     return def.transitiveCompileScopes;
   }
 
@@ -486,7 +480,7 @@ export function transitiveScopesFor<T>(
 
     // When this module imports another, the imported module's exported directives and pipes are
     // added to the compilation scope of this module.
-    const importedScope = transitiveScopesFor(importedType, forceRecalc);
+    const importedScope = transitiveScopesFor(importedType);
     importedScope.exported.directives.forEach(entry => scopes.compilation.directives.add(entry));
     importedScope.exported.pipes.forEach(entry => scopes.compilation.pipes.add(entry));
   });
@@ -505,7 +499,7 @@ export function transitiveScopesFor<T>(
     if (isNgModule(exportedType)) {
       // When this module exports another, the exported module's exported directives and pipes are
       // added to both the compilation and exported scopes of this module.
-      const exportedScope = transitiveScopesFor(exportedType, forceRecalc);
+      const exportedScope = transitiveScopesFor(exportedType);
       exportedScope.exported.directives.forEach(entry => {
         scopes.compilation.directives.add(entry);
         scopes.exported.directives.add(entry);
@@ -521,9 +515,7 @@ export function transitiveScopesFor<T>(
     }
   });
 
-  if (!forceRecalc) {
+  def.transitiveCompileScopes = scopes;
-    def.transitiveCompileScopes = scopes;
-  }
   return scopes;
 }
 

packages/core/test/test_bed_spec.ts

@@ -889,15 +889,23 @@ describe('TestBed', () => {
               {set: {template: `<span someDirective>{{'hello' | somePipe}}</span>`}});
           TestBed.createComponent(SomeComponent);
 
-          const defBeforeReset = (SomeComponent as any).ɵcmp;
+          const cmpDefBeforeReset = (SomeComponent as any).ɵcmp;
-          expect(defBeforeReset.pipeDefs().length).toEqual(1);
+          expect(cmpDefBeforeReset.pipeDefs().length).toEqual(1);
-          expect(defBeforeReset.directiveDefs().length).toEqual(2);  // directive + component
+          expect(cmpDefBeforeReset.directiveDefs().length).toEqual(2);  // directive + component
+
+          const modDefBeforeReset = (SomeModule as any).ɵmod;
+          const transitiveScope = modDefBeforeReset.transitiveCompileScopes.compilation;
+          expect(transitiveScope.pipes.size).toEqual(1);
+          expect(transitiveScope.directives.size).toEqual(2);
 
           TestBed.resetTestingModule();
 
-          const defAfterReset = (SomeComponent as any).ɵcmp;
+          const cmpDefAfterReset = (SomeComponent as any).ɵcmp;
-          expect(defAfterReset.pipeDefs).toBe(null);
+          expect(cmpDefAfterReset.pipeDefs).toBe(null);
-          expect(defAfterReset.directiveDefs).toBe(null);
+          expect(cmpDefAfterReset.directiveDefs).toBe(null);
+
+          const modDefAfterReset = (SomeModule as any).ɵmod;
+          expect(modDefAfterReset.transitiveCompileScopes).toBe(null);
         });
 
         it('should cleanup ng defs for classes with no ng annotations (in case of inheritance)',

packages/core/testing/src/r3_test_bed_compiler.ts

@@ -355,8 +355,11 @@ export class R3TestBedCompiler {
             // are present, always re-calculate transitive scopes to have the most up-to-date
             // information available. The `moduleToScope` map avoids repeated re-calculation of
             // scopes for the same module.
-            const forceRecalc = !isTestingModule && this.hasModuleOverrides;
+            if (!isTestingModule && this.hasModuleOverrides) {
-            moduleToScope.set(moduleType, transitiveScopesFor(realType, forceRecalc));
+              this.storeFieldOfDefOnType(moduleType as any, NG_MOD_DEF, 'transitiveCompileScopes');
+              (moduleType as any)[NG_MOD_DEF].transitiveCompileScopes = null;
+            }
+            moduleToScope.set(moduleType, transitiveScopesFor(realType));
           }
           return moduleToScope.get(moduleType) !;
         };