rework state from parent container, add test

packages/state/src/container.ts

@@ -2,21 +2,24 @@ import { getOwner, type Owner } from 'solid-js';
 import { ExtractStore, GenericStoreApi, StoreApiDefinition } from './types';
 import { $CREATOR, resolve } from './api';
 import { runInSubRoot } from '~/root';
-import { getOptionalStateContext } from '~/solid/provider';
+import { StateBuilderError } from '~/error';
 
 export class Container {
   private readonly states = new Map<string, GenericStoreApi>();
 
-  protected constructor(private readonly owner: Owner) {}
+  protected constructor(
+    private readonly owner: Owner,
+    private readonly parent?: Container,
+  ) {}
 
-  static create(owner?: Owner) {
+  static create(owner?: Owner, parentContainer?: Container) {
     const resolvedOwner = owner ?? getOwner()!;
     if (!resolvedOwner) {
       console.warn(
         '[statebuilder] Using StateContainer without <StateProvider/> or `createRoot()` context is discouraged',
       );
     }
-    return new Container(resolvedOwner);
+    return new Container(resolvedOwner, parentContainer);
   }
 
   remove<TStoreDefinition extends StoreApiDefinition<any, any>>(
@@ -30,13 +33,13 @@ export class Container {
   ): ExtractStore<TStoreDefinition> {
     type TypedStore = ExtractStore<TStoreDefinition>;
     if (!state[$CREATOR]) {
-      throw new Error('[statebuilder] No state $CREATOR found.', {
+      throw new StateBuilderError('No state $CREATOR found.', {
         cause: { state: state },
       });
     }
     try {
       const name = state[$CREATOR].name;
-      const instance = this.recursivelySearchStateFromContainer(name);
+      const instance = this.#retrieveInstance(name);
       if (instance) {
         return instance as TypedStore;
       }
@@ -45,8 +48,8 @@ export class Container {
       return store as TypedStore;
     } catch (exception) {
       if (exception instanceof Error) throw exception;
-      throw new Error(
-        '[statebuilder] An error occurred during store initialization',
+      throw new StateBuilderError(
+        'An error occurred during store initialization',
         { cause: exception },
       );
     }
@@ -59,21 +62,19 @@ export class Container {
     return runInSubRoot(() => resolve(state, this), owner);
   }
 
-  private recursivelySearchStateFromContainer(
-    name: string,
-  ): GenericStoreApi | null {
-    let instance: GenericStoreApi | null;
-    instance = this.states.get(name) ?? null;
-    if (!instance && this.owner?.owner) {
-      const parentContainer = runInSubRoot((dispose) => {
-        const value = getOptionalStateContext();
-        dispose();
-        return value;
-      }, this.owner.owner);
-      if (parentContainer) {
-        instance = parentContainer.recursivelySearchStateFromContainer(name);
+  #retrieveInstance(name: string): GenericStoreApi | null {
+    let instance: GenericStoreApi | null = this.states.get(name) ?? null;
+    if (instance) {
+      return instance;
+    }
+    let currentParent = this.parent;
+    while (currentParent) {
+      instance = currentParent.states.get(name) ?? null;
+      if (!!instance) {
+        return instance;
       }
+      currentParent = currentParent.parent;
     }
-    return instance || null;
+    return instance;
   }
 }

packages/state/src/solid/provider.ts

@@ -18,7 +18,8 @@ export function StateProvider(props: FlowProps) {
       'Owner is missing. Cannot construct instance of Container',
     );
   }
-  const container = Container.create(owner);
+  const parentContainer = useContext(StateProviderContext);
+  const container = Container.create(owner, parentContainer);
   return createComponent(StateProviderContext.Provider, {
     value: container,
     get children() {
@@ -27,10 +28,6 @@ export function StateProvider(props: FlowProps) {
   });
 }
 
-export function getOptionalStateContext() {
-  return useContext(StateProviderContext) ?? null;
-}
-
 export function getStateContext() {
   const container = useContext(StateProviderContext);
   if (!container) {

packages/state/test/container.test.ts

@@ -1,22 +1,38 @@
-import { describe, expect, it } from 'vitest';
+import { assert, describe, it } from 'vitest';
 import { Container, defineStore } from '../src';
-import { getOwner } from 'solid-js';
+import { createRoot, getOwner } from 'solid-js';
 
 describe('Container', () => {
-  const owner = getOwner()!;
+  it('should create container', () =>
+    createRoot(() => {
+      const owner = getOwner()!;
+      const container = Container.create(owner);
+      assert.instanceOf(container, Container);
+    }));
 
-  it('should create container', function() {
-    const container = Container.create(owner);
-    expect(container).toBeInstanceOf(Container);
-  });
+  it('should create state', () =>
+    createRoot(() => {
+      const owner = getOwner()!;
+      const container = Container.create(owner);
+      const stateDef = defineStore(() => ({}));
 
-  it('should create state', function() {
-    const container = Container.create(owner);
-    const stateDef = defineStore(() => ({}));
+      const state = container.get(stateDef);
 
-    const state = container.get(stateDef);
+      assert.instanceOf(state, Function);
+      assert.ok(container['states'].size === 1);
+    }));
 
-    expect(state).toBeInstanceOf(Function);
-    expect(container['states'].size).toEqual(1);
+  it('should inject state from parent container', () => {
+    createRoot(() => {
+      const owner = getOwner()!;
+      const parentContainer = Container.create(owner);
+      const container = Container.create(owner, parentContainer);
+      const def = defineStore(() => ({}));
+      const stateFromParentContainer = parentContainer.get(def);
+      const stateFromContainer = container.get(def);
+      assert.strictEqual(stateFromContainer, stateFromParentContainer);
+      assert.isTrue(container['states'].size === 0);
+      assert.isTrue(parentContainer['states'].size === 1);
+    });
   });
 });