fix(ivy): compute transitive scopes from NgModuleDef only (#24334)

Previously, the transitive scopes of an NgModuleDef were computed
during execution of the @NgModule decorator. This meant that JIT-
compiled modules could only import other JIT-compiled modules, as
the import mechanism relied on the calculation of transitive scopes
to already have happened for the imported module.

This change moves computation of transitive scopes to a function
`transitiveScopesFor` (and makes it lazy). This opens the door for
AOT -> JIT or JIT -> AOT imports, as transitive scopes for AOT
modules can be calculated when needed by JIT, and AOT modules can
also write expressions that call `transitiveScopesFor` when
importing a JIT-compiled module.

PR Close #24334

packages/core/src/metadata/ng_module.ts

@@ -13,6 +13,20 @@ import {R3_COMPILE_NGMODULE} from '../ivy_switch';
 import {Type} from '../type';
 import {TypeDecorator, makeDecorator} from '../util/decorators';
 
+/**
+ * Represents the expansion of an `NgModule` into its scopes.
+ *
+ * A scope is a set of directives and pipes that are visible in a particular context. Each
+ * `NgModule` has two scopes. The `compilation` scope is the set of directives and pipes that will
+ * be recognized in the templates of components declared by the module. The `exported` scope is the
+ * set of directives and pipes exported by a module (that is, module B's exported scope gets added
+ * to module A's compilation scope when module A imports B).
+ */
+export interface NgModuleTransitiveScopes {
+  compilation: {directives: Set<any>; pipes: Set<any>;};
+  exported: {directives: Set<any>; pipes: Set<any>;};
+}
+
 export interface NgModuleDef<T> {
   type: T;
   bootstrap: Type<any>[];
@@ -20,7 +34,12 @@ export interface NgModuleDef<T> {
   imports: Type<any>[];
   exports: Type<any>[];
 
-  transitiveCompileScope: {directives: any[]; pipes: any[];}|undefined;
+  /**
+   * Cached value of computed `transitiveCompileScopes` for this module.
+   *
+   * This should never be read directly, but accessed via `transitiveScopesFor`.
+   */
+  transitiveCompileScopes: NgModuleTransitiveScopes|null;
 }
 
 export function defineNgModule<T>(def: {type: T} & Partial<NgModuleDef<T>>): never {
@@ -30,7 +49,7 @@ export function defineNgModule<T>(def: {type: T} & Partial<NgModuleDef<T>>): nev
     declarations: def.declarations || [],
     imports: def.imports || [],
     exports: def.exports || [],
-    transitiveCompileScope: undefined,
+    transitiveCompileScopes: null,
   };
   return res as never;
 }

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

@@ -13,6 +13,8 @@ import {ReflectionCapabilities} from '../../reflection/reflection_capabilities';
 import {Type} from '../../type';
 
 import {angularCoreEnv} from './environment';
+import {NG_COMPONENT_DEF, NG_DIRECTIVE_DEF} from './fields';
+import {patchComponentDefWithScope} from './module';
 import {getReflect, reflectDependencies} from './util';
 
 let _pendingPromises: Promise<void>[] = [];
@@ -34,7 +36,7 @@ export function compileComponent(type: Type<any>, metadata: Component): Promise<
   const templateStr = metadata.template;
 
   let def: any = null;
-  Object.defineProperty(type, 'ngComponentDef', {
+  Object.defineProperty(type, NG_COMPONENT_DEF, {
     get: () => {
       if (def === null) {
         // The ConstantPool is a requirement of the JIT'er.
@@ -61,6 +63,14 @@ export function compileComponent(type: Type<any>, metadata: Component): Promise<
 
         def = jitExpression(
             res.expression, angularCoreEnv, `ng://${type.name}/ngComponentDef.js`, constantPool);
+
+        // If component compilation is async, then the @NgModule annotation which declares the
+        // component may execute and set an ngSelectorScope property on the component type. This
+        // allows the component to patch itself with directiveDefs from the module after it finishes
+        // compiling.
+        if (hasSelectorScope(type)) {
+          patchComponentDefWithScope(def, type.ngSelectorScope);
+        }
       }
       return def;
     },
@@ -69,6 +79,11 @@ export function compileComponent(type: Type<any>, metadata: Component): Promise<
   return null;
 }
 
+function hasSelectorScope<T>(component: Type<T>): component is Type<T>&
+    {ngSelectorScope: Type<any>} {
+  return (component as{ngSelectorScope?: any}).ngSelectorScope !== undefined;
+}
+
 /**
  * Compile an Angular directive according to its decorator metadata, and patch the resulting
  * ngDirectiveDef onto the component type.
@@ -78,7 +93,7 @@ export function compileComponent(type: Type<any>, metadata: Component): Promise<
  */
 export function compileDirective(type: Type<any>, directive: Directive): Promise<void>|null {
   let def: any = null;
-  Object.defineProperty(type, 'ngDirectiveDef', {
+  Object.defineProperty(type, NG_DIRECTIVE_DEF, {
     get: () => {
       if (def === null) {
         const constantPool = new ConstantPool();

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

@@ -0,0 +1,16 @@
+/**
+ * @license
+ * Copyright Google Inc. 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 {getClosureSafeProperty} from '../../util/property';
+
+const TARGET = {} as any;
+
+export const NG_COMPONENT_DEF = getClosureSafeProperty({ngComponentDef: TARGET}, TARGET);
+export const NG_DIRECTIVE_DEF = getClosureSafeProperty({ngDirectiveDef: TARGET}, TARGET);
+export const NG_PIPE_DEF = getClosureSafeProperty({ngPipeDef: TARGET}, TARGET);
+export const NG_MODULE_DEF = getClosureSafeProperty({ngModuleDef: TARGET}, TARGET);

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

@@ -8,80 +8,26 @@
 
 import {Expression, R3NgModuleMetadata, WrappedNodeExpr, compileNgModule as compileR3NgModule, jitExpression} from '@angular/compiler';
 
-import {ModuleWithProviders, NgModule, NgModuleDef} from '../../metadata/ng_module';
+import {ModuleWithProviders, NgModule, NgModuleDef, NgModuleTransitiveScopes} from '../../metadata/ng_module';
 import {Type} from '../../type';
 import {ComponentDef} from '../interfaces/definition';
-import {flatten} from '../util';
 
 import {angularCoreEnv} from './environment';
+import {NG_COMPONENT_DEF, NG_DIRECTIVE_DEF, NG_MODULE_DEF, NG_PIPE_DEF} from './fields';
 
 const EMPTY_ARRAY: Type<any>[] = [];
 
 export function compileNgModule(type: Type<any>, ngModule: NgModule): void {
-  const meta: R3NgModuleMetadata = {
-    type: wrap(type),
-    bootstrap: flatten(ngModule.bootstrap || EMPTY_ARRAY).map(wrap),
-    declarations: flatten(ngModule.declarations || EMPTY_ARRAY).map(wrap),
-    imports: flatten(ngModule.imports || EMPTY_ARRAY).map(expandModuleWithProviders).map(wrap),
-    exports: flatten(ngModule.exports || EMPTY_ARRAY).map(expandModuleWithProviders).map(wrap),
-    emitInline: true,
-  };
-
-  // Compute transitiveCompileScope
-  const transitiveCompileScope = {
-    directives: new Set<any>(),
-    pipes: new Set<any>(),
-    modules: new Set<any>(),
-  };
-
-  function addExportsFrom(module: Type<any>& {ngModuleDef: NgModuleDef<any>}): void {
-    if (!transitiveCompileScope.modules.has(module)) {
-      module.ngModuleDef.exports.forEach((exp: any) => {
-        if (isNgModule(exp)) {
-          addExportsFrom(exp);
-        } else if (exp.ngPipeDef) {
-          transitiveCompileScope.pipes.add(exp);
-        } else {
-          transitiveCompileScope.directives.add(exp);
-        }
-      });
-    }
-  }
-
-  flatten([
-    (ngModule.imports || EMPTY_ARRAY), (ngModule.exports || EMPTY_ARRAY)
-  ]).forEach(importExport => {
-    const maybeModule = expandModuleWithProviders(importExport);
-    if (isNgModule(maybeModule)) {
-      addExportsFrom(maybeModule);
-    }
-  });
-
-  flatten(ngModule.declarations || EMPTY_ARRAY).forEach(decl => {
-    if (decl.ngPipeDef) {
-      transitiveCompileScope.pipes.add(decl);
-    } else if (decl.ngDirectiveDef) {
-      transitiveCompileScope.directives.add(decl);
-    } else if (decl.ngComponentDef) {
-      transitiveCompileScope.directives.add(decl);
-      patchComponentWithScope(decl, type as any);
-    } else {
-      // A component that has not been compiled yet because the template is being fetched
-      // we need to store a reference to the module to update the selector scope after
-      // the component gets compiled
-      transitiveCompileScope.directives.add(decl);
-      decl.ngSelectorScope = type;
-    }
-  });
+  const declarations: Type<any>[] = flatten(ngModule.declarations || EMPTY_ARRAY);
 
   let def: any = null;
-  Object.defineProperty(type, 'ngModuleDef', {
+  Object.defineProperty(type, NG_MODULE_DEF, {
     get: () => {
       if (def === null) {
         const meta: R3NgModuleMetadata = {
           type: wrap(type),
           bootstrap: flatten(ngModule.bootstrap || EMPTY_ARRAY).map(wrap),
-          declarations: flatten(ngModule.declarations || EMPTY_ARRAY).map(wrap),
+          declarations: declarations.map(wrap),
           imports:
               flatten(ngModule.imports || EMPTY_ARRAY).map(expandModuleWithProviders).map(wrap),
           exports:
@@ -90,25 +36,141 @@ export function compileNgModule(type: Type<any>, ngModule: NgModule): void {
         };
         const res = compileR3NgModule(meta);
         def = jitExpression(res.expression, angularCoreEnv, `ng://${type.name}/ngModuleDef.js`);
-        def.transitiveCompileScope = {
-          directives: Array.from(transitiveCompileScope.directives),
-          pipes: Array.from(transitiveCompileScope.pipes),
-        };
       }
       return def;
     },
   });
+
+  declarations.forEach(declaration => {
+    // Some declared components may be compiled asynchronously, and thus may not have their
+    // ngComponentDef set yet. If this is the case, then a reference to the module is written into
+    // the `ngSelectorScope` property of the declared type.
+    if (declaration.hasOwnProperty(NG_COMPONENT_DEF)) {
+      // An `ngComponentDef` field exists - go ahead and patch the component directly.
+      patchComponentDefWithScope(
+          (declaration as Type<any>& {ngComponentDef: ComponentDef<any>}).ngComponentDef, type);
+    } else if (
+        !declaration.hasOwnProperty(NG_DIRECTIVE_DEF) && !declaration.hasOwnProperty(NG_PIPE_DEF)) {
+      // Set `ngSelectorScope` for future reference when the component compilation finishes.
+      (declaration as Type<any>& {ngSelectorScope?: any}).ngSelectorScope = type;
+    }
+  });
 }
 
-export function patchComponentWithScope<C, M>(
-    component: Type<C>& {ngComponentDef: ComponentDef<C>},
-    module: Type<M>& {ngModuleDef: NgModuleDef<M>}) {
-  component.ngComponentDef.directiveDefs = () =>
-      module.ngModuleDef.transitiveCompileScope !.directives
-          .map(dir => dir.ngDirectiveDef || dir.ngComponentDef)
-          .filter(def => !!def);
-  component.ngComponentDef.pipeDefs = () =>
-      module.ngModuleDef.transitiveCompileScope !.pipes.map(pipe => pipe.ngPipeDef);
+/**
+ * Patch the definition of a component with directives and pipes from the compilation scope of
+ * a given module.
+ */
+export function patchComponentDefWithScope<C, M>(componentDef: ComponentDef<C>, module: Type<M>) {
+  componentDef.directiveDefs = () => Array.from(transitiveScopesFor(module).compilation.directives)
+                                         .map(dir => dir.ngDirectiveDef || dir.ngComponentDef)
+                                         .filter(def => !!def);
+  componentDef.pipeDefs = () =>
+      Array.from(transitiveScopesFor(module).compilation.pipes).map(pipe => pipe.ngPipeDef);
+}
+
+/**
+ * Compute the pair of transitive scopes (compilation scope and exported scope) for a given module.
+ *
+ * This operation is memoized and the result is cached on the module's definition. It can be called
+ * on modules with components that have not fully compiled yet, but the result should not be used
+ * until they have.
+ */
+export function transitiveScopesFor<T>(moduleType: Type<T>): NgModuleTransitiveScopes {
+  if (!isNgModule(moduleType)) {
+    throw new Error(`${moduleType.name} does not have an ngModuleDef`);
+  }
+  const def = moduleType.ngModuleDef;
+
+  if (def.transitiveCompileScopes !== null) {
+    return def.transitiveCompileScopes;
+  }
+
+  const scopes: NgModuleTransitiveScopes = {
+    compilation: {
+      directives: new Set<any>(),
+      pipes: new Set<any>(),
+    },
+    exported: {
+      directives: new Set<any>(),
+      pipes: new Set<any>(),
+    },
+  };
+
+  def.declarations.forEach(declared => {
+    const declaredWithDefs = declared as Type<any>& { ngPipeDef?: any; };
+
+    if (declaredWithDefs.ngPipeDef !== undefined) {
+      scopes.compilation.pipes.add(declared);
+    } else {
+      // Either declared has an ngComponentDef or ngDirectiveDef, or it's a component which hasn't
+      // had its template compiled yet. In either case, it gets added to the compilation's
+      // directives.
+      scopes.compilation.directives.add(declared);
+    }
+  });
+
+  def.imports.forEach(<I>(imported: Type<I>) => {
+    let importedTyped = imported as Type<I>& {
+      // If imported is an @NgModule:
+      ngModuleDef?: NgModuleDef<I>;
+    };
+
+    if (!isNgModule<I>(importedTyped)) {
+      throw new Error(`Importing ${importedTyped.name} which does not have an ngModuleDef`);
+    }
+
+    // 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(importedTyped);
+    importedScope.exported.directives.forEach(entry => scopes.compilation.directives.add(entry));
+    importedScope.exported.pipes.forEach(entry => scopes.compilation.pipes.add(entry));
+  });
+
+  def.exports.forEach(<E>(exported: Type<E>) => {
+    const exportedTyped = exported as Type<E>& {
+      // Components, Directives, NgModules, and Pipes can all be exported.
+      ngComponentDef?: any;
+      ngDirectiveDef?: any;
+      ngModuleDef?: NgModuleDef<E>;
+      ngPipeDef?: any;
+    };
+
+    // Either the type is a module, a pipe, or a component/directive (which may not have an
+    // ngComponentDef as it might be compiled asynchronously).
+    if (isNgModule(exportedTyped)) {
+      // 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(exportedTyped);
+      exportedScope.exported.directives.forEach(entry => {
+        scopes.compilation.directives.add(entry);
+        scopes.exported.directives.add(entry);
+      });
+      exportedScope.exported.pipes.forEach(entry => {
+        scopes.compilation.pipes.add(entry);
+        scopes.exported.pipes.add(entry);
+      });
+    } else if (exportedTyped.ngPipeDef !== undefined) {
+      scopes.exported.pipes.add(exportedTyped);
+    } else {
+      scopes.exported.directives.add(exportedTyped);
+    }
+  });
+
+  def.transitiveCompileScopes = scopes;
+  return scopes;
+}
+
+function flatten<T>(values: any[]): T[] {
+  const out: T[] = [];
+  values.forEach(value => {
+    if (Array.isArray(value)) {
+      out.push(...flatten<T>(value));
+    } else {
+      out.push(value);
+    }
+  });
+  return out;
 }
 
 function expandModuleWithProviders(value: Type<any>| ModuleWithProviders): Type<any> {
@@ -123,9 +185,9 @@ function wrap(value: Type<any>): Expression {
 }
 
 function isModuleWithProviders(value: any): value is ModuleWithProviders {
-  return value.ngModule !== undefined;
+  return (value as{ngModule?: any}).ngModule !== undefined;
 }
 
-function isNgModule(value: any): value is Type<any>&{ngModuleDef: NgModuleDef<any>} {
-  return value.ngModuleDef !== undefined;
+function isNgModule<T>(value: Type<T>): value is Type<T>&{ngModuleDef: NgModuleDef<T>} {
+  return (value as{ngModuleDef?: NgModuleDef<T>}).ngModuleDef !== undefined;
 }

packages/core/test/render3/ivy/BUILD.bazel

@@ -28,6 +28,9 @@ jasmine_node_test(
     bootstrap = [
         "angular/packages/core/test/render3/load_domino",
     ],
+    tags = [
+        "ivy-jit",
+    ],
     deps = [
         ":ivy_node_lib",
     ],