scoped stores contain their scope (#79)

Author: dmaskasky

src/ScopeProvider/ScopeProvider.tsx

@@ -1,27 +1,16 @@
-import {
-  type PropsWithChildren,
-  createContext,
-  useContext,
-  useEffect,
-  useState,
-} from 'react'
+import { type PropsWithChildren, useEffect, useState } from 'react'
 import { Provider, useStore } from 'jotai/react'
 import { useHydrateAtoms } from 'jotai/utils'
-import type {
-  AnyAtom,
-  AnyAtomFamily,
-  AtomDefault,
-  Scope,
-  Store,
+import {
+  type AnyAtom,
+  type AnyAtomFamily,
+  type AtomDefault,
+  SCOPE,
+  ScopedStore,
+  type Store,
 } from '../types'
-import { createPatchedStore, isTopLevelScope } from './patchedStore'
 import { createScope } from './scope'
 
-const ScopeContext = createContext<{
-  scope: Scope | undefined
-  baseStore: Store | undefined
-}>({ scope: undefined, baseStore: undefined })
-
 type ScopeProviderBaseProps = PropsWithChildren<{
   atoms?: Iterable<AnyAtom | AtomDefault>
   atomFamilies?: Iterable<AnyAtomFamily>
@@ -50,16 +39,9 @@ export function ScopeProvider({
   atoms: atomsOrTuples = [],
   atomFamilies,
   children,
-  debugName,
+  debugName: scopeName,
 }: ScopeProviderBaseProps) {
-  const parentStore: Store = useStore()
-  let { scope: parentScope, baseStore = parentStore } = useContext(ScopeContext)
-  // if this ScopeProvider is the first descendant scope under Provider then it is the top level scope
-  // https://github.com/jotaijs/jotai-scope/pull/33#discussion_r1604268003
-  if (isTopLevelScope(parentStore)) {
-    parentScope = undefined
-    baseStore = parentStore
-  }
+  const parentStore: Store | ScopedStore = useStore()
 
   const atoms = Array.from(atomsOrTuples, (a) => (Array.isArray(a) ? a[0] : a))
 
@@ -68,19 +50,20 @@ export function ScopeProvider({
   const atomFamilySet = new Set(atomFamilies)
 
   function initialize() {
-    const scope = createScope(atomSet, atomFamilySet, parentScope, debugName)
     return {
-      patchedStore: createPatchedStore(baseStore, scope),
-      scopeContext: { scope, baseStore },
+      scopedStore: createScope({
+        atomSet,
+        atomFamilySet,
+        parentStore,
+        scopeName,
+      }),
       hasChanged(current: {
-        baseStore: Store
-        parentScope: Scope | undefined
+        parentStore: Store | ScopedStore
         atomSet: Set<AnyAtom>
         atomFamilySet: Set<AnyAtomFamily>
       }) {
         return (
-          parentScope !== current.parentScope ||
-          baseStore !== current.baseStore ||
+          parentStore !== current.parentStore ||
           !isEqualSet(atomSet, current.atomSet) ||
           !isEqualSet(atomFamilySet, current.atomFamilySet)
         )
@@ -89,21 +72,17 @@ export function ScopeProvider({
   }
 
   const [state, setState] = useState(initialize)
-  const { hasChanged, scopeContext, patchedStore } = state
-  if (hasChanged({ parentScope, atomSet, atomFamilySet, baseStore })) {
-    scopeContext.scope?.cleanup()
+  const { hasChanged, scopedStore } = state
+  if (hasChanged({ atomSet, atomFamilySet, parentStore })) {
+    scopedStore[SCOPE].cleanup()
     setState(initialize)
   }
   useHydrateAtoms(
     Array.from(atomsOrTuples).filter(Array.isArray) as AtomDefault[],
-    { store: patchedStore }
-  )
-  useEffect(() => scopeContext.scope.cleanup, [scopeContext.scope])
-  return (
-    <ScopeContext.Provider value={scopeContext}>
-      <Provider store={patchedStore}>{children}</Provider>
-    </ScopeContext.Provider>
+    { store: scopedStore }
   )
+  useEffect(() => scopedStore[SCOPE].cleanup, [scopedStore])
+  return <Provider store={scopedStore}>{children}</Provider>
 }
 
 function isEqualSet(a: Set<unknown>, b: Set<unknown>) {

src/ScopeProvider/patchedStore.ts

@@ -1,21 +1,35 @@
 import { type Atom, atom } from 'jotai'
 import { __DEV__ } from '../env'
-import type { AnyAtom, AnyAtomFamily, AnyWritableAtom, Scope } from '../types'
+import {
+  type AnyAtom,
+  type AnyAtomFamily,
+  type AnyWritableAtom,
+  SCOPE,
+  type Scope,
+  type ScopedStore,
+  type Store,
+} from '../types'
 
 const globalScopeKey: { name?: string } = {}
 if (__DEV__) {
   globalScopeKey.name = 'unscoped'
-  globalScopeKey.toString = toString
+  globalScopeKey.toString = toNameString
 }
 
 type GlobalScopeKey = typeof globalScopeKey
 
-export function createScope(
-  atoms: Set<AnyAtom> = new Set(),
-  atomFamilies: Set<AnyAtomFamily> = new Set(),
-  parentScope: Scope | undefined,
-  scopeName?: string | undefined
-): Scope {
+export function createScope({
+  atomSet = new Set(),
+  atomFamilySet = new Set(),
+  parentStore,
+  scopeName,
+}: {
+  atomSet?: Set<AnyAtom>
+  atomFamilySet?: Set<AnyAtomFamily>
+  parentStore: Store | ScopedStore
+  scopeName?: string
+}): ScopedStore {
+  const parentScope = SCOPE in parentStore ? parentStore[SCOPE] : undefined
   const explicit = new WeakMap<AnyAtom, [AnyAtom, Scope?]>()
   const implicit = new WeakMap<AnyAtom, [AnyAtom, Scope?]>()
   type ScopeMap = WeakMap<AnyAtom, [AnyAtom, Scope?]>
@@ -53,16 +67,16 @@ export function createScope(
 
   if (scopeName && __DEV__) {
     currentScope.name = scopeName
-    currentScope.toString = toString
+    currentScope.toString = toNameString
   }
 
   // populate explicitly scoped atoms
-  for (const anAtom of atoms) {
+  for (const anAtom of atomSet) {
     explicit.set(anAtom, [cloneAtom(anAtom, currentScope), currentScope])
   }
 
   const cleanupFamiliesSet = new Set<() => void>()
-  for (const atomFamily of atomFamilies) {
+  for (const atomFamily of atomFamilySet) {
     for (const param of atomFamily.getParams()) {
       const anAtom = atomFamily(param)
       if (!explicit.has(anAtom)) {
@@ -72,7 +86,7 @@ export function createScope(
     const cleanupFamily = atomFamily.unstable_listen((e) => {
       if (e.type === 'CREATE' && !explicit.has(e.atom)) {
         explicit.set(e.atom, [cloneAtom(e.atom, currentScope), currentScope])
-      } else if (!atoms.has(e.atom)) {
+      } else if (!atomSet.has(e.atom)) {
         explicit.delete(e.atom)
       }
     })
@@ -245,7 +259,8 @@ export function createScope(
     }
   }
 
-  return currentScope
+  const scopedStore = createPatchedStore(parentStore, currentScope)
+  return scopedStore
 }
 
 function isWritableAtom(anAtom: AnyAtom): anAtom is AnyWritableAtom {
@@ -254,7 +269,7 @@ function isWritableAtom(anAtom: AnyAtom): anAtom is AnyWritableAtom {
 
 const { read: defaultRead, write: defaultWrite } = atom<unknown>(null)
 
-function toString(this: { name: string }) {
+function toNameString(this: { name: string }) {
   return this.name
 }
 
@@ -265,3 +280,46 @@ function combineVoidFunctions(...fns: (() => void)[]) {
     }
   }
 }
+
+function PatchedStore() {}
+
+/**
+ * @returns a patched store that intercepts get and set calls to apply the scope
+ */
+function createPatchedStore(baseStore: Store, scope: Scope): ScopedStore {
+  const store: ScopedStore = {
+    ...baseStore,
+    get(anAtom, ...args) {
+      const [scopedAtom] = scope.getAtom(anAtom)
+      return baseStore.get(scopedAtom, ...args)
+    },
+    set(anAtom, ...args) {
+      const [scopedAtom, implicitScope] = scope.getAtom(anAtom)
+      const restore = scope.prepareWriteAtom(
+        scopedAtom,
+        anAtom,
+        implicitScope,
+        scope
+      )
+      try {
+        return baseStore.set(scopedAtom, ...args)
+      } finally {
+        restore?.()
+      }
+    },
+    sub(anAtom, ...args) {
+      const [scopedAtom] = scope.getAtom(anAtom)
+      return baseStore.sub(scopedAtom, ...args)
+    },
+    [SCOPE]: scope,
+    // TODO: update this patch to support devtools
+  }
+  return Object.assign(Object.create(PatchedStore.prototype), store)
+}
+
+/**
+ * @returns true if the current scope is the first descendant scope under Provider
+ */
+export function isTopLevelScope(parentStore: Store) {
+  return !(parentStore instanceof PatchedStore)
+}

src/ScopeProvider/scope.ts

@@ -3,6 +3,10 @@ import type { AtomFamily } from 'jotai/vanilla/utils/atomFamily'
 
 export type Store = ReturnType<typeof createStore>
 
+export type ScopedStore = Store & {
+  [SCOPE]: Scope
+}
+
 export type AnyAtom = Atom<any> | AnyWritableAtom
 
 export type AnyAtomFamily = AtomFamily<any, AnyAtom>
@@ -45,4 +49,6 @@ export type Scope = {
   toString?: () => string
 }
 
+export const SCOPE = Symbol('scope')
+
 export type AtomDefault = readonly [AnyWritableAtom, unknown]

src/types.ts

@@ -8,7 +8,6 @@ import {
 } from 'jotai'
 import { describe, expect, test } from 'vitest'
 import { ScopeProvider } from '../../src'
-import { createPatchedStore } from '../../src/ScopeProvider/patchedStore'
 import { createScope } from '../../src/ScopeProvider/scope'
 import { AnyAtom } from '../../src/types'
 import { clickButton, getTextContents } from '../utils'
@@ -183,17 +182,12 @@ describe('scope chains', () => {
   c.debugLabel = 'c'
   function createScopes(atoms: AnyAtom[] = []) {
     const s0 = createStore()
-    const s1 = createPatchedStore(
-      s0,
-      createScope(
-        ...(Object.values({
-          atoms: new Set(atoms),
-          atomFamilies: undefined,
-          parentScope: undefined,
-          scopeName: 'S1',
-        }) as Parameters<typeof createScope>)
-      )
-    )
+    const s1 = createScope({
+      atomSet: new Set(atoms),
+      atomFamilySet: undefined,
+      parentStore: s0,
+      scopeName: 'S1',
+    })
     return { s0, s1 }
   }
   test('S1[a]: a1, b0(,a1), c0(,b0(,a1))', () => {