[WIP] Refactor to use the new jotai store exposed internal methods (Attempt 11)

src/ScopeProvider/ScopeProvider.tsx

@@ -1,14 +1,10 @@
 import { type PropsWithChildren, useEffect, useState } from 'react'
 import { Provider, useStore } from 'jotai/react'
 import { useHydrateAtoms } from 'jotai/utils'
-import {
-  type AnyAtom,
-  type AnyAtomFamily,
-  type AtomDefault,
-  SCOPE,
-  ScopedStore,
-  type Store,
-} from '../types'
+import type { INTERNAL_Store as Store } from 'jotai/vanilla/internals'
+import type { AnyAtom, AnyAtomFamily, AtomDefault, ScopedStore } from '../types'
+import { SCOPE } from '../types'
+import { isEqualSet } from '../utils'
 import { createScope } from './scope'
 
 type ScopeProviderBaseProps = PropsWithChildren<{
@@ -84,7 +80,3 @@ export function ScopeProvider({
   useEffect(() => scopedStore[SCOPE].cleanup, [scopedStore])
   return <Provider store={scopedStore}>{children}</Provider>
 }
-
-function isEqualSet(a: Set<unknown>, b: Set<unknown>) {
-  return a === b || (a.size === b.size && Array.from(a).every(b.has.bind(b)))
-}

src/ScopeProvider/scope.ts

@@ -1,15 +1,23 @@
 import { type Atom, atom } from 'jotai'
+import {
+  INTERNAL_Mounted,
+  INTERNAL_buildStoreRev1 as INTERNAL_buildStore,
+  INTERNAL_getBuildingBlocksRev1 as INTERNAL_getBuildingBlocks,
+  INTERNAL_isSelfAtom,
+  type INTERNAL_Store as Store,
+} from 'jotai/vanilla/internals'
 import { __DEV__ } from '../env'
 import type {
   AnyAtom,
   AnyAtomFamily,
   AnyWritableAtom,
+  BuildingBlocks,
+  CloneAtom,
   Scope,
   ScopedStore,
-  Store,
-  WithOriginal,
 } from '../types'
-import { SCOPE } from '../types'
+import { CONSUMER, EXPLICIT, SCOPE } from '../types'
+import { isCloneAtom, isEqualSet } from '../utils'
 
 const globalScopeKey: { name?: string } = {}
 if (__DEV__) {
@@ -73,7 +81,10 @@ export function createScope({
 
   // populate explicitly scoped atoms
   for (const anAtom of atomSet) {
-    explicit.set(anAtom, [cloneAtom(anAtom, currentScope), currentScope])
+    explicit.set(anAtom, [
+      cloneAtom(anAtom, currentScope, EXPLICIT),
+      currentScope,
+    ])
   }
 
   const cleanupFamiliesSet = new Set<() => void>()
@@ -183,13 +194,18 @@ export function createScope({
   /**
    * @returns a scoped copy of the atom
    */
-  function cloneAtom<T>(originalAtom: Atom<T>, implicitScope?: Scope) {
-    // avoid reading `init` to preserve lazy initialization
-    const scopedAtom: WithOriginal<Atom<T>> = Object.create(
+  function cloneAtom<T>(
+    originalAtom: Atom<T>,
+    implicitScope?: Scope,
+    cloneType?: EXPLICIT | CONSUMER
+  ) {
+    const scopedAtom: CloneAtom<Atom<T>> = Object.create(
+      // avoid reading `init` to preserve lazy initialization
       Object.getPrototypeOf(originalAtom),
       Object.getOwnPropertyDescriptors(originalAtom)
     )
-    scopedAtom.originalAtom = originalAtom
+    scopedAtom.o = originalAtom
+    scopedAtom.x = cloneType
 
     if (scopedAtom.read !== defaultRead) {
       scopedAtom.read = createScopedRead<typeof scopedAtom>(
@@ -263,6 +279,125 @@ export function createScope({
 
   const scopedStore = createPatchedStore(parentStore, currentScope)
   return scopedStore
+
+  /**
+   * @returns a patched store that intercepts get and set calls to apply the scope
+   */
+  function createPatchedStore(baseStore: Store, scope: Scope): ScopedStore {
+    const baseBuildingBlocks = INTERNAL_getBuildingBlocks(baseStore)
+    const [atomStateMap, mountedMap, invalidatedAtoms, changedAtoms] =
+      baseBuildingBlocks
+    const ensureAtomState = baseBuildingBlocks[11]
+    const readAtomState = baseBuildingBlocks[14]
+    const buildingBlocks: BuildingBlocks = [
+      atomStateMap,
+      undefined,
+      invalidatedAtoms,
+      changedAtoms,
+    ]
+    const internalMountedMap = new WeakMap<AnyAtom, INTERNAL_Mounted>()
+    buildingBlocks[1] = {
+      get: (atom) => {
+        if (!isCloneAtom(atom)) return mountedMap.get(atom)
+        if (!checkConsumer(atom)) return mountedMap.get(atom.o)
+        return internalMountedMap.get(atom)
+      },
+      set: (atom, mounted) => {
+        if (!isCloneAtom(atom)) return mountedMap.set(atom, mounted)
+        if (!checkConsumer(atom)) return mountedMap.set(atom.o, mounted)
+        return internalMountedMap.set(atom, mounted)
+      },
+      has: (atom) => {
+        if (!isCloneAtom(atom)) return mountedMap.has(atom)
+        if (!checkConsumer(atom)) return mountedMap.has(atom.o)
+        return internalMountedMap.has(atom)
+      },
+      delete: (atom) => {
+        if (!isCloneAtom(atom)) return mountedMap.delete(atom)
+        if (!checkConsumer(atom)) return mountedMap.delete(atom.o)
+        return internalMountedMap.delete(atom)
+      },
+    }
+    buildingBlocks[14] = (atom) => {
+      checkConsumer(atom)
+      const deps = new Set(ensureAtomState(atom).d.keys())
+      if (isCloneAtom(atom) && atom.x === undefined) {
+        const newAtomState = readAtomState(atom.o)
+        // deps changed?
+        const newDeps = new Set(newAtomState.d.keys())
+        if (!isEqualSet(deps, newDeps)) {
+          checkConsumer(atom)
+        }
+        return newAtomState
+      }
+      return readAtomState(atom)
+    }
+    const wrappedBaseStore = INTERNAL_buildStore(...buildingBlocks)
+    const storeShim: ScopedStore = {
+      get(anAtom, ...args) {
+        const [scopedAtom] = scope.getAtom(anAtom)
+        return wrappedBaseStore.get(scopedAtom, ...args)
+      },
+      set(anAtom, ...args) {
+        const [scopedAtom, implicitScope] = scope.getAtom(anAtom)
+        const restore = scope.prepareWriteAtom(
+          scopedAtom,
+          anAtom,
+          implicitScope,
+          scope
+        )
+        try {
+          return wrappedBaseStore.set(scopedAtom, ...args)
+        } finally {
+          restore?.()
+        }
+      },
+      sub(anAtom, ...args) {
+        const [scopedAtom] = scope.getAtom(anAtom)
+        return wrappedBaseStore.sub(scopedAtom, ...args)
+      },
+      [SCOPE]: scope,
+    }
+    return Object.assign(wrappedBaseStore, storeShim) as ScopedStore
+
+    /**
+     * Check if the atom is a consumer.
+     * Looks at the atom's dependencies to determine if it is a consumer.
+     * Updates the atom's clone type with the new value if it changed.
+     * Recursively checks the dependents if mounted.
+     * @param atom
+     * @returns true if the atom is a consumer
+     */
+    function checkConsumer(atom: AnyAtom): boolean {
+      let atomState = ensureAtomState(atom)
+      const mountedState = mountedMap.get(atom)
+      if (!isCloneAtom(atom) || atom.x === EXPLICIT) {
+        return false
+      }
+
+      if (!mountedState && mountedMap.has(atom.o)) {
+        atomState = ensureAtomState(atom.o)
+      }
+
+      const dependencies = Array.from(atomState.d.keys()).filter(
+        (a) => !INTERNAL_isSelfAtom(atom, a)
+      )
+
+      const isConsumer = dependencies.some(
+        (atom) =>
+          (isCloneAtom(atom) && (atom.x === CONSUMER || atom.x === EXPLICIT)) ||
+          explicit.has(atom) // TODO: a consumer can also read consumers and inherited too.
+      )
+      if (atom.x === CONSUMER || atom.x === undefined) {
+        const newValue = isConsumer ? CONSUMER : undefined
+        if (atom.x !== newValue) {
+          atom.x = newValue
+          mountedState?.t.forEach(checkConsumer)
+        }
+      }
+      return isConsumer
+    }
+  }
 }
 
 function isWritableAtom(anAtom: AnyAtom): anAtom is AnyWritableAtom {
@@ -282,46 +417,3 @@ 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/createIsolation.tsx

@@ -8,7 +8,8 @@ import {
 } from 'jotai/react'
 import { useHydrateAtoms } from 'jotai/react/utils'
 import { createStore } from 'jotai/vanilla'
-import type { AnyWritableAtom, Store } from './types'
+import { INTERNAL_Store as Store } from 'jotai/vanilla/internals'
+import type { AnyWritableAtom } from './types'
 
 type CreateIsolationResult = {
   Provider: (props: {

src/types.ts

@@ -1,11 +1,11 @@
-import type { Atom, WritableAtom, createStore } from 'jotai/vanilla'
+import type { Atom, WritableAtom } from 'jotai/vanilla'
+import {
+  INTERNAL_getBuildingBlocksRev1 as INTERNAL_getBuildingBlocks,
+  INTERNAL_Store as Store,
+} from 'jotai/vanilla/internals'
 import type { AtomFamily } from 'jotai/vanilla/utils/atomFamily'
 
-export type Store = ReturnType<typeof createStore>
-
-export type ScopedStore = Store & {
-  [SCOPE]: Scope
-}
+export type ScopedStore = Store & { [SCOPE]: Scope }
 
 export type AnyAtom = Atom<any> | AnyWritableAtom
 
@@ -53,6 +53,21 @@ export const SCOPE = Symbol('scope')
 
 export type AtomDefault = readonly [AnyWritableAtom, unknown]
 
-export type WithOriginal<T extends AnyAtom> = T & {
-  originalAtom: T
+export type CloneAtom<T extends AnyAtom> = T & {
+  /** original atom */
+  o: T
+  /** clone type */
+  x: EXPLICIT | CONSUMER | undefined
+}
+
+export const EXPLICIT = Symbol('explicit')
+export const CONSUMER = Symbol('consumer')
+export type EXPLICIT = typeof EXPLICIT
+export type CONSUMER = typeof CONSUMER
+
+type Mutable<T> = {
+  -readonly [K in keyof T]: T[K]
 }
+export type BuildingBlocks = Partial<
+  Mutable<ReturnType<typeof INTERNAL_getBuildingBlocks>>
+>

src/utils.ts

@@ -0,0 +1,9 @@
+import { AnyAtom, CloneAtom } from './types'
+
+export function isCloneAtom<T extends AnyAtom>(atom: T): atom is CloneAtom<T> {
+  return 'o' in atom && 'x' in atom
+}
+
+export function isEqualSet(a: Set<unknown>, b: Set<unknown>) {
+  return a === b || (a.size === b.size && Array.from(a).every(b.has.bind(b)))
+}