Fix resultEqualityCheck behavior in weakMapMemoize

src/weakMapMemoize.ts

@@ -227,25 +227,29 @@ export function weakMapMemoize<Func extends AnyFunction>(
       // Allow errors to propagate
       result = func.apply(null, arguments as unknown as any[])
       resultsCount++
-    }
 
-    terminatedNode.s = TERMINATED
+      if (resultEqualityCheck) {
+        const lastResultValue = lastResult?.deref?.() ?? lastResult
 
-    if (resultEqualityCheck) {
-      const lastResultValue = lastResult?.deref?.() ?? lastResult
-      if (
-        lastResultValue != null &&
-        resultEqualityCheck(lastResultValue as ReturnType<Func>, result)
-      ) {
-        result = lastResultValue
-        resultsCount !== 0 && resultsCount--
-      }
+        if (
+          lastResultValue != null &&
+          resultEqualityCheck(lastResultValue as ReturnType<Func>, result)
+        ) {
+          result = lastResultValue
+
+          resultsCount !== 0 && resultsCount--
+        }
+
+        const needsWeakRef =
+          (typeof result === 'object' && result !== null) ||
+          typeof result === 'function'
 
-      const needsWeakRef =
-        (typeof result === 'object' && result !== null) ||
-        typeof result === 'function'
-      lastResult = needsWeakRef ? new Ref(result) : result
+        lastResult = needsWeakRef ? new Ref(result) : result
+      }
     }
+
+    terminatedNode.s = TERMINATED
+
     terminatedNode.v = result
     return result
   }

test/benchmarks/resultEqualityCheck.bench.ts

@@ -0,0 +1,190 @@
+import type { AnyFunction } from '@internal/types'
+import type { OutputSelector, Selector } from 'reselect'
+import {
+  createSelector,
+  lruMemoize,
+  referenceEqualityCheck,
+  weakMapMemoize
+} from 'reselect'
+import type { Options } from 'tinybench'
+import { bench } from 'vitest'
+import {
+  logSelectorRecomputations,
+  setFunctionNames,
+  setupStore,
+  toggleCompleted,
+  type RootState
+} from '../testUtils'
+
+describe('memoize functions performance with resultEqualityCheck set to referenceEqualityCheck vs. without resultEqualityCheck', () => {
+  describe('comparing selectors created with createSelector', () => {
+    const store = setupStore()
+
+    const arrayOfNumbers = Array.from({ length: 1_000 }, (num, index) => index)
+
+    const commonOptions: Options = {
+      iterations: 10_000,
+      time: 0
+    }
+
+    const runSelector = <S extends Selector>(selector: S) => {
+      arrayOfNumbers.forEach(num => {
+        selector(store.getState())
+      })
+    }
+
+    const createAppSelector = createSelector.withTypes<RootState>()
+
+    const selectTodoIdsWeakMap = createAppSelector(
+      [state => state.todos],
+      todos => todos.map(({ id }) => id)
+    )
+
+    const selectTodoIdsWeakMapWithResultEqualityCheck = createAppSelector(
+      [state => state.todos],
+      todos => todos.map(({ id }) => id),
+      {
+        memoizeOptions: { resultEqualityCheck: referenceEqualityCheck },
+        argsMemoizeOptions: { resultEqualityCheck: referenceEqualityCheck }
+      }
+    )
+
+    const selectTodoIdsLru = createAppSelector(
+      [state => state.todos],
+      todos => todos.map(({ id }) => id),
+      { memoize: lruMemoize, argsMemoize: lruMemoize }
+    )
+
+    const selectTodoIdsLruWithResultEqualityCheck = createAppSelector(
+      [state => state.todos],
+      todos => todos.map(({ id }) => id),
+      {
+        memoize: lruMemoize,
+        memoizeOptions: { resultEqualityCheck: referenceEqualityCheck },
+        argsMemoize: lruMemoize,
+        argsMemoizeOptions: { resultEqualityCheck: referenceEqualityCheck }
+      }
+    )
+
+    const selectors = {
+      selectTodoIdsWeakMap,
+      selectTodoIdsWeakMapWithResultEqualityCheck,
+      selectTodoIdsLru,
+      selectTodoIdsLruWithResultEqualityCheck
+    }
+
+    setFunctionNames(selectors)
+
+    const createOptions = <S extends OutputSelector>(selector: S) => {
+      const options: Options = {
+        setup: (task, mode) => {
+          if (mode === 'warmup') return
+
+          task.opts = {
+            beforeEach: () => {
+              store.dispatch(toggleCompleted(1))
+            },
+
+            afterAll: () => {
+              logSelectorRecomputations(selector)
+            }
+          }
+        }
+      }
+      return { ...commonOptions, ...options }
+    }
+
+    Object.values(selectors).forEach(selector => {
+      bench(
+        selector,
+        () => {
+          runSelector(selector)
+        },
+        createOptions(selector)
+      )
+    })
+  })
+
+  describe('comparing selectors created with memoize functions', () => {
+    const store = setupStore()
+
+    const arrayOfNumbers = Array.from(
+      { length: 100_000 },
+      (num, index) => index
+    )
+
+    const commonOptions: Options = {
+      iterations: 1000,
+      time: 0
+    }
+
+    const runSelector = <S extends Selector>(selector: S) => {
+      arrayOfNumbers.forEach(num => {
+        selector(store.getState())
+      })
+    }
+
+    const selectTodoIdsWeakMap = weakMapMemoize((state: RootState) =>
+      state.todos.map(({ id }) => id)
+    )
+
+    const selectTodoIdsWeakMapWithResultEqualityCheck = weakMapMemoize(
+      (state: RootState) => state.todos.map(({ id }) => id),
+      { resultEqualityCheck: referenceEqualityCheck }
+    )
+
+    const selectTodoIdsLru = lruMemoize((state: RootState) =>
+      state.todos.map(({ id }) => id)
+    )
+
+    const selectTodoIdsLruWithResultEqualityCheck = lruMemoize(
+      (state: RootState) => state.todos.map(({ id }) => id),
+      { resultEqualityCheck: referenceEqualityCheck }
+    )
+
+    const memoizedFunctions = {
+      selectTodoIdsWeakMap,
+      selectTodoIdsWeakMapWithResultEqualityCheck,
+      selectTodoIdsLru,
+      selectTodoIdsLruWithResultEqualityCheck
+    }
+
+    setFunctionNames(memoizedFunctions)
+
+    const createOptions = <
+      Func extends AnyFunction & { resultsCount: () => number }
+    >(
+      memoizedFunction: Func
+    ) => {
+      const options: Options = {
+        setup: (task, mode) => {
+          if (mode === 'warmup') return
+
+          task.opts = {
+            beforeEach: () => {
+              store.dispatch(toggleCompleted(1))
+            },
+
+            afterAll: () => {
+              console.log(
+                memoizedFunction.name,
+                memoizedFunction.resultsCount()
+              )
+            }
+          }
+        }
+      }
+      return { ...commonOptions, ...options }
+    }
+
+    Object.values(memoizedFunctions).forEach(memoizedFunction => {
+      bench(
+        memoizedFunction,
+        () => {
+          runSelector(memoizedFunction)
+        },
+        createOptions(memoizedFunction)
+      )
+    })
+  })
+})

test/computationComparisons.spec.tsx

@@ -304,18 +304,18 @@ describe('resultEqualityCheck in weakMapMemoize', () => {
     expect(memoized.resultsCount()).toBe(5)
 
     expect(memoizedShallow(state)).toBe(memoizedShallow(state))
-    expect(memoizedShallow.resultsCount()).toBe(0)
+    expect(memoizedShallow.resultsCount()).toBe(1)
     expect(memoizedShallow({ ...state })).toBe(memoizedShallow(state))
-    expect(memoizedShallow.resultsCount()).toBe(0)
+    expect(memoizedShallow.resultsCount()).toBe(1)
     expect(memoizedShallow({ ...state })).toBe(memoizedShallow(state))
     // We spread the state to force the function to re-run but the
     // result maintains the same reference because of `resultEqualityCheck`.
     const first = memoizedShallow({ ...state })
-    expect(memoizedShallow.resultsCount()).toBe(0)
+    expect(memoizedShallow.resultsCount()).toBe(1)
     memoizedShallow({ ...state })
-    expect(memoizedShallow.resultsCount()).toBe(0)
+    expect(memoizedShallow.resultsCount()).toBe(1)
     const second = memoizedShallow({ ...state })
-    expect(memoizedShallow.resultsCount()).toBe(0)
+    expect(memoizedShallow.resultsCount()).toBe(1)
     expect(first).toBe(second)
   })
 })

test/inputStabilityCheck.spec.ts

@@ -1,5 +1,12 @@
 import { shallowEqual } from 'react-redux'
-import { createSelector, lruMemoize, setGlobalDevModeChecks } from 'reselect'
+import {
+  createSelector,
+  lruMemoize,
+  referenceEqualityCheck,
+  setGlobalDevModeChecks
+} from 'reselect'
+import type { RootState } from './testUtils'
+import { localTest } from './testUtils'
 
 describe('inputStabilityCheck', () => {
   const consoleSpy = vi.spyOn(console, 'warn').mockImplementation(() => {})
@@ -164,3 +171,105 @@ describe('inputStabilityCheck', () => {
     expect(consoleSpy).not.toHaveBeenCalled()
   })
 })
+
+describe('the effects of inputStabilityCheck with resultEqualityCheck', () => {
+  const createAppSelector = createSelector.withTypes<RootState>()
+
+  const resultEqualityCheck = vi
+    .fn(referenceEqualityCheck)
+    .mockName('resultEqualityCheck')
+
+  afterEach(() => {
+    resultEqualityCheck.mockClear()
+  })
+
+  localTest(
+    'resultEqualityCheck should not be called with empty objects when inputStabilityCheck is set to once and input selectors are stable',
+    ({ store }) => {
+      const selectTodoIds = createAppSelector(
+        [state => state.todos],
+        todos => todos.map(({ id }) => id),
+        {
+          memoizeOptions: { resultEqualityCheck },
+          devModeChecks: { inputStabilityCheck: 'once' }
+        }
+      )
+
+      const firstResult = selectTodoIds(store.getState())
+
+      expect(resultEqualityCheck).not.toHaveBeenCalled()
+
+      const secondResult = selectTodoIds(store.getState())
+
+      expect(firstResult).toBe(secondResult)
+
+      expect(resultEqualityCheck).not.toHaveBeenCalled()
+
+      const thirdResult = selectTodoIds(store.getState())
+
+      expect(secondResult).toBe(thirdResult)
+
+      expect(resultEqualityCheck).not.toHaveBeenCalled()
+    }
+  )
+
+  localTest(
+    'resultEqualityCheck should not be called with empty objects when inputStabilityCheck is set to always and input selectors are stable',
+    ({ store }) => {
+      const selectTodoIds = createAppSelector(
+        [state => state.todos],
+        todos => todos.map(({ id }) => id),
+        {
+          memoizeOptions: { resultEqualityCheck },
+          devModeChecks: { inputStabilityCheck: 'always' }
+        }
+      )
+
+      const firstResult = selectTodoIds(store.getState())
+
+      expect(resultEqualityCheck).not.toHaveBeenCalled()
+
+      const secondResult = selectTodoIds(store.getState())
+
+      expect(firstResult).toBe(secondResult)
+
+      expect(resultEqualityCheck).not.toHaveBeenCalled()
+
+      const thirdResult = selectTodoIds(store.getState())
+
+      expect(secondResult).toBe(thirdResult)
+
+      expect(resultEqualityCheck).not.toHaveBeenCalled()
+    }
+  )
+
+  localTest(
+    'resultEqualityCheck should not be called with empty objects when inputStabilityCheck is set to never and input selectors are unstable',
+    ({ store }) => {
+      const selectTodoIds = createAppSelector(
+        [state => [...state.todos]],
+        todos => todos.map(({ id }) => id),
+        {
+          memoizeOptions: { resultEqualityCheck },
+          devModeChecks: { inputStabilityCheck: 'never' }
+        }
+      )
+
+      const firstResult = selectTodoIds(store.getState())
+
+      expect(resultEqualityCheck).not.toHaveBeenCalled()
+
+      const secondResult = selectTodoIds(store.getState())
+
+      expect(firstResult).toBe(secondResult)
+
+      expect(resultEqualityCheck).not.toHaveBeenCalled()
+
+      const thirdResult = selectTodoIds(store.getState())
+
+      expect(secondResult).toBe(thirdResult)
+
+      expect(resultEqualityCheck).not.toHaveBeenCalled()
+    }
+  )
+})