Register contexts at construct time (#35)

.changeset/chilly-cameras-greet.md

@@ -0,0 +1,5 @@
+---
+"ember-provide-consume-context": minor
+---
+
+Register contexts at construct time

ember-provide-consume-context/src/-private/@glimmer/opcodes.ts

@@ -7,6 +7,7 @@
 // This is safe, because the opcodes should be stable.
 // https://github.com/glimmerjs/glimmer-vm/blob/68d371bdccb41bc239b8f70d832e956ce6c349d8/packages/%40glimmer/vm/lib/opcodes.ts#L196
 export const enum Op {
+  CreateComponent = 87,
   GetComponentSelf = 90,
   DidRenderLayout = 100,
 }

ember-provide-consume-context/src/-private/override-glimmer-runtime-classes.ts

@@ -31,6 +31,13 @@ function overrideVM(runtime: any) {
     try {
       const { type, op1 } = opcode;
 
+      if (type === Op.CreateComponent) {
+        // Let the container know we're instantiating a new component
+        this.env.provideConsumeContextContainer?.createComponent();
+        // No need to register "updateWith", a component only instantiates
+        // once, and we don't need to run any further updates
+      }
+
       if (type === Op.GetComponentSelf) {
         // Get the component instance from the VM
         // (that's the VM's component instance, not the Glimmer Component one)

ember-provide-consume-context/src/-private/provide-consume-context-container.ts

@@ -58,6 +58,8 @@ export class ProvideConsumeContextContainer {
   // providing context values in tests.
   #globalContexts: Contexts | null = null;
 
+  #isCreatingComponent = false;
+
   begin(): void {
     this.reset();
   }
@@ -103,6 +105,11 @@ export class ProvideConsumeContextContainer {
   };
 
   enter(instance: ComponentInstance): void {
+    // When "enter" is called, a component instance has already been created.
+    // Update the flag to reflect that.
+    // See the "contextsFor" method below for how this flag is used.
+    this.#isCreatingComponent = false;
+
     const actualComponentInstance = (instance?.state as any)?.component;
 
     if (actualComponentInstance != null) {
@@ -128,25 +135,11 @@ export class ProvideConsumeContextContainer {
   }
 
   private registerProvider(provider: any) {
-    const { current } = this;
-
-    let providerContexts: Contexts = {};
-
-    // If global contexts are defined, make sure providers can read them
-    if (this.#globalContexts != null) {
-      providerContexts = { ...this.#globalContexts };
-    }
-
-    if (this.contexts.has(current)) {
-      // If a provider is nested within another provider, we merge their
-      // contexts
-      const context = this.contexts.get(current);
-      if (context != null) {
-        providerContexts = { ...providerContexts, ...context };
-      }
-    }
+    const providerContexts: Contexts = this.currentContexts();
 
     const registeredContexts = provider[EMBER_PROVIDE_CONSUME_CONTEXT_KEY];
+    // If the provider has registered contexts, store references
+    // to them on the current contexts object
     if (registeredContexts != null) {
       Object.entries(
         registeredContexts as Record<keyof ContextRegistry, string>,
@@ -164,15 +157,45 @@ export class ProvideConsumeContextContainer {
   }
 
   private registerComponent(component: any) {
+    const currentContexts = this.currentContexts();
+
+    // If a current context reference or global contexts exist, register them to the component
+    if (Object.keys(currentContexts).length > 0) {
+      this.contexts.set(component, currentContexts);
+    }
+  }
+
+  currentContexts() {
     const { current } = this;
 
     const globalContexts = this.#globalContexts ?? {};
 
-    // If a current context reference or global contexts exist, register them to the component
     if (this.contexts.has(current) || Object.keys(globalContexts).length > 0) {
       const context = this.contexts.get(current);
-      const mergedContexts = { ...globalContexts, ...context };
-      this.contexts.set(component, mergedContexts);
+      return { ...globalContexts, ...context };
+    }
+
+    return {};
+  }
+
+  contextsFor(component: any) {
+    if (this.contexts.has(component)) {
+      return this.contexts.get(component);
+    }
+
+    // If a context for this component is not yet registered, but
+    // we're in the phase of initializing a component, return
+    // the current contexts, so that the values can be read in constructors.
+    if (this.#isCreatingComponent) {
+      return this.currentContexts();
     }
+
+    return null;
+  }
+
+  createComponent() {
+    // Indicates that a component instance is being created, see
+    // "contextsFor" above for how we use this.
+    this.#isCreatingComponent = true;
   }
 }

ember-provide-consume-context/src/-private/utils.ts

@@ -29,7 +29,7 @@ export function getProvider(owner: any, contextKey: keyof ContextRegistry) {
     return null;
   }
 
-  const contextsObject = provideConsumeContextContainer.contexts.get(owner);
+  const contextsObject = provideConsumeContextContainer.contextsFor(owner);
   return contextsObject?.[contextKey];
 }