Some atoms do not re-render if other atoms are conditionally read them

[ics-matsumoto]

Bug Description

Hi,

I found atoms some times not updated dependancies if I use early return or other conditional statements in derived atom.

I have created the following atoms: 1 primitive atom and 3 derived atoms. The last one conditionally depends on the other 2 derived atoms. I think it's not a particularly special use case.

// Primitive atom: Login user data
const userAtom = atom<User | undefined>(undefined);

// Derived atoms
const nameAtom = atom((get) => get(userAtom)?.name);
const ageAtom = atom((get) => get(userAtom)?.age);

// Derived atom with conditional dependencies
const ageIfNameAtom = atom((get) => {
  const name = get(nameAtom);
  if (!name) return;
  return get(ageAtom);
});

At the initial render, the values of all 3 derived atoms are undefined.
Then I set a value to the primitive atom.

// Set from components
const loginAction = atom(null, (get, set) => {
  set(userAtom, { name: "Alice", age: 12 });
});

I expected nameAtom and ageAtom to be both recalculated. However, the value of useAtomValue(ageAtom) remains undefined, not 12.

---

This problem occurs under:

• jotai 2.9.2
• jotai 2.8.1 (store2)

It does not occur under:

• jotai 2.8.1
• jotai 2.6.5

---

I don't know much about the core store code but traced what happened in recomputeDependents. This might be of some help.

1. When step 1 is finished, all 4 atoms are listed in topsortedAtoms:

topsortedAtoms = [
  ageAtom,
  ageIfNameAtom,
  nameAtom,
  userAtom
]

2. Recompute topsortedAtoms from the last one.

In this process, ageIfNameAtom is recomputed before ageAtom, and ageAtom is updated to 12 but not added to pending at this time. In the next loop of topsortedAtoms, ageAtom is read again, but because the epochNumber has already been incremented by ageIfNameAtom, ageAtom is not added to pending. Eventually, the value of the ageAtom atom state is updated but does not cause a re-render of useAtomValue(ageAtom).

---

If I comment out the early return from ageIfNameAtom,

const ageIfNameAtom = atom((get) => {
  const name = get(nameAtom);
  // if (!name) return;
  return get(ageAtom);
});

the order of topsortedAtoms is changed, and ageAtom is evaluated before ageIfNameAtom.

topsortedAtoms = [
  ageIfNameAtom,
  ageAtom,
  nameAtom,
  userAtom
]

However, I think that we can't detect changes in ageIfNameAtom's dependencies before reading it, making it impossible to list ageIfNameAtom before ageAtom in topsortedAtoms.

---

I think we might need to check for changes in epoch numbers not only in topsortedAtoms[i] itself but also in its dependencies. The following patch (looks) works in my case:

// Get all epochNumbers before starting to reevaluate topsortedAtoms
const prevEpochNumbers: number[] = topsortedAtoms.map(a => getAtomState(a).n);

for (let i = topsortedAtoms.length - 1; i >= 0; --i) {
  ...
  const prevEpochNumber = prevEpochNumbers[i];
  ...
}

But I can't determine its overall effect and can't judge this approach is right or not.

Reproduction Link

https://codesandbox.io/p/sandbox/jotai-2-9-d7yfkn

[rothsandro]

The sandbox isn't public.

[ics-matsumoto]

thanks. I just set to public.

[dai-shi]

So, it's clearly a regression in store2 introduced in v2.9.0.

I don't know much about the core store code but traced what happened in recomputeDependents

I was hoping to keep the logic of recomputeDependents from the previous store implementation, but probably something got missed.

I think we might need to check for changes in epoch numbers not only in topsortedAtoms[i] itself but also in its dependencies.

I don't look it very closely, but it sounds fair.

But I can't determine its overall effect and can't judge this approach is right or not.

I can't judge that either. Can you open a PR and see if how it goes?

I would like you to first add a failing test and open a PR (which runs CI and fails),
and then add another commit for your patch.

[yuneco]

I checked the test I have added at #2700 on some versions.
This test failed only between 2.9.0 and 2.9.2, passed on 2.6.3, 2.6.4, 2.8.4(not using experimental store2).
So I think this problem itself is regression with new store implementation introduced at 2.9.

[dai-shi]

Thanks for confirming. Hm, maybe we are talking about two different issues and the single fix applies to the two.

[yuneco]

As a separate issue from this bug fix, it might be worth considering adding more information about conditional atom reading and writing to the documentation and examples, as cltnschlosser pointed out. Ideally, I hope, any branching or early return in atoms should be allowed as long as the final state is uniquely determined without looping. However, if there are any limitations, they should be clearly stated in the documentation.

[cltnschlosser]

I think it’s the same bug, but maybe it became easier to reproduce with the new store. You’ll notice Eddie’s test case is a bit more complex than the one you added, where there is multiple layers of derived atoms.

[dai-shi]

Could you guys help this and send PRs to improve test coverage and docs?