Bug with ScopeProvider falsely capturing an immer atom

[dbmikus]

When you define an immer atom and the create a ScopeProvider that should not capture the immer atom, it still ends up capturing the immer atom inside the ScopeProvider.

Here is an example where we have 3 atoms:

• a scoped regular atom
• an unscoped regular atom
• a scoped regular atom

We start out with the mounted ScopeProvider. Try incrementing all 3 atoms and then press the button to unmount the ScopeProvider. You would expect to see results:

Counter scoped: 0
Counter immer: 101
Counter unscoped: 1001

but instead you see:

Counter scoped: 0
Counter immer: 100
Counter unscoped: 1001

For some reason, the immer atom ends up scoped within the ScopeProvider.

Here is the code example:

"use client";

import React from "react";
import { useAtom, atom } from "jotai";
import { atomWithImmer } from "jotai-immer";
import { ScopeProvider } from "jotai-scope";

const regAtom = atom({counter: 1000})
const scopedAtom = atom({ counter: 0 });
const immerAtom = atomWithImmer({ counter: 100 });

export default function JotaiBug() {
  const [innerMounted, setInnerMounted] = React.useState(true);

  return (
    <div>
      <div>Are we mounted? {innerMounted}</div>
      {innerMounted ? (
        <ScopeProvider atoms={[scopedAtom]}>
          <div>
            <Inner />
            <div>
              <button onClick={() => setInnerMounted(false)}>
                unmount scoped provider
              </button>
            </div>
          </div>
        </ScopeProvider>
      ) : (
        <CounterDisplay />
      )}
    </div>
  );
}

function Inner() {
  return <CounterDisplay />;
}

function CounterDisplay() {
  const [counterScoped, setState] = useAtom(scopedAtom);
  const [counterImmer, setState2] = useAtom(immerAtom);
  const [counterReg, setState3] = useAtom(regAtom);

  const incScoped = () => {
    setState((prev) => {
      return { counter: prev.counter + 1 };
    });
  };
  const incImmer = () => {
    setState2((prev) => {
      prev.counter++;
    });
  };
  const incReg = () => {
    setState3((prev) => {
      return { counter: prev.counter + 1 };
    });
  }

  return (
    <div>
      <div>
        <span>Counter scoped: {counterScoped.counter}</span>
        <button style={{ marginLeft: 20 }} onClick={incScoped}>
          increment
        </button>
      </div>
      <div>
        <span>Counter immer: {counterImmer.counter}</span>
        <button style={{ marginLeft: 20 }} onClick={incImmer}>
          increment
        </button>
      </div>
      <div>
        <span>Counter unscoped: {counterReg.counter}</span>
        <button style={{ marginLeft: 20 }} onClick={incReg}>
          increment
        </button>
      </div>
    </div>
  );
}

[dai-shi]

Thanks for reporting!
I wonder if someone can dig into it.
(But, someone only I know is @yf-yang )

@dbmikus One thing I want you to try is if you can reproduce the same bug with atomWithReducer.

[yf-yang]

Just saw the issue :), I can take a look when I have the time, but I cannot guarantee when I will do it. Ping me in 2 weeks if I forget it.

[dbmikus]

@dbmikus One thing I want you to try is if you can reproduce the same bug with atomWithReducer.

Unfortunately, I can't put time into trying to reproduce or fix right now. I just stopped scoping the atom for now. If I have time free up to investigate/attempt a fix, I will post here.

[yf-yang]

@dai-shi Check this out: codesandbox

I've added two lines:

// App.js
const immerAtom = atomWithImmer({ counter: 100 });
immerAtom.debugLabel = "immerAtom";
immerAtom.temp = 1;

// patched-jotai-scope.js, in the patched store, when the new atom is created
if (scopedAtom.debugLabel === "immerAtom") {
  scopedAtom.temp = 1000;
}

Open the log console, click "increment" right next to "Counter Immer", and you'll get the point.
However, I am not so clear of how jotai calls atom.read, and my intuition is the point is how the "thisArg" is bind. Another intuition is, current solution does not work for all the atom, but the "pure" atom bypass the flaw in some way, so both reducer and immer do not work.
Do you have any ideas?

[yf-yang]

Hmmm, I have some clues. Take jotai-immer for example (I think reducer is similar)
https://github.com/jotaijs/jotai-immer/blob/2748abcf203daa96ac07cf6ae609a815fcb6eaca/src/atomWithImmer.ts#L13

export function atomWithImmer<Value>(
  initialValue: Value
): WritableAtom<Value, [Value | ((draft: Draft<Value>) => void)], void> {
  const anAtom: any = atom(
    initialValue,
    (get, set, fn: Value | ((draft: Draft<Value>) => void)) =>
      set(
        anAtom,
        produce(
          get(anAtom),
          typeof fn === 'function'
            ? (fn as (draft: Draft<Value>) => void)
            : () => fn
        )
      )
  )
  return anAtom
}

Here, the first argument of the anAtom.write will always be anAtom itself (Actually those wrappers move the anAtom from its external scope into its current scope). Therefore, when we call store.set, which in jotai-scope calls

jotai-scope/src/ScopeProvider.tsx

Lines 53 to 61 in ab16ee3

	...('write' in anAtom && {
	write(get, set, ...args) {
	return anAtom.write.call(
	this,
	(a) => get(getAtom(this, anAtom, a)),
	(a, ...v) => set(getAtom(this, anAtom, a), ...v),
	...args,
	);
	},

We expect the a at line 57 be the copied atom (scopedAtom), but due to the implementation above, it is the original atom. Therefore,

jotai-scope/src/ScopeProvider.tsx

Lines 32 to 41 in ab16ee3

	const getAtom = <A extends AnyAtom>(
	thisArg: AnyAtom,
	orig: AnyAtom,
	target: A,
	): A => {
	if (target === thisArg) {
	return delegate ? getParentScopedAtom(orig as A) : target;
	}
	return getScopedAtom(target);
	};

will Line 37 evaluates to false and the function will return the scopedAtom instead of the correct default atom. Then, in the internal atomStateMap, the scopedAtom's value is incremented, while the original atom's value is unchanged at all. When we remove the ScopeProvider, the original atom is exposed again, so it rewinds back to the untouched original atom's (initial) value.

[dai-shi]

ah, nice catch... that explains the #16 issue too.

[yf-yang]

Maybe refactor getAtom can fix the issue (check if mappings or atomSet has the atom), but frankly speaking I am not so clear what target, thisArg, orig, delegate mean in different scenarios.

[dai-shi]

Yeah, it's confusing to me too. I'll draft a PR later.

[dai-shi]

After thinking about this, this is the issue of atomWithImmer and atomWithReducer implementation. Let me draft PRs there.

[dai-shi]

#16 (comment)

Well, now, I don't know how to fix it...

[yf-yang]

I'll leave my comment here. I am subject to pmndrs/jotai#2351 fix, because that is quite a common technique. I use this trick/recipe a lot, too. It is hard to tell people why we need to do so, and jotai-scope is relative not so broadly used.

Let me think about the implementation.

[yf-yang]

@dai-shi I find an intuitive solution (for this specific case), haven't have it tested with other scenarios and fully consider its validity yet. See if that can give you some inspirations.

I just add a recursive call of getAtom for the return value.

The intuition (my guess) is, the fail case only happens when target is the same atom with this (but it is not scoped because its implementation bypass the ScopeProvider's tracking mechanism), so we explicitly try to track it by calling getScopedAtom(target) and give getAtom another try.

Not sure if the recursion can always correctly terminate.

Besides, I quickly drafted a doc for getAtom, writing it hastily without careful consideration, and its wording still needs to be improved.

/**
 * When an scoped atom call get(anotherAtom) or set(anotherAtom, value), we ensure `anotherAtom` be
 * scoped by calling this function.
 * @param thisArg The scoped atom.
 * @param orig The unscoped original atom of this scoped atom.
 * @param target The `anotherAtom` that this atom is accessing.
 * @returns The scoped target if needed.
 *
 * Check the example below, when calling useAtomValue, jotai-scope first finds the anonymous
 * scoped atom of `anAtom` (we call it `anAtomScoped`). Then, `anAtomScoped.read(dependencyAtom)`
 * becomes `getAtom(anAtomScoped, anAtom, dependencyAtom)`
 * @example
 * const anAtom = atom(get => get(dependencyAtom))
 * const Component = () => {
 *  useAtomValue(anAtom);
 * }
 * const App = () => {
 *   return (
 *    <ScopeProvider atoms={[anAtom]}>
 *      <Component />
 *    </ScopeProvider>
 *   );
 * }
 */
const getAtom = (thisArg, orig, target) => {
  if (thisArg.debugLabel === "immerAtom") {
    console.log(
      `GETATOM thisArg.temp=${thisArg.temp} target.temp=${target.temp} delegate=${delegate}`
    );
  }
  if (target === thisArg) {
    return delegate ? getParentScopedAtom(orig) : target;
  }
  return getAtom(thisArg, orig, getScopedAtom(target));
};

[yf-yang]

Hmmm, another try:

Let me first show my understanding of those two conditions:
delegate === true => this atom is not scoped by ScopeProvider, so all its dependencies should not be scoped, too.
thisArg === target => the atom is trying to read (not possible at all) or write itself (very common in primitive atoms).

Those two conditions form four combinations:
delegate === true && thisArg === target => An unscoped primitive atom, then we should find target in the parent scope, until we find the original global one.
delegate === true && thisArg !== target => An unscoped atom with custom read/write functions. We should find target in the parent scope, until we find the original global one.
delegate === false && thisArg === target => An scoped primitive atom, update the actual value of the scoped one (thisArg itself).
delegate === false && thisArg !== target => An scoped atom with custom read/write functions. Acquire the value of the scoped target.

Therefore, the correct function should be:

const getAtom = (thisArg, orig, target) => {
  if (delegate) {
    return getParentScopedAtom(orig);
  }
  if (target === thisArg) {
    return target;
  }
  return getScopedAtom(target);
};

This analysis sounds more reasonable, it does solve this case, too. I think it won't break existing tests because we only change the behavior of the second case. Still, I'm not sure if it is thorough.

[yf-yang]

I also suggest we change those internal parameter names (delegate, thisArg, orig, target), use clearer terms.

[yf-yang]

OK, that doesn't fix #16, needs more investigation.

[dai-shi]

I also suggest we change those internal parameter names (delegate, thisArg, orig, target), use clearer terms.

Totally agree. If you understand how it's working, can you suggest the terms?

[yf-yang]

Fixed in 0.5.0