fix(react): handle unstable inline selector results in useAtomSelector

packages/react/src/hooks/useAtomSelector.ts

@@ -47,14 +47,30 @@ export const useAtomSelector = <T, Args extends any[]>(
   const resolvedArgs = argsChanged ? args : (existingCache.args as Args)
   const renderedResult = ecosystem.select(selectorOrConfig, ...resolvedArgs)
 
+  // When an inline selector returns a referentially unstable result every run,
+  // we have to ignore the subsequent update. Do that using a "state machine"
+  // that goes from 0 -> 1 -> 2. This machine ensures that the ignored update
+  // occurs after the component rerenders and the effect reruns after that
+  // render. This works with strict mode on or off. Use the stable `render`
+  // function as a "ref" :O
+  if ((render as any).ignorePhase === 1) {
+    ;(render as any).ignorePhase = 2
+  }
+
   useEffect(() => {
     // Don't cache the selector until this effect runs. Sadly, this means that
     // all selectors that are first invoked from React will be double-invoked.
     // There's really nothing (performant and good) that we can do about this.
     // React is really just missing lots of features for external stores.
     const cache = selectors.getCache(selectorOrConfig, resolvedArgs)
 
-    if (_graph.nodes[cache.id]?.dependents.get(dependentKey)) return
+    if (_graph.nodes[cache.id]?.dependents.get(dependentKey)) {
+      return () => {
+        if ((render as any).ignorePhase !== 1) {
+          _graph.removeEdge(dependentKey, cache.id)
+        }
+      }
+    }
 
     _graph.addEdge(dependentKey, cache.id, OPERATION, External, () =>
       render({})
@@ -63,14 +79,23 @@ export const useAtomSelector = <T, Args extends any[]>(
     // an unmounting component's effect cleanup can force-destroy the selector
     // or update its dependencies before this component is mounted. If that
     // happened, trigger a rerender to recache the selector and/or get its new
-    // result
-    if (cache.result !== renderedResult) render({})
+    // result. On the rerender, ignore changes
+    if (cache.result !== renderedResult && !(render as any).ignorePhase) {
+      ;(render as any).ignorePhase = 1
+      render({})
+    }
+
+    if ((render as any).ignorePhase === 2) {
+      ;(render as any).ignorePhase = 0
+    }
 
-    // use the referentially stable render function as a ref :O
-    // ;(render as any).mounted = true
     cacheRef.current = { args: resolvedArgs }
 
-    return () => _graph.removeEdge(dependentKey, cache.id)
+    return () => {
+      if ((render as any).ignorePhase !== 1) {
+        _graph.removeEdge(dependentKey, cache.id)
+      }
+    }
   }, [selectorOrConfig, resolvedArgs])
 
   return renderedResult

packages/react/test/units/__snapshots__/useAtomSelector.test.tsx.snap

@@ -0,0 +1,141 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`useAtomSelector inline selector that returns a different object reference every time only triggers one extra rerender (strict mode off) 1`] = `
+{
+  "1": {
+    "dependencies": Map {},
+    "dependents": Map {
+      "@@selector-unnamed-1" => {
+        "callback": undefined,
+        "createdAt": 123456789,
+        "flags": 0,
+        "operation": "get",
+      },
+    },
+    "isSelector": undefined,
+    "refCount": 1,
+    "weight": 1,
+  },
+  "@@selector-unnamed-1": {
+    "dependencies": Map {
+      "1" => true,
+    },
+    "dependents": Map {
+      "Test-:rb:" => {
+        "callback": [Function],
+        "createdAt": 123456789,
+        "flags": 2,
+        "operation": "useAtomSelector",
+      },
+    },
+    "isSelector": true,
+    "refCount": 1,
+    "weight": 2,
+  },
+}
+`;
+
+exports[`useAtomSelector inline selector that returns a different object reference every time only triggers one extra rerender (strict mode off) 2`] = `
+{
+  "1": {
+    "dependencies": Map {},
+    "dependents": Map {
+      "@@selector-unnamed-3" => {
+        "callback": undefined,
+        "createdAt": 123456789,
+        "flags": 0,
+        "operation": "get",
+      },
+    },
+    "isSelector": undefined,
+    "refCount": 1,
+    "weight": 1,
+  },
+  "@@selector-unnamed-3": {
+    "dependencies": Map {
+      "1" => true,
+    },
+    "dependents": Map {
+      "Test-:rb:" => {
+        "callback": [Function],
+        "createdAt": 123456789,
+        "flags": 2,
+        "operation": "useAtomSelector",
+      },
+    },
+    "isSelector": true,
+    "refCount": 1,
+    "weight": 2,
+  },
+}
+`;
+
+exports[`useAtomSelector inline selector that returns a different object reference every time only triggers one extra rerender (strict mode on) 1`] = `
+{
+  "1": {
+    "dependencies": Map {},
+    "dependents": Map {
+      "@@selector-unnamed-1" => {
+        "callback": undefined,
+        "createdAt": 123456789,
+        "flags": 0,
+        "operation": "get",
+      },
+    },
+    "isSelector": undefined,
+    "refCount": 1,
+    "weight": 1,
+  },
+  "@@selector-unnamed-1": {
+    "dependencies": Map {
+      "1" => true,
+    },
+    "dependents": Map {
+      "Test-:rd:" => {
+        "callback": [Function],
+        "createdAt": 123456789,
+        "flags": 2,
+        "operation": "useAtomSelector",
+      },
+    },
+    "isSelector": true,
+    "refCount": 1,
+    "weight": 2,
+  },
+}
+`;
+
+exports[`useAtomSelector inline selector that returns a different object reference every time only triggers one extra rerender (strict mode on) 2`] = `
+{
+  "1": {
+    "dependencies": Map {},
+    "dependents": Map {
+      "@@selector-unnamed-3" => {
+        "callback": undefined,
+        "createdAt": 123456789,
+        "flags": 0,
+        "operation": "get",
+      },
+    },
+    "isSelector": undefined,
+    "refCount": 1,
+    "weight": 1,
+  },
+  "@@selector-unnamed-3": {
+    "dependencies": Map {
+      "1" => true,
+    },
+    "dependents": Map {
+      "Test-:rd:" => {
+        "callback": [Function],
+        "createdAt": 123456789,
+        "flags": 2,
+        "operation": "useAtomSelector",
+      },
+    },
+    "isSelector": true,
+    "refCount": 1,
+    "weight": 2,
+  },
+}
+`;

packages/react/test/units/useAtomSelector.test.tsx

@@ -474,4 +474,74 @@ describe('useAtomSelector', () => {
       }
     `)
   })
+
+  test('inline selector that returns a different object reference every time only triggers one extra rerender (strict mode off)', async () => {
+    const atom1 = atom('1', () => ({ val: 1 }))
+    let renders = 0
+
+    function Test() {
+      renders++
+      const { val } = useAtomSelector(({ get }) => ({ val: get(atom1).val }))
+
+      return (
+        <>
+          <div data-testid="text">{val}</div>
+        </>
+      )
+    }
+
+    const { findByTestId } = renderInEcosystem(<Test />)
+
+    const div = await findByTestId('text')
+
+    expect(div.innerHTML).toBe('1')
+    expect(renders).toBe(2)
+    expect(ecosystem._graph.nodes).toMatchSnapshot()
+
+    act(() => {
+      ecosystem.getInstance(atom1).setState({ val: 2 })
+    })
+
+    await Promise.resolve()
+
+    expect(div.innerHTML).toBe('2')
+    expect(renders).toBe(4)
+    expect(ecosystem._graph.nodes).toMatchSnapshot()
+  })
+
+  test('inline selector that returns a different object reference every time only triggers one extra rerender (strict mode on)', async () => {
+    const atom1 = atom('1', () => ({ val: 1 }))
+    let renders = 0
+
+    function Test() {
+      renders++
+      const { val } = useAtomSelector(({ get }) => ({ val: get(atom1).val }))
+
+      return (
+        <>
+          <div data-testid="text">{val}</div>
+        </>
+      )
+    }
+
+    const { findByTestId } = renderInEcosystem(<Test />, {
+      useStrictMode: true,
+    })
+
+    const div = await findByTestId('text')
+
+    expect(div.innerHTML).toBe('1')
+    expect(renders).toBe(4) // 2 rerenders + 2 for strict mode
+    expect(ecosystem._graph.nodes).toMatchSnapshot()
+
+    act(() => {
+      ecosystem.getInstance(atom1).setState({ val: 2 })
+    })
+
+    await Promise.resolve()
+
+    expect(div.innerHTML).toBe('2')
+    expect(renders).toBe(8) // 4 rerenders + 4 for strict mode
+    expect(ecosystem._graph.nodes).toMatchSnapshot()
+  })
 })