feat(atoms)!: make injectMemo reactive when no deps are passed (#171)

packages/atoms/src/classes/SelectorInstance.ts

@@ -97,7 +97,10 @@ export const swapSelectorRefs = <G extends SelectorGenerics>(
 ) => {
   const baseKey = ecosystem.b.get(oldInstance.t)
 
-  if (!baseKey) return // TODO: remove
+  // TODO: remove. This is currently needed for selectors created outside the
+  // ecosystem (e.g. via `new SelectorInstance`). Only the ecosystem `getNode*`
+  // methods add the selector ref to `ecosystem.b`aseKeys. Change that.
+  if (!baseKey) return
 
   ecosystem.b.set(newRef, baseKey)
   ecosystem.b.delete(oldInstance.t)

packages/atoms/src/injectors/injectMemo.ts

@@ -1,14 +1,21 @@
+import { SelectorInstance } from '../classes/SelectorInstance'
 import { InjectorDeps } from '../types/index'
 import { untrack } from '../utils/evaluationContext'
 import { compare } from '../utils/general'
 import { injectPrevDescriptor, setNextInjector } from './injectPrevDescriptor'
+import { injectSelf } from './injectSelf'
 
 interface MemoValue<T> {
   /**
    * `d`eps - the cached `injectMemo` deps array
    */
   d: InjectorDeps
 
+  /**
+   * `n`ode - the cached selector instance for an injectMemo call with no deps
+   */
+  n: SelectorInstance | undefined
+
   /**
    * `v`alue - the cached `injectMemo` result
    */
@@ -26,17 +33,32 @@ export const injectMemo = <T = any>(
   valueFactory: () => T,
   deps?: InjectorDeps
 ): T => {
+  const instance = injectSelf()
   const prevDescriptor = injectPrevDescriptor<MemoValue<T>>(TYPE)
+  const depsUnchanged = compare(prevDescriptor?.v.d, deps)
+
+  if (depsUnchanged) return setNextInjector(prevDescriptor!).v.v
+
+  // define this object first so the callback has access to it
+  const value: MemoValue<T> = {
+    d: deps,
+    n: prevDescriptor?.v.n,
+    v: undefined as T,
+  }
+
+  value.v = deps
+    ? untrack(valueFactory)
+    : (value.n ??= new SelectorInstance(
+        instance.e,
+        instance.e._idGenerator.generateId(`injectMemo(${instance.id})-`),
+        valueFactory,
+        []
+      )).get()
 
   return setNextInjector({
     c: undefined,
     i: undefined,
     t: TYPE,
-    v: {
-      d: deps,
-      v: compare(prevDescriptor?.v.d, deps)
-        ? prevDescriptor!.v.v
-        : untrack(valueFactory),
-    },
+    v: value,
   }).v.v
 }

packages/atoms/src/utils/graph.ts

@@ -15,7 +15,7 @@ import {
 
 /**
  * Actually add an edge to the graph. When we buffer graph updates, we're
- * really just deferring the calling of this method.
+ * really just deferring the calling of this function.
  */
 export const addEdge = (
   observer: GraphNode,
@@ -46,7 +46,7 @@ export const addEdge = (
 
 export const destroyNodeStart = (node: GraphNode, force?: boolean) => {
   // If we're not force-destroying, don't destroy if there are dependents. Also
-  // don't destroy of `node.K`eep is set
+  // don't destroy if `node.K`eep is set
   if (node.l === 'Destroyed' || (!force && node.o.size - (node.L ? 1 : 0))) {
     return
   }

packages/react/test/integrations/injectors.test.tsx

@@ -54,6 +54,9 @@ describe('injectors', () => {
 
     const atom1 = atom('1', () => {
       const signal = injectSignal('a')
+
+      // no deps is the exception - this does not mimic React, it turns the
+      // callback into a new reactive context that auto-tracks signal usages
       const val1 = injectMemo(() => signal.get())
       const val2 = injectMemo(() => signal.get(), [])
       const val3 = injectMemo(() => signal.get(), [signal.get()])
@@ -62,6 +65,11 @@ describe('injectors', () => {
       const injectedRef = injectRef(ref)
       refs.push(injectedRef.current)
 
+      // no deps is the exception - this does not mimic React, it turns the
+      // callback into a new reactive context that auto-tracks signal usages.
+      // This is an example of a bad `injectCallback` usage - when using
+      // reactive mode, you can't change the callback function body like this.
+      // You also can't access any non-reactive, closed-over values:
       const cb1 = injectCallback(signal.get() === 'a' ? cbA : cbB)
       const cb2 = injectCallback(signal.get() === 'a' ? cbA : cbB, [])
       const cb3 = injectCallback(signal.get() === 'a' ? cbA : cbB, [
@@ -83,18 +91,19 @@ describe('injectors', () => {
     instance.set('b')
 
     // all those `get` calls shouldn't add any edges besides the one already
-    // added by `injectSignal`:
-    expect(instance.s.size).toBe(1)
+    // added by `injectSignal`. The single `injectMemo` and `injectCallback`
+    // calls with no deps will also each add one source:
+    expect(instance.s.size).toBe(3)
     expect(vals).toEqual(['a', 'a', 'a', 'b', 'a', 'b'])
-    expect(cbs).toEqual(['aa', 'aa', 'aa', 'bb', 'aa', 'bb'])
+    expect(cbs).toEqual(['aa', 'aa', 'aa', 'aa', 'aa', 'bb'])
     expect(effects).toEqual(['a', 'b'])
     expect(cleanups).toEqual(['b'])
     expect(refs).toEqual([ref, ref])
 
     instance.set('c')
 
     expect(vals).toEqual(['a', 'a', 'a', 'b', 'a', 'b', 'c', 'a', 'c'])
-    expect(cbs).toEqual(['aa', 'aa', 'aa', 'bb', 'aa', 'bb', 'bb', 'aa', 'bb'])
+    expect(cbs).toEqual(['aa', 'aa', 'aa', 'aa', 'aa', 'bb', 'aa', 'aa', 'bb'])
     expect(effects).toEqual(['a', 'b', 'c'])
     expect(cleanups).toEqual(['b', 'c'])
     expect(refs).toEqual([ref, ref, ref])
@@ -308,6 +317,40 @@ describe('injectors', () => {
     expect(node1.get()).toBe(0)
   })
 
+  test('when called with no deps, injectMemo() tracks used signals', () => {
+    const signal1 = ecosystem.signal(1)
+    const signal2 = ecosystem.signal(2)
+    const calls: any[] = []
+
+    const atom1 = atom('1', () => {
+      const val2 = signal2.get()
+
+      const memoVal = injectMemo(() => {
+        const val = signal1.get()
+        calls.push(val)
+
+        return val
+      })
+
+      return val2 + memoVal
+    })
+
+    const node1 = ecosystem.getNode(atom1)
+
+    expect(node1.get()).toBe(3)
+    expect(calls).toEqual([1])
+
+    signal1.set(11)
+
+    expect(node1.get()).toBe(13)
+    expect(calls).toEqual([1, 11])
+
+    signal2.set(22)
+
+    expect(node1.get()).toBe(33)
+    expect(calls).toEqual([1, 11])
+  })
+
   test('injectEffect() callback does not track signal usages', () => {
     const signal = ecosystem.signal(0)
 

packages/react/test/stores/injectors.test.tsx

@@ -80,16 +80,16 @@ describe('injectors', () => {
 
     instance.setState('b')
 
-    expect(vals).toEqual(['a', 'a', 'a', 'b', 'a', 'b'])
-    expect(cbs).toEqual(['aa', 'aa', 'aa', 'bb', 'aa', 'bb'])
+    expect(vals).toEqual(['a', 'a', 'a', 'a', 'a', 'b'])
+    expect(cbs).toEqual(['aa', 'aa', 'aa', 'aa', 'aa', 'bb'])
     expect(effects).toEqual(['a', 'b'])
     expect(cleanups).toEqual(['b'])
     expect(refs).toEqual([ref, ref])
 
     instance.setState('c')
 
-    expect(vals).toEqual(['a', 'a', 'a', 'b', 'a', 'b', 'c', 'a', 'c'])
-    expect(cbs).toEqual(['aa', 'aa', 'aa', 'bb', 'aa', 'bb', 'bb', 'aa', 'bb'])
+    expect(vals).toEqual(['a', 'a', 'a', 'a', 'a', 'b', 'a', 'a', 'c'])
+    expect(cbs).toEqual(['aa', 'aa', 'aa', 'aa', 'aa', 'bb', 'aa', 'aa', 'bb'])
     expect(effects).toEqual(['a', 'b', 'c'])
     expect(cleanups).toEqual(['b', 'c'])
     expect(refs).toEqual([ref, ref, ref])