[FEATURE set-component-template] Initial impl

Co-authored-by: Robert Jackson <me@rwjblue.com>

packages/@ember/-internals/glimmer/index.ts

@@ -395,3 +395,4 @@ export { setComponentManager, getComponentManager } from './lib/utils/custom-com
 export { setModifierManager, getModifierManager } from './lib/utils/custom-modifier-manager';
 export { capabilities as modifierCapabilties } from './lib/modifiers/custom';
 export { isSerializationFirstNode } from './lib/utils/serialization-first-node-helpers';
+export { setComponentTemplate, getComponentTemplate } from './lib/utils/component-template';

packages/@ember/-internals/glimmer/lib/component-managers/curly.ts

@@ -583,7 +583,7 @@ export class CurlyComponentDefinition implements ComponentDefinition {
     public name: string,
     public ComponentClass: any,
     public handle: Option<VMHandle>,
-    public template?: OwnedTemplate,
+    public template: Option<OwnedTemplate>,
     args?: CurriedArgs
   ) {
     const layout = template && template.asLayout();

packages/@ember/-internals/glimmer/lib/environment.ts

@@ -1,5 +1,5 @@
 import { OWNER, Owner } from '@ember/-internals/owner';
-import { constructStyleDeprecationMessage, lookupComponent } from '@ember/-internals/views';
+import { constructStyleDeprecationMessage } from '@ember/-internals/views';
 import { warn } from '@ember/debug';
 import { DEBUG } from '@glimmer/env';
 import { Option, Simple } from '@glimmer/interfaces';
@@ -57,10 +57,6 @@ export default class Environment extends GlimmerEnvironment {
     return s;
   }
 
-  lookupComponent(name: string, meta: any) {
-    return lookupComponent(meta.owner, name, meta);
-  }
-
   toConditionalReference(reference: UpdatableReference): VersionedReference<boolean> {
     return ConditionalReference.create(reference);
   }

packages/@ember/-internals/glimmer/lib/resolver.ts

@@ -1,10 +1,11 @@
 import { privatize as P } from '@ember/-internals/container';
 import { ENV } from '@ember/-internals/environment';
-import { FactoryClass, LookupOptions, Owner } from '@ember/-internals/owner';
-import { lookupComponent, lookupPartial, OwnedTemplateMeta } from '@ember/-internals/views';
+import { Factory, FactoryClass, LookupOptions, Owner } from '@ember/-internals/owner';
+import { lookupPartial, OwnedTemplateMeta } from '@ember/-internals/views';
 import {
   EMBER_GLIMMER_ANGLE_BRACKET_BUILT_INS,
   EMBER_GLIMMER_FN_HELPER,
+  EMBER_GLIMMER_SET_COMPONENT_TEMPLATE,
   EMBER_MODULE_UNIFICATION,
 } from '@ember/canary-features';
 import { assert } from '@ember/debug';
@@ -48,6 +49,8 @@ import OnModifierManager from './modifiers/on';
 import { populateMacros } from './syntax';
 import { mountHelper } from './syntax/mount';
 import { outletHelper } from './syntax/outlet';
+import { Factory as TemplateFactory, OwnedTemplate } from './template';
+import { getComponentTemplate } from './utils/component-template';
 import { getModifierManager } from './utils/custom-modifier-manager';
 import { getManager } from './utils/managers';
 import { ClassBasedHelperReference, SimpleHelperReference } from './utils/references';
@@ -63,6 +66,117 @@ function makeOptions(moduleName: string, namespace?: string): LookupOptions {
   };
 }
 
+function componentFor(
+  name: string,
+  owner: Owner,
+  options?: LookupOptions
+): Option<Factory<{}, {}>> {
+  let fullName = `component:${name}`;
+  return owner.factoryFor(fullName, options) || null;
+}
+
+function layoutFor(name: string, owner: Owner, options?: LookupOptions): Option<OwnedTemplate> {
+  let templateFullName = `template:components/${name}`;
+
+  return owner.lookup(templateFullName, options) || null;
+}
+
+function lookupModuleUnificationComponentPair(
+  owner: Owner,
+  name: string,
+  options?: LookupOptions
+): Option<LookupResult> {
+  let localComponent = componentFor(name, owner, options);
+  let localLayout = layoutFor(name, owner, options);
+
+  let globalComponent = componentFor(name, owner);
+  let globalLayout = layoutFor(name, owner);
+
+  // TODO: we shouldn't have to recheck fallback, we should have a lookup that doesn't fallback
+  if (
+    localComponent !== null &&
+    globalComponent !== null &&
+    globalComponent.class === localComponent.class
+  ) {
+    localComponent = null;
+  }
+  if (
+    localLayout !== null &&
+    globalLayout !== null &&
+    localLayout.referrer.moduleName === globalLayout.referrer.moduleName
+  ) {
+    localLayout = null;
+  }
+
+  if (localComponent !== null || localLayout !== null) {
+    return { component: localComponent, layout: localLayout } as LookupResult;
+  } else if (globalComponent !== null || globalLayout !== null) {
+    return { component: globalComponent, layout: globalLayout } as LookupResult;
+  } else {
+    return null;
+  }
+}
+
+type LookupResult =
+  | {
+      component: Factory<{}, {}>;
+      layout: TemplateFactory;
+    }
+  | {
+      component: Factory<{}, {}>;
+      layout: null;
+    }
+  | {
+      component: null;
+      layout: TemplateFactory;
+    };
+
+function lookupComponentPair(
+  owner: Owner,
+  name: string,
+  options?: LookupOptions
+): Option<LookupResult> {
+  let component = componentFor(name, owner, options);
+
+  if (EMBER_GLIMMER_SET_COMPONENT_TEMPLATE) {
+    if (component !== null && component.class !== undefined) {
+      let layout = getComponentTemplate(component.class);
+
+      if (layout !== null) {
+        return { component, layout };
+      }
+    }
+  }
+
+  let layout = layoutFor(name, owner, options);
+
+  if (component === null && layout === null) {
+    return null;
+  } else {
+    return { component, layout } as LookupResult;
+  }
+}
+
+function lookupComponent(owner: Owner, name: string, options: LookupOptions): Option<LookupResult> {
+  if (options.source || options.namespace) {
+    if (EMBER_MODULE_UNIFICATION) {
+      return lookupModuleUnificationComponentPair(owner, name, options);
+    }
+
+    let pair = lookupComponentPair(owner, name, options);
+
+    if (pair !== null) {
+      return pair;
+    }
+  }
+
+  if (EMBER_MODULE_UNIFICATION) {
+    return lookupModuleUnificationComponentPair(owner, name);
+  }
+
+  return lookupComponentPair(owner, name);
+}
+
 interface IBuiltInHelpers {
   [name: string]: Helper | undefined;
 }
@@ -113,7 +227,6 @@ export default class RuntimeResolver implements IRuntimeResolver<OwnedTemplateMe
   private builtInModifiers: IBuiltInModifiers;
 
   private componentDefinitionCache: Map<object, ComponentDefinition | null> = new Map();
-  private customManagerCache: Map<string, ManagerDelegate<Opaque>> = new Map();
 
   public componentDefinitionCount = 0;
   public helperDefinitionCount = 0;
@@ -309,51 +422,62 @@ export default class RuntimeResolver implements IRuntimeResolver<OwnedTemplateMe
       name = parsed.name;
       namespace = parsed.namespace;
     }
-    let { layout, component } = lookupComponent(owner, name, makeOptions(moduleName, namespace));
+    let pair = lookupComponent(owner, name, makeOptions(moduleName, namespace));
+    if (pair === null) {
+      return null;
+    }
 
-    let key = component === undefined ? layout : component;
+    let layout: Option<OwnedTemplate> = null;
+    let key: object;
 
-    if (key === undefined) {
-      return null;
+    if (pair.component === null) {
+      key = layout = pair.layout!(owner);
+    } else {
+      key = pair.component;
     }
 
     let cachedComponentDefinition = this.componentDefinitionCache.get(key);
     if (cachedComponentDefinition !== undefined) {
       return cachedComponentDefinition;
     }
 
+    if (layout === null && pair.layout !== null) {
+      layout = pair.layout(owner);
+    }
+
     let finalizer = _instrumentStart('render.getComponentDefinition', instrumentationPayload, name);
 
     let definition: Option<ComponentDefinition> = null;
 
-    if (layout !== undefined && component === undefined && ENV._TEMPLATE_ONLY_GLIMMER_COMPONENTS) {
-      definition = new TemplateOnlyComponentDefinition(layout(owner));
+    if (pair.component === null && ENV._TEMPLATE_ONLY_GLIMMER_COMPONENTS) {
+      definition = new TemplateOnlyComponentDefinition(layout!);
     }
 
-    if (component !== undefined && component.class !== undefined) {
-      let wrapper = getManager(component.class);
+    if (pair.component !== null) {
+      assert(`missing component class ${name}`, pair.component.class !== undefined);
+
+      let ComponentClass = pair.component.class!;
+      let wrapper = getManager(ComponentClass);
 
       if (wrapper !== null && wrapper.type === 'component') {
         let { factory } = wrapper;
 
         if (wrapper.internal) {
-          assert(`missing layout for internal component ${name}`, layout !== undefined);
+          assert(`missing layout for internal component ${name}`, pair.layout !== null);
 
           definition = new InternalComponentDefinition(
             factory(owner) as InternalComponentManager<Opaque>,
-            component.class,
-            layout!(owner)
+            ComponentClass as Factory<any, any>,
+            layout!
           );
         } else {
-          if (layout === undefined) {
-            layout = owner.lookup(P`template:components/-default`);
-          }
-
           definition = new CustomManagerDefinition(
             name,
-            component,
+            pair.component,
             factory(owner) as ManagerDelegate<Opaque>,
-            layout!(owner)
+            layout !== null
+              ? layout
+              : owner.lookup<TemplateFactory>(P`template:components/-default`)!(owner)
           );
         }
       }
@@ -362,25 +486,14 @@ export default class RuntimeResolver implements IRuntimeResolver<OwnedTemplateMe
     if (definition === null) {
       definition = new CurlyComponentDefinition(
         name,
-        component || owner.factoryFor(P`component:-default`),
+        pair.component || owner.factoryFor(P`component:-default`),
         null,
-        layout !== undefined ? layout(owner) : undefined
+        layout
       );
     }
 
     finalizer();
     this.componentDefinitionCache.set(key, definition);
     return definition;
   }
-
-  _lookupComponentManager(owner: Owner, managerId: string): ManagerDelegate<Opaque> {
-    if (this.customManagerCache.has(managerId)) {
-      return this.customManagerCache.get(managerId)!;
-    }
-    let delegate = owner.lookup<ManagerDelegate<Opaque>>(`component-manager:${managerId}`);
-
-    this.customManagerCache.set(managerId, delegate);
-
-    return delegate;
-  }
 }

packages/@ember/-internals/glimmer/lib/template.ts

@@ -32,7 +32,7 @@ export default function template(json: StaticTemplate): Factory {
 
     if (result === undefined) {
       counters.cacheMiss++;
-      let compiler: LazyCompiler<StaticTemplateMeta> = owner.lookup(TEMPLATE_COMPILER_MAIN);
+      let compiler = owner.lookup<LazyCompiler<StaticTemplateMeta>>(TEMPLATE_COMPILER_MAIN)!;
       result = glimmerFactory.create(compiler, { owner });
       cache.set(owner, result);
     } else {

packages/@ember/-internals/glimmer/lib/utils/component-template.ts

@@ -0,0 +1,39 @@
+import { toString } from '@ember/-internals/utils';
+import { assert } from '@ember/debug';
+import { Option } from '@glimmer/interfaces';
+import { Factory as TemplateFactory } from '../template';
+
+const TEMPLATES: WeakMap<object, TemplateFactory> = new WeakMap();
+
+const getPrototypeOf = Object.getPrototypeOf;
+
+export function setComponentTemplate(factory: TemplateFactory, obj: object) {
+  assert(
+    `Cannot call \`setComponentTemplate\` on \`${toString(obj)}\``,
+    obj !== null && (typeof obj === 'object' || typeof obj === 'function')
+  );
+
+  assert(
+    `Cannot call \`setComponentTemplate\` multiple times on the same class (\`${obj}\`)`,
+    !TEMPLATES.has(obj)
+  );
+
+  TEMPLATES.set(obj, factory);
+
+  return obj;
+}
+
+export function getComponentTemplate(obj: object): Option<TemplateFactory> {
+  let pointer = obj;
+  while (pointer !== undefined && pointer !== null) {
+    let template = TEMPLATES.get(pointer);
+
+    if (template !== undefined) {
+      return template;
+    }
+
+    pointer = getPrototypeOf(pointer);
+  }
+
+  return null;
+}

packages/@ember/-internals/glimmer/lib/utils/managers.ts

@@ -1,7 +1,7 @@
 import { Owner } from '@ember/-internals/owner';
-import { Opaque, Option } from '@glimmer/interfaces';
+import { Option } from '@glimmer/interfaces';
 
-const MANAGERS: WeakMap<any, ManagerWrapper<Opaque>> = new WeakMap();
+const MANAGERS: WeakMap<object, ManagerWrapper<unknown>> = new WeakMap();
 
 const getPrototypeOf = Object.getPrototypeOf;
 
@@ -13,16 +13,18 @@ export interface ManagerWrapper<ManagerDelegate> {
   type: 'component' | 'modifier';
 }
 
-export function setManager<ManagerDelegate>(wrapper: ManagerWrapper<ManagerDelegate>, obj: any) {
+export function setManager<ManagerDelegate>(wrapper: ManagerWrapper<ManagerDelegate>, obj: object) {
   MANAGERS.set(obj, wrapper);
   return obj;
 }
 
-export function getManager<ManagerDelegate>(obj: any): Option<ManagerWrapper<ManagerDelegate>> {
+export function getManager<ManagerDelegate>(obj: object): Option<ManagerWrapper<ManagerDelegate>> {
   let pointer = obj;
   while (pointer !== undefined && pointer !== null) {
-    if (MANAGERS.has(pointer)) {
-      return MANAGERS.get(pointer) as ManagerWrapper<ManagerDelegate>;
+    let manager = MANAGERS.get(pointer);
+
+    if (manager !== undefined) {
+      return manager as ManagerWrapper<ManagerDelegate>;
     }
 
     pointer = getPrototypeOf(pointer);