add tests to understand atomState

__tests__/derive/derivedStore.ts

@@ -1,10 +1,30 @@
 import { createStore as baseCreateStore } from '../../jotai';
+import { AnyAtom } from './types';
 
-export function createStore() {
+export function createStore(scopedAtoms: Set<AnyAtom> = new Set()) {
   const store = baseCreateStore();
   const derivedStore = store.unstable_derive((getAtomState) => {
+    // return [
+    //   (atom, originAtomState) => {
+    //     return getAtomState(atom, originAtomState);
+    //   },
+    // ];
+    const scopedAtomStateMap = new WeakMap();
+    const scopedAtomStateSet = new WeakSet();
     return [
-      (atom, originAtomState) => {
+      function customGetAtomState(atom, originAtomState) {
+        if (
+          scopedAtomStateSet.has(originAtomState as never) ||
+          scopedAtoms.has(atom)
+        ) {
+          let atomState = scopedAtomStateMap.get(atom);
+          if (!atomState) {
+            atomState = { d: new Map(), p: new Set(), n: 0 };
+            scopedAtomStateMap.set(atom, atomState);
+            scopedAtomStateSet.add(atomState);
+          }
+          return atomState;
+        }
         return getAtomState(atom, originAtomState);
       },
     ];

__tests__/derive/types.ts

@@ -16,14 +16,17 @@ export type AtomState<Value = AnyValue> = {
    * The map value is the epoch number of the dependency.
    */
   readonly d: Map<AnyAtom, number>;
+
   /**
    * Set of atoms with pending promise that depend on the atom.
    *
    * This may cause memory leaks, but it's for the capability to continue promises
    */
   readonly p: Set<AnyAtom>;
+
   /** The epoch number of the atom. */
   n: number;
+
   /**
    * Object to store mounted state of the atom.
    *
@@ -38,15 +41,20 @@ export type AtomState<Value = AnyValue> = {
   m?: {
     /** Set of listeners to notify when the atom value changes. */
     readonly l: Set<() => void>;
+
     /** Set of mounted atoms that the atom depends on. */
     readonly d: Set<AnyAtom>;
+
     /** Set of mounted atoms that depends on the atom. */
     readonly t: Set<AnyAtom>;
+
     /** Function to run when the atom is unmounted. */
     u?: OnUnmount;
   };
+
   /** Atom value */
   v?: Value;
+
   /** Atom error */
   e?: AnyError;
 };

__tests__/derive/understandingAtomState/understandingAtomState.test.ts

@@ -0,0 +1,153 @@
+import { atom, createStore } from '../../../jotai';
+import { assertIsDevStore } from '../../utils';
+
+const store = createStore();
+assertIsDevStore(store);
+const stateMap = store.dev4_get_internal_weak_map();
+type AtomState = NonNullable<ReturnType<typeof stateMap.get>>;
+
+const atomA = atom(0);
+atomA.debugLabel = 'atomA';
+const atomB = atom((get) => String(get(atomA)));
+atomB.debugLabel = 'atomB';
+let atomAState: AtomState;
+let atomBState: AtomState;
+let unsub: () => void;
+
+it('sets d when an atom has a dependency', () => {
+  store.get(atomB);
+
+  atomAState = stateMap.get(atomA)!;
+  atomBState = stateMap.get(atomB)!;
+  /*
+    AtomA state:
+    { d: Map(0) {}, p: Set(0) {}, n: 1, v: 0 }
+
+    AtomB state:
+    {
+      d: Map(1) { atomA => 1 },
+      p: Set(0) {},
+      n: 1,
+      v: '0',
+    }
+  */
+  expect(atomBState.d.has(atomA)).toBe(true);
+});
+
+it('mounts the atoms when an atom has a subscriber', () => {
+  function onAUnmount() {}
+  function onAMount() {
+    return onAUnmount;
+  }
+  atomA.onMount = onAMount;
+
+  function subscribeB() {}
+  unsub = store.sub(atomB, subscribeB);
+
+  /*
+    AtomA state:
+    {
+      d: Map(0) {},
+      p: Set(0) {},
+      n: 1,
+      v: 0,
+      m: {
+        l: Set(0) {},
+        d: Set(0) {},
+        t: Set(1) { [atomB] }
+        u: [Function onAUnmount]
+      }
+    }
+
+    AtomB state:
+    {
+      d: Map(1) { atomA => 1 },
+      p: Set(0) {},
+      n: 1,
+      v: '0',
+      m: {
+        l: Set(1) { [Function: subscribeB] },
+        d: Set(1) { [atomA] },
+        t: Set(0) {}
+      }
+    }
+  */
+
+  expect(atomBState.d.has(atomA)).toBe(true);
+  expect(atomBState.m!.d.has(atomA)).toBe(true);
+  expect(atomAState.m!.t.has(atomB)).toBe(true);
+  expect(atomBState.m!.l.has(subscribeB)).toBe(true);
+  expect(atomAState.m!.u!).toBe(onAUnmount);
+});
+
+it('increments the epoch number when an atom is updated', () => {
+  store.set(atomA, 1);
+  /*
+    AtomA state:
+    {
+      d: Map(0) {},
+      p: Set(0) {},
+      n: 2,
+      v: 1,
+      m: {
+        l: Set(0) {},
+        d: Set(0) {},
+        t: Set(1) { [atomB] }
+        u: [Function onAUnmount]
+      }
+    }
+
+    AtomB state:
+    {
+      d: Map(1) { atomA => 1 },
+      p: Set(0) {},
+      n: 2,
+      v: '1',
+      m: {
+        l: Set(1) { [Function: subscribeB] },
+        d: Set(1) { [atomA] },
+        t: Set(0) {}
+      }
+    }
+  */
+
+  expect(atomAState.n).toBe(2);
+  expect(atomBState.n).toBe(2);
+});
+
+it('unmounts the atoms when there are no subscribers', () => {
+  unsub();
+  /*
+    AtomA state:
+    { d: Map(0) {}, p: Set(0) {}, n: 1, v: 0 }
+
+    AtomB state:
+    {
+      d: Map(1) { atomA => 1 },
+      p: Set(0) {},
+      n: 1,
+      v: '0',
+    }
+  */
+  expect(atomBState.m).toBeUndefined();
+  expect(atomAState.m).toBeUndefined();
+});
+
+it('does not increment the epoch number when an atom is not mounted', () => {
+  store.set(atomA, 2);
+  /*
+    AtomA state:
+    { d: Map(0) {}, p: Set(0) {}, n: 3, v: 2 }
+
+    AtomB state:
+    {
+      d: Map(1) { atomA => 1 },
+      p: Set(0) {},
+      n: 2,
+      v: '1',
+    }
+  */
+
+  expect(atomAState.n).toBe(3);
+  expect(atomBState.n).toBe(2);
+});

__tests__/utils.ts

@@ -1,4 +1,9 @@
 import { fireEvent } from '@testing-library/react';
+import {
+  INTERNAL_DevStoreRev4,
+  INTERNAL_PrdStore,
+} from '../jotai/vanilla/store';
+import { Store } from 'src/ScopeProvider/types';
 
 function getElements(
   container: HTMLElement,
@@ -29,3 +34,30 @@ export function clickButton(container: HTMLElement, querySelector: string) {
   }
   fireEvent.click(button);
 }
+
+export function getDevStore(
+  store: Store,
+): INTERNAL_PrdStore & INTERNAL_DevStoreRev4 {
+  if (!isDevStore(store)) {
+    throw new Error('Store is not a dev store');
+  }
+  return store;
+}
+
+export function isDevStore(
+  store: Store,
+): store is INTERNAL_PrdStore & INTERNAL_DevStoreRev4 {
+  return (
+    'dev4_get_internal_weak_map' in store &&
+    'dev4_get_mounted_atoms' in store &&
+    'dev4_restore_atoms' in store
+  );
+}
+
+export function assertIsDevStore(
+  store: Store,
+): asserts store is INTERNAL_PrdStore & INTERNAL_DevStoreRev4 {
+  if (!isDevStore(store)) {
+    throw new Error('Store is not a dev store');
+  }
+}