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

[dmaskasky]

Goals

1. Don't copy scoped atoms.
Currently, we copy scoped atoms because all atoms (scoped and unscoped) use the same base store. When scoped atoms are copied, their init property is accessed an extra time (per copy). This causes issues with atoms such as atomWithLazy. The new approach will be to use separate stores for each scope layer and to create a routing function that interprets the atom correctly.

2. Don't copy unscoped derived atoms.
Currently, we copy all derived atoms. This allows us shim the getter so derived atoms can read scoped atoms. The issue with doing this is that it causes another mount event for the copy. This causes issues with atoms such as atomWithObservable. The plan is to use internal store methods to resolve the correct atom / atomState rather than shimming the getter. We can also create our own readAtomState, but that might be too heavy-handed.

3. Remove the writable hack.
Currently, writable atoms may read or write to scoped atoms, so we need to shim their write function's getter and setter. However, we cannot copy writable atoms because they might also hold a meaningful value. So instead we mutate them temporarily. I'd like to avoid this.

Summary of changes

Update ScopeProvider to use the new exposed internals api (pmndrs/jotai#2911).

The new code should address #25 and #36.

TODO

1. figure out how unscoped derived can read scoped atoms
2. simplify tests

Prior art

#56
#33

---

Overview of Expected Behavior

atoms: a, b, c, d, A(c + d), B(a + b + A)

S1[c]: a0, b0, c1, d0, A0(c1 + d0), B0(a0 + b0 + A0(c1 + d0))
S2[B]: a0, b0, c1, d0, A0(c1 + d0), B2(a2 + b2 + A2(c2 + d2))
S3[a]: a3, b0, c1, d0, A0(c1 + d0), B2(a3 + b2 + A2(c2 + d2))
S4[A]: a3, b0, c1, d0, A4(c4 + d4), B2(a3 + b2 + A4(c4 + d4))
S5[d]: a3, b0, c1, d5, A4(c4 + d5), B2(a3 + b2 + A4(c4 + d5))
-------------------------------------------------------------
S1:
  explicit: [c1]
  implicit: []
  inherited: [0 => [a0, b0, d0, B0, A0]]
S2:
  explicit: [B2]
  implicit: [a2, b2, c2, d2, A2],
  inherited: [0 => [a0, b0, d0, A0], 1 => [c1]]
S3:
  explicit: [a3]
  implicit: []
  inherited: [0 => [b0, d0, A0], 1 => [c1], 2 => [b2, c2, d2, A2, B2]]
S4:
  explicit: [A4]
  implicit: [c4, d4]
  inherited: [0 => [b0, d0], 1 => [c1], 2 => [b2, c2, d2, B2], 3 => [a3]]
S5:
  explicit: [d5]
  implicit: []
  inherited: [0 => [b0], 1 => [c1], 2 => [b2, B2], 3 => [a3],  4 => [c4, A4]]

Let us assume we have base atoms a, b, c, d and derived atoms A and B; { where A depends on c and d, and B depends on a, b, and A }

Scope S0 is the global scope under the nearest jotai Provider or defaultStore.
Scopes S1-Sn are nested where S1 is the first level and S5 is a descendant of S1 at the fifth level.

In scope S1, primitive atom c is explicitly scoped. This means that c1 corresponding to c in S1 holds an independent value. All derived atoms in this scope will use c1. All descendant scopes will use c1 unless c is explicitly defined in a nearer ancestor.

In scope S2, derived atom B is explicitly scoped. This means that all of its dependents are implicitly scoped. Explicitly and implicitly scoped atoms are copied and do not inherit their value from the parent scope. This is why all dependents of B2 are denoted with a 2 suffix.

In scope S3, primitive atom a is explicitly scoped. All other atoms are inherited from S2 where B is explicitly scoped. Even though B2 is inherited, it still will use a3 in S3. This is a similar behavior to B0 using c1 in S1 with the only difference being B0 is unscoped vs B2 is inherited. To me, there is no difference between inherited and unscoped.

And so on... All subsequent scopes feature similar behavior as described above.

FAQ

Types of atoms

// primitive
// { read: defaultRead, write: defaultWrite }
atom(0)

// derived
// { read: customRead }
atom(function customRead(get) {})

// writable
// { read: defaultRead, write: customWrite }
atom(0, function customWrite(get, set, ...args) {})

// derived writable
// { read: customRead, write: customWrite }
atom(
  function customRead(get) {},
  function customWrite(get, set, ...args) {}
)

atoms type A & B can hold a value. Type B can also reference atoms in the current scope. When a scope inherits a valued atom (A or B), the inherited value is used.

For Type B, we currently:

1. use the original atom config to keep using the same value
2. override the write function

I'm generally not a fan of mutating configs but I don't see any other way, unless Jotai natively supports creating a cloned copy of an atom whose config references the original atom's value in the store.

Scoping Priority

1. Explicitly scoped - An atom is explicitly scoped (passed in to ScopeProvider) in the current scope.

// a is explicitly scoped in S1
a
S1[a]: a1

2. Implicitly scoped - An atom is being called by an Explicitly or Implicitly scoped derived atom.

// a in B is implicitly scoped in S1
a, B(a)
S1[B]: a0, B1(a1)

3. Inherited scoped - An atom is explicitly scoped in an ancestor ScopeProvider. The nearest scope that explicitly scopes the atom takes priority. Atom dependencies of inherited scoped atoms are implicitly scoped to the ancestor scope

// a is inherited scoped in S2
a
S1[a]: a1
S2[ ]: a1

4. Unscoped - An atom is not scoped, but it may still read scoped atoms.

// B is unscoped scoped in S1
a, B
S1[a]: a1, B0(a1)

---

Existing Implementation

Scope

createScope handles everything dealing with scope and is moved to a separate file.

WeakMaps

1. explicit - atoms added to ScopeProvider
2. implicit - dependencies of explicit and implicit atoms are scoped
3. unscoped / inherited

• scoped atoms of the parent scope (explicit, implicit, inherited)
• derived readable and derived writeable atoms, needed for enabling access to scoped atoms

getAtom

from the originalAtom, gets the explicit, implicit, inherited, unscoped derived, and unscoped primitive atoms. As necessary, creates scoped copies (one-time) and saves them in their appropriate weakmap.

Inheriting Atoms

Atoms are inherited from the ancestor scopes. The top ancestor is the global scope where the original atom is used.

Inherited Derived Atoms

To inherit atoms that have a custom read, they are copied so that the custom read and write functions can use that scope's getAtom function.

Inheriting Primitive Atoms

To inherit primitive atoms, those atoms are stored directly and not copied. Since the atom itself is the key to its value, we need to store the original in the weakmap.

Inherited Write Only Atoms

const valuedWritableAtom = atom(0, (get, set) => {
  set(valuedWritableAtom, get(someAtom))
})

Where someAtom could be scoped.

⚠️ Dirty Hack
For write-only atoms, since these atoms also hold a value, we also want to preserve the originals. However the write function may access atoms in the current scope. To work around this, we modify the original write method synchronous before baseStore.set and restore the original write method synchronous after in a finally block.

What I'd love is a utility to clone an atom but have the clone also resolve the same value as the original. That way I can copy atoms like below and wrap their write method to resolve atoms in the current scope.

---

Summary of Changes

TBC

Tests Results

01_basic_spec

❌ 01. ScopeProvider does not provide isolation for unscoped primitive atoms

base
S0[]: base0
S1[]: base0

❌ 02. unscoped derived atoms are unaffected in ScopeProvider

base, Derived(base)
S0[]: base0 Derived0(base0)
S1[]: base0 Derived0(base0)

❌ 03. ScopeProvider provides isolation for scoped primitive atoms

base
S0[base]: base0
S1[base]: base1

❌ 04. unscoped derived can read and write to scoped primitive atoms

base, derived(base)
S0[base]: derived0(base0)
S1[base]: derived0(base1)

❌ 05. unscoped derived can read both scoped and unscoped atoms

base, notScoped, derived(base + notScoped)
S0[base]: derived0(base0 + notScoped0)
S1[base]: derived0(base1 + notScoped0)

❌ 06. dependencies of scoped derived are implicitly scoped

base, derived(base),
S0[derived]: derived0(base0)
S1[derived]: derived1(base1)

❌ 07. scoped derived atoms can share implicitly scoped dependencies

base, derivedA(base), derivemB(base)
S0[derivedA, derivedB]: derivedA0(base0), derivedB0(base0)
S1[derivedA, derivedB]: derivedA1(base1), derivedB1(base1)

❌ 08. nested scopes provide isolation for primitive atoms at every level

base, derivedA(base), derivedB(base)
S0[base]: base0
S1[base]: base1
S2[base]: base2
S3[base]: base3

❌ 09. unscoped derived atoms in nested scoped can read and write to scoped primitive atoms at every level

baseA, baseB, baseC, derived(baseA + baseB + baseC),
S0[     ]: derived(baseA0 + baseB0 + baseC0)
S1[baseB]: derived(baseA0 + baseB1 + baseC0)
S2[baseC]: derived(baseA0 + baseB1 + baseC2)

❌ 10. inherited scoped derived atoms can read and write to scoped primitive atoms at every nested level

baseA, baseB, derived(baseA + baseB)
S1[baseB, derived]: derived1(baseA1 + baseB1)
S2[baseB]: derived1(baseA1 + baseB2)

02_removeScope

❌ atom get correct value when ScopeProvider is added/removed

03_nested

❌ nested primitive atoms are correctly scoped

baseA, baseB, baseC
S1[baseA]: baseA1 baseB0 baseC0
S2[baseB]: baseA1 baseB2 baseC0

04_derived

❌ parent scope's derived atom is prior to nested scope's scoped base

base, derivedA(base), derivedB(base)
case1[base]: base1, derivedA0(base1), derivedB0(base1)
case2[derivedA, derivedB]: base0, derivedA1(base1), derivedB1(base1)
layer1[derivedA]: base0, derivedA1(base1), derivedB0(base0)
layer2[base, derivedB]: base2, derivedA1(base2), derivedB2(base2)

05_derived_self

❌ derived dep scope is preserved in self reference

baseA, derivedB(baseA, derivedB)
S1[baseA]: baseA1, derivedB0(baseA1, derivedB0)

06_implicit_parent

❌ level 1: "BD" and level 2: "BD"
❌ level 1: "BD" and level 2: "DB"
❌ level 1: "DB" and level 2: "BD"
❌ level 1: "DB" and level 2: "DB"

07_writable

❌ "writableAtom" updates its value in both scoped and unscoped and read scoped atom (27 ms)
❌ "thisWritableAtom" updates its value in both scoped and unscoped and read scoped atom (5 ms)

writable=w(,w + s), base=b
S0[ ]: b0, w0(,w0 + b0)
S1[b]: b1, w0(,w0 + b1)