From 5b80856eda59abfb81d3daf3be2d5d07b36e1b9a Mon Sep 17 00:00:00 2001
From: Riccardo Perra <perrariccardo0@gmail.com>
Date: Sat, 28 Oct 2023 11:22:57 +0200
Subject: [PATCH] allow to inject states from parent context

---
 examples/counter/src/components/Counter.tsx | 33 ++++++++++-
 packages/state/src/container.ts             | 63 ++++++++++-----------
 packages/state/src/index.ts                 |  7 ++-
 packages/state/src/root.ts                  | 36 ++++++++++++
 packages/state/src/solid/index.ts           | 10 +++-
 packages/state/src/solid/provider.ts        | 31 +++++-----
 packages/state/src/solid/resource.ts        | 23 +++++---
 packages/state/test/resource.test.ts        | 11 ++--
 8 files changed, 147 insertions(+), 67 deletions(-)
 create mode 100644 packages/state/src/root.ts

diff --git a/examples/counter/src/components/Counter.tsx b/examples/counter/src/components/Counter.tsx
index 98bc8fc..5cc7e45 100644
--- a/examples/counter/src/components/Counter.tsx
+++ b/examples/counter/src/components/Counter.tsx
@@ -1,5 +1,10 @@
 import './Counter.css';
-import { defineStore, InjectFlags, provideState } from 'statebuilder';
+import {
+  defineSignal,
+  defineStore,
+  provideState,
+  StateProvider,
+} from 'statebuilder';
 import { withProxyCommands } from 'statebuilder/commands';
 import { withReduxDevtools } from 'statebuilder/devtools';
 import { withAsyncAction } from 'statebuilder/asyncAction';
@@ -26,6 +31,11 @@ function appReducer(state: AppState, action: AppActions) {
   }
 }
 
+const GlobalCount = defineSignal(() => 1).extend((_, context) => {
+  context.hooks.onInit(() => console.log('init count2'));
+  context.hooks.onDestroy(() => console.log('destroy count2'));
+});
+
 const CountStore = defineStore(() => ({
   count: 0,
 }))
@@ -54,11 +64,19 @@ const CountStore = defineStore(() => ({
     };
   });
 
-export default function Counter() {
-  const store = provideState(CountStore, InjectFlags.global);
+function Counter() {
+  const store = provideState(CountStore);
+  const globalCount = provideState(GlobalCount);
 
   return (
     <>
+      <button
+        class="increment"
+        onClick={() => globalCount.set((count) => count + 1)}
+      >
+        Global Count Clicks: {globalCount()}
+      </button>
+
       <button
         class="increment"
         onClick={() => store.actions.increment()}
@@ -83,3 +101,12 @@ export default function Counter() {
     </>
   );
 }
+
+export default function CounterRoot() {
+  const globalCount = provideState(GlobalCount);
+  return (
+    <StateProvider>
+      <Counter />
+    </StateProvider>
+  );
+}
diff --git a/packages/state/src/container.ts b/packages/state/src/container.ts
index e0ad59d..92c9730 100644
--- a/packages/state/src/container.ts
+++ b/packages/state/src/container.ts
@@ -1,18 +1,15 @@
-import { getOwner, type Owner, runWithOwner } from 'solid-js';
+import { getOwner, type Owner } from 'solid-js';
 import { ExtractStore, GenericStoreApi, StoreApiDefinition } from './types';
 import { $CREATOR, resolve } from './api';
-
-export const enum InjectFlags {
-  global,
-  local,
-}
+import { runInSubRoot } from '~/root';
+import { getOptionalStateContext } from '~/solid/provider';
 
 export class Container {
   private readonly states = new Map<string, GenericStoreApi>();
 
-  protected constructor(private readonly owner: Owner) {}
+  protected constructor(private readonly owner: typeof Owner) {}
 
-  static create(owner?: Owner) {
+  static create(owner?: typeof Owner) {
     const resolvedOwner = owner ?? getOwner()!;
     if (!resolvedOwner) {
       console.warn(
@@ -30,18 +27,20 @@ export class Container {
 
   get<TStoreDefinition extends StoreApiDefinition<any, any>>(
     state: TStoreDefinition,
-    flags?: InjectFlags,
   ): ExtractStore<TStoreDefinition> {
     type TypedStore = ExtractStore<TStoreDefinition>;
-
+    if (!state[$CREATOR]) {
+      throw new Error('[statebuilder] No state $CREATOR found.', {
+        cause: { state: state },
+      });
+    }
     try {
       const name = state[$CREATOR].name;
-      const instance = this.states.get(name);
+      const instance = this.recursivelySearchStateFromContainer(name);
       if (instance) {
-        return instance as unknown as TypedStore;
+        return instance as TypedStore;
       }
-      const owner = this.#resolveOwner(flags ?? InjectFlags.global);
-      const store = this.#resolveStore(owner!, state);
+      const store = this.#resolveStore(this.owner, state);
       this.states.set(name, store!);
       return store as TypedStore;
     } catch (exception) {
@@ -54,27 +53,27 @@ export class Container {
   }
 
   #resolveStore<TStoreDefinition extends StoreApiDefinition<any, any>>(
-    owner: Owner,
+    owner: typeof Owner,
     state: TStoreDefinition,
-  ) {
-    let error: Error | undefined;
-    const store = runWithOwner(owner, () => {
-      try {
-        return resolve(state, this);
-      } catch (e) {
-        error = e as Error;
-      }
-    });
-    if (error) throw error;
-    return store;
+  ): GenericStoreApi {
+    return runInSubRoot(() => resolve(state, this), owner);
   }
 
-  #resolveOwner(flags: InjectFlags) {
-    switch (flags) {
-      case InjectFlags.global:
-        return this.owner;
-      case InjectFlags.local:
-        return getOwner();
+  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);
+      }
     }
+    return instance || null;
   }
 }
diff --git a/packages/state/src/index.ts b/packages/state/src/index.ts
index 3a8832c..188ab82 100644
--- a/packages/state/src/index.ts
+++ b/packages/state/src/index.ts
@@ -1,13 +1,16 @@
-export { Container, InjectFlags } from './container';
+export { Container } from './container';
 
 export {
   StateProvider,
   provideState,
+  getStateContext,
   defineSignal,
   defineStore,
+  ɵdefineResource,
+  ɵWithResourceStorage,
 } from './solid';
 
-export type { Signal, Store } from './solid';
+export type { Signal, Store, Resource } from './solid';
 
 export { create, makePlugin } from './api';
 
diff --git a/packages/state/src/root.ts b/packages/state/src/root.ts
new file mode 100644
index 0000000..c1c444e
--- /dev/null
+++ b/packages/state/src/root.ts
@@ -0,0 +1,36 @@
+import {
+  createRoot,
+  getOwner,
+  onCleanup,
+  type Owner,
+  runWithOwner,
+} from 'solid-js';
+
+type Disposable<T> = (dispoe: () => void) => T;
+
+function createSubRoot<T>(
+  fn: Disposable<T>,
+  owner: typeof Owner = getOwner(),
+): T {
+  return createRoot((dispose) => {
+    owner && runWithOwner(owner, onCleanup.bind(void 0, dispose));
+    return fn(dispose);
+  }, owner!);
+}
+
+export function runInSubRoot<T>(fn: Disposable<T>, owner?: typeof Owner): T {
+  let error: unknown;
+  const result = createSubRoot((dispose) => {
+    try {
+      return fn(dispose);
+    } catch (e) {
+      error = e;
+      dispose();
+      throw e;
+    }
+  }, owner);
+  if (error) {
+    throw error;
+  }
+  return result;
+}
diff --git a/packages/state/src/solid/index.ts b/packages/state/src/solid/index.ts
index ca21089..8328b67 100644
--- a/packages/state/src/solid/index.ts
+++ b/packages/state/src/solid/index.ts
@@ -1,8 +1,16 @@
-export { provideState, StateProvider } from './provider';
+export { provideState, StateProvider, getStateContext } from './provider';
 
 export { defineStore, type Store, type StoreValue } from './store';
+
 export {
   defineSignal,
   type Signal,
   type SignalDefinitionCreator,
 } from './signal';
+
+export {
+  ɵdefineResource,
+  ɵWithResourceStorage,
+  type ResourceActions,
+  type Resource,
+} from './resource';
diff --git a/packages/state/src/solid/provider.ts b/packages/state/src/solid/provider.ts
index 23638a0..5b86625 100644
--- a/packages/state/src/solid/provider.ts
+++ b/packages/state/src/solid/provider.ts
@@ -5,15 +5,17 @@ import {
   getOwner,
   useContext,
 } from 'solid-js';
-import { Container, InjectFlags } from '~/container';
-import { StoreApiDefinition, ExtractStore } from '~/types';
+import { Container } from '~/container';
+import { ExtractStore, StoreApiDefinition } from '~/types';
 
 const StateProviderContext = createContext<Container>();
 
 export function StateProvider(props: FlowProps) {
   const owner = getOwner();
   if (!owner) {
-    throw new Error('Owner missing');
+    throw new Error(
+      '[statebuilder] Owner is missing. Cannot construct instance of Container',
+    );
   }
   const container = Container.create(owner);
   return createComponent(StateProviderContext.Provider, {
@@ -24,20 +26,21 @@ export function StateProvider(props: FlowProps) {
   });
 }
 
-function useStateContext() {
-  const ctx = useContext(StateProviderContext);
-  if (ctx) {
-    return ctx;
+export function getOptionalStateContext() {
+  return useContext(StateProviderContext) ?? null;
+}
+
+export function getStateContext() {
+  const container = useContext(StateProviderContext);
+  if (!container) {
+    throw new Error('No <StateProvider> found in component tree');
   }
-  throw new Error('No <StateProvider> found in component tree');
+  return container;
 }
 
 export function provideState<
   TStoreDefinition extends StoreApiDefinition<any, any>,
->(
-  definition: TStoreDefinition,
-  flags?: InjectFlags,
-): ExtractStore<TStoreDefinition> {
-  const context = useStateContext();
-  return context.get(definition, flags);
+>(definition: TStoreDefinition): ExtractStore<TStoreDefinition> {
+  const context = getStateContext();
+  return context.get(definition);
 }
diff --git a/packages/state/src/solid/resource.ts b/packages/state/src/solid/resource.ts
index 8fa442c..f8a7a3d 100644
--- a/packages/state/src/solid/resource.ts
+++ b/packages/state/src/solid/resource.ts
@@ -1,15 +1,24 @@
 import {
+  createResource,
   NoInfer,
+  Resource as InternalResource,
+  ResourceActions as InternalResourceActions,
   ResourceFetcher,
   ResourceOptions,
   Setter,
   Signal,
-  createResource,
-  Resource as InternalResource,
 } from 'solid-js';
-import { GenericStoreApi, create } from '..';
+import { GenericStoreApi } from '~/types';
+import { create } from '~/api';
+
+export interface ResourceActions<T> {
+  set: InternalResourceActions<T>['mutate'];
+  refetch: InternalResourceActions<T>['refetch'];
+}
 
-export type Resource<T> = GenericStoreApi<T, Setter<T>> & InternalResource<T>;
+export type Resource<T> = GenericStoreApi<T, Setter<T>> &
+  InternalResource<T> &
+  ResourceActions<T>;
 
 function makeResource<T>(
   resourceFetcher: ResourceFetcher<true, T, true>,
@@ -26,11 +35,9 @@ function makeResource<T>(
   return state as unknown as Resource<T>;
 }
 
-export const experimental__defineResource = create('resource', makeResource);
+export const ɵdefineResource = create('resource', makeResource);
 
-export function experimental__withResourceStorage<T>(
-  store: GenericStoreApi<T>,
-) {
+export function ɵWithResourceStorage<T>(store: GenericStoreApi<T>) {
   return function (_: T | undefined): Signal<T | undefined> {
     return [
       () => store(),
diff --git a/packages/state/test/resource.test.ts b/packages/state/test/resource.test.ts
index 5912afc..3c32ea3 100644
--- a/packages/state/test/resource.test.ts
+++ b/packages/state/test/resource.test.ts
@@ -1,10 +1,7 @@
 import { describe, expect, it } from 'vitest';
 import { $CREATOR } from '~/api';
 import { Container } from '~/container';
-import {
-  experimental__defineResource,
-  experimental__withResourceStorage,
-} from '~/solid/resource';
+import { ɵdefineResource, ɵWithResourceStorage } from '~/solid/resource';
 import { createResource, createRoot } from 'solid-js';
 import { defineSignal } from '~/solid';
 
@@ -16,7 +13,7 @@ interface Todo {
 
 describe('Resource', () => {
   it('should define state with params', () => {
-    const def = experimental__defineResource(
+    const def = ɵdefineResource(
       () => async () =>
         Promise.resolve({
           id: 1,
@@ -29,7 +26,7 @@ describe('Resource', () => {
   });
 
   it('should fetch data async', async () => {
-    const def = experimental__defineResource(async () => {
+    const def = ɵdefineResource(async () => {
       await new Promise((r) => setTimeout(r, 2000));
       return Promise.resolve({
         id: 1,
@@ -60,7 +57,7 @@ describe('Resource', () => {
       const [stateResource, { mutate }] = createResource(
         () => Promise.resolve(5),
         {
-          storage: experimental__withResourceStorage(store),
+          storage: ɵWithResourceStorage(store),
         },
       );
       return {