lazy recomputeDependents

[dmaskasky]

Related Bug Reports or Discussions

#2782

Summary

Currently, recomputeDependents runs for each call to set. This causes overlapping dependents to recompute multiple times.

We still need to recompute some dependents if they are accessed with get. But the remaining dependents can be recomputed at the end of the transaction in flushPending.

const a = atom(0)
const b = atom((get) => get(a))
const fetch = vi.fn()
const c = atom((get) => fetch(get(a)))
const w = atom(null, (get, set) => {
  set(a, 1)
  expect(get(b)).toBe(1)
  expect(fetch).toHaveBeenCalledTimes(0)
})
const store = createStore()
store.sub(b, () => {})
store.sub(c, () => {})
fetch.mockClear()
store.set(w)
expect(fetch).toHaveBeenCalledOnce()
expect(fetch).toBeCalledWith(1)
expect(store.get(a)).toBe(1)

Proposed Fix

1. on set
   • keep a list of all atoms that have been changed
   • mark the atom's state and its dependents' state as dirty
2. on get, if the atom's state is dirty
   • recompute the atom and all it's dependencies that are dirty
3. on flushPending, recompute all remaining dirty dependents

More complex example

Suppose you write to node 17. Itself and it's dependents are {17, 15, 16, 12, 13, 14, 9, 10, 11, 6, 3, 1}.
We save these nodes as pending recompute.
Then suppose you read from node 1. Itself and it's dependencies are {1, 2, 3, 4, 5, 6, 7, 8, 9, 12, 15, 17}.
The intersection of these two sets is {1, 3, 6, 9, 12, 15, 17}.
On get, only these nodes will be eagerly recomputed.
On flushPending, the remaining nodes {16, 13, 14, 10, 11} will be recomputed.

Description of Core Functions

getMountedDependents

Was previously getDependents. Is called by getAllDependents. Returns a set of all mounted dependents.

getAllDependents

Is called by markRecomputePending and getSortedDependents. Returns a map of atoms to a set of their dependents (Map<AnyAtom, Set<AnyAtom>>).

markRecomputePending

Is called by recomputeDependents and setter. Sets atomState.x to true.

getSortedDependents

Is called by recomputeDependents. Returns a topological sorted list of atoms and its dependent atoms (deep).

recomputeDependents

Is called by recomputeDependencies and flushPending. Receives a set of atoms for which to recompute dependents on. Iteratively calls readAtomState and mountDependencies on all atoms and their dependents (deep) whose atomState.x is true.

recomputeDependencies

Is called by the getter in writeAtomState. Recomputes the dependents (deep) of all dependencies (deep) whose atomState.x is true.

flushPending

Is called in the following places

1. In the readAtomState getter finally block if getter is called async.
2. In complete which is called when a valueOrPromise is a promise and has settled.
3. In setter finally block if setter is called async.
4. In writeAtom finally block.
5. In setAtom finally block if setAtom is called async.
6. At the end of subscribeAtom
7. At the end of unsubscribeAtom

Check List

• pnpm run prettier for formatting code and docs

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Comments	Updated (UTC)
jotai	✅ Ready (Inspect)	Visit Preview	💬 Add feedback	Dec 18, 2024 0:49am

[codesandbox-ci]

This pull request is automatically built and testable in CodeSandbox.

To see build info of the built libraries, click here or the icon next to each commit SHA.

[github-actions]

Preview in LiveCodes

Latest commit: c758a69
Last updated: Dec 18, 2024 12:49am (UTC)

Playground	Link
React demo	https://livecodes.io?x=id/4ACRS53X4

---

See documentations for usage instructions.

[dai-shi]

• Unrelated: aState.m.l subscribers are not deduped in flushPending Fixed

The commit link is not found. (because of force pushing?)

[dmaskasky]

Actual change

- while (pending[1].size || pending[2].size) {
-   pending[0].clear()
+ while (pending[0].size || pending[1].size || pending[2].size) {
+   recomputeDependents(pending, new Set(pending[0].keys()))

[dai-shi]

Thanks so much for working on it. The commit history tells how challenging this has been.

I only check the diff, and assuming that this keeps the basic logic and only defers the execution, it's pretty close to what I would do. What's blocking is the reasonable understanding of the store.test.tsx change and failed CI tests.

You can stop it here and I can take it over, if you like.

[dai-shi]

While I'm not sure which is better, I would try keeping it outside. I'm not sure about the future, and we may put "pending" in buildStore. In general, I prefer isolating it.

[dai-shi]

I'm not sure why this is necessary.

[dmaskasky]

Perhaps I should leave a comment here.

Before we get any atom, we need to recompute the dependents of any of its dependencies that have been marked as dirty, to ensure we are returning the latest value. This is the exception to the lazy recomputeDependents rule.

Example:

const a = atom(0)
const b = atom(0)
const c = atom((get) => [get(a), get(b)])
const d = atom((get) => get(c))
const e = atom((get) => get(a))
const w = atom(null, (get, set) => {
  get(d) // [0, 0]

  // does not immediately recompute,
  // marks `a` and all its dependents as dirty
  set(a, 1)

  // setA: dependencies of `d` are {c,a,b}
  // setB: dirty atoms of setA are {a}
  // setC: dependents setB are {a,c,e}
  // setD: intersection of setC and setB are {a,c}
  // recompute setD
  get(d) // [1, 0]

  // recompute remaining dependents of dirty atoms in flushPending {e}
})
store.set(w)

[dai-shi]

I'm not convinced but maybe I need to play with it to get better understanding. It's my problem.

[dmaskasky]

I made this graphic to help me reason about this better. Red is setD in this example.

[dmaskasky]

I think I figured out why this passes the peek test. Here's a better test.
This test also fails on main, but passes on fix/deep-peek. This is good to know, but I don't think this is something we want to fix.

it('can read in write function without changing dependencies', () => {
  const a = atom(0)
  let bReadCount = 0
  const b = atom(
    (get) => {
      ++bReadCount
      return get(a)
    },
    () => {},
  )
  let bMountCalled = false
  b.onMount = () => {
    bMountCalled = true
  }
  const c = atom((get) => get(a))
  const w = atom(null, (get, set) => {
    expect(bReadCount).toBe(0)
    get(b)
    expect(bReadCount).toBe(1)
    set(a, 1)
    expect(bReadCount).toBe(1)
  })
  const store = createStore()
  store.sub(c, () => {}) // mounts c,a
  store.set(w)
  
  expect(bMountCalled).toBe(false) // ✅ PASSES (should fail if b atomState.m is defined
  expect(bReadCount).toBe(1) // ❌ FAILS - received 2
  expect(getAtomState(b)!.m).toBeUndefined() // ✅ PASSES (could fail)
  expect(getAtomState(b)!.d.size).toBe(0) // ❌ FAILS - 1 (a atomState)
  expect(getAtomState(a)!.m!.d.has(b)).toBe(false) // ❌ FAILS - true
})

[pkg-pr-new]

Open in Stackblitz

More templates

• hacker_news
• hello
• mega-form
• text_length
• todos
• todos_with_atomFamily

npm i https://pkg.pr.new/jotai@2827

commit: c758a69

[dai-shi]

Thanks so much for your effort! This is a big improvement.