Try adding resultEqualityCheck to weakMapMemoize

package.json

@@ -50,8 +50,11 @@
   },
   "license": "MIT",
   "devDependencies": {
-    "@reduxjs/toolkit": "^1.9.3",
+    "@reduxjs/toolkit": "^2.0.0-rc.1",
+    "@testing-library/react": "^14.1.2",
     "@types/lodash": "^4.14.175",
+    "@types/react": "^18.2.38",
+    "@types/react-dom": "^18.2.17",
     "@types/shelljs": "^0.8.11",
     "@typescript-eslint/eslint-plugin": "5.1.0",
     "@typescript-eslint/eslint-plugin-tslint": "5.1.0",
@@ -60,13 +63,14 @@
     "eslint": "^8.0.1",
     "eslint-plugin-react": "^7.26.1",
     "eslint-plugin-typescript": "0.14.0",
+    "jsdom": "^23.0.0",
     "lodash.memoize": "^4.1.2",
     "memoize-one": "^6.0.0",
     "micro-memoize": "^4.0.9",
     "prettier": "^2.7.1",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",
-    "react-redux": "^8.0.2",
+    "react-redux": "^9.0.0-rc.0",
     "rimraf": "^3.0.2",
     "shelljs": "^0.8.5",
     "tsup": "^6.7.0",

src/autotrackMemoize/autotrackMemoize.ts

@@ -4,12 +4,8 @@ import type { Node } from './tracking'
 import {
   createCacheKeyComparator,
   defaultEqualityCheck
-} from '@internal/defaultMemoize'
-import type {
-  AnyFunction,
-  DefaultMemoizeFields,
-  Simplify
-} from '@internal/types'
+} from '../defaultMemoize'
+import type { AnyFunction, DefaultMemoizeFields, Simplify } from '../types'
 import { createCache } from './autotracking'
 
 /**

src/autotrackMemoize/autotracking.ts

@@ -3,8 +3,8 @@
 // Additional references:
 // - https://www.pzuraq.com/blog/how-autotracking-works
 // - https://v5.chriskrycho.com/journal/autotracking-elegant-dx-via-cutting-edge-cs/
-import type { EqualityFn } from '@internal/types'
-import { assertIsFunction } from '@internal/utils'
+import type { EqualityFn } from '../types'
+import { assertIsFunction } from '../utils'
 
 // The global revision clock. Every time state changes, the clock increments.
 export let $REVISION = 0

src/defaultMemoize.ts

@@ -5,12 +5,12 @@ import type {
   Simplify
 } from './types'
 
+import type { NOT_FOUND_TYPE } from './utils'
+import { NOT_FOUND } from './utils'
+
 // Cache implementation based on Erik Rasmussen's `lru-memoize`:
 // https://github.com/erikras/lru-memoize
 
-const NOT_FOUND = 'NOT_FOUND'
-type NOT_FOUND_TYPE = typeof NOT_FOUND
-
 interface Entry {
   key: unknown
   value: unknown
@@ -182,6 +182,8 @@ export function defaultMemoize<Func extends AnyFunction>(
 
   const comparator = createCacheKeyComparator(equalityCheck)
 
+  let resultsCount = 0
+
   const cache =
     maxSize === 1
       ? createSingletonCache(comparator)
@@ -193,6 +195,7 @@ export function defaultMemoize<Func extends AnyFunction>(
     if (value === NOT_FOUND) {
       // @ts-ignore
       value = func.apply(null, arguments)
+      resultsCount++
 
       if (resultEqualityCheck) {
         const entries = cache.getEntries()
@@ -202,6 +205,7 @@ export function defaultMemoize<Func extends AnyFunction>(
 
         if (matchingEntry) {
           value = matchingEntry.value
+          resultsCount--
         }
       }
 
@@ -212,6 +216,13 @@ export function defaultMemoize<Func extends AnyFunction>(
 
   memoized.clearCache = () => {
     cache.clear()
+    memoized.resetResultsCount()
+  }
+
+  memoized.resultsCount = () => resultsCount
+
+  memoized.resetResultsCount = () => {
+    resultsCount = 0
   }
 
   return memoized as Func & Simplify<DefaultMemoizeFields>

src/types.ts

@@ -401,6 +401,8 @@ export type DefaultMemoizeFields = {
    * that future calls to the function recompute the results.
    */
   clearCache: () => void
+  resultsCount: () => number
+  resetResultsCount: () => void
 }
 
 /*
@@ -464,7 +466,7 @@ export type FunctionType<T> = Extract<T, AnyFunction>
  */
 export type ExtractReturnType<FunctionsArray extends readonly AnyFunction[]> = {
   [Index in keyof FunctionsArray]: FunctionsArray[Index] extends FunctionsArray[number]
-    ? FallbackIfUnknown<FallbackIfUnknown<ReturnType<FunctionsArray[Index]>, any>, any>
+    ? FallbackIfUnknown<ReturnType<FunctionsArray[Index]>, any>
     : never
 }
 

src/utils.ts

@@ -6,6 +6,9 @@ import type {
   UnknownMemoizer
 } from './types'
 
+export const NOT_FOUND = 'NOT_FOUND'
+export type NOT_FOUND_TYPE = typeof NOT_FOUND
+
 /**
  * Assert that the provided value is a function. If the assertion fails,
  * a `TypeError` is thrown with an optional custom error message.

src/weakMapMemoize.ts

@@ -1,7 +1,21 @@
 // Original source:
 // - https://github.com/facebook/react/blob/0b974418c9a56f6c560298560265dcf4b65784bc/packages/react/src/ReactCache.js
 
-import type { AnyFunction, DefaultMemoizeFields, Simplify } from './types'
+import type {
+  AnyFunction,
+  DefaultMemoizeFields,
+  EqualityFn,
+  Simplify
+} from './types'
+
+class StrongRef<T> {
+  constructor(private value: T) {}
+  deref() {
+    return this.value
+  }
+}
+
+const Ref = WeakRef ?? StrongRef
 
 const UNTERMINATED = 0
 const TERMINATED = 1
@@ -55,6 +69,22 @@ function createCacheNode<T>(): CacheNode<T> {
   }
 }
 
+/**
+ * @public
+ */
+export interface WeakMapMemoizeOptions {
+  /**
+   * If provided, used to compare a newly generated output value against previous values in the cache.
+   * If a match is found, the old value is returned. This addresses the common
+   * ```ts
+   * todos.map(todo => todo.id)
+   * ```
+   * use case, where an update to another field in the original data causes a recalculation
+   * due to changed references, but the output is still effectively the same.
+   */
+  resultEqualityCheck?: EqualityFn
+}
+
 /**
  * Creates a tree of `WeakMap`-based cache nodes based on the identity of the
  * arguments it's been called with (in this case, the extracted values from your input selectors).
@@ -128,8 +158,16 @@ function createCacheNode<T>(): CacheNode<T> {
  * @public
  * @experimental
  */
-export function weakMapMemoize<Func extends AnyFunction>(func: Func) {
+export function weakMapMemoize<Func extends AnyFunction>(
+  func: Func,
+  options: WeakMapMemoizeOptions = {}
+) {
   let fnNode = createCacheNode()
+  const { resultEqualityCheck } = options
+
+  let lastResult: WeakRef<object> | undefined
+
+  let resultsCount = 0
 
   function memoized() {
     let cacheNode = fnNode
@@ -167,19 +205,46 @@ export function weakMapMemoize<Func extends AnyFunction>(func: Func) {
         }
       }
     }
+
+    const terminatedNode = cacheNode as unknown as TerminatedCacheNode<any>
+
+    let result
+
     if (cacheNode.s === TERMINATED) {
-      return cacheNode.v
+      result = cacheNode.v
+    } else {
+      // Allow errors to propagate
+      result = func.apply(null, arguments as unknown as any[])
+      resultsCount++
     }
-    // Allow errors to propagate
-    const result = func.apply(null, arguments as unknown as any[])
-    const terminatedNode = cacheNode as unknown as TerminatedCacheNode<any>
+
     terminatedNode.s = TERMINATED
+
+    if (resultEqualityCheck) {
+      const lastResultValue = lastResult?.deref() ?? lastResult
+      if (lastResultValue != null && resultEqualityCheck(lastResultValue, result)) {
+        result = lastResultValue
+        resultsCount !== 0 && resultsCount--
+      }
+
+      const needsWeakRef =
+        (typeof result === 'object' && result !== null) ||
+        typeof result === 'function'
+      lastResult = needsWeakRef ? new Ref(result) : result
+    }
     terminatedNode.v = result
     return result
   }
 
   memoized.clearCache = () => {
     fnNode = createCacheNode()
+    memoized.resetResultsCount()
+  }
+
+  memoized.resultsCount = () => resultsCount
+
+  memoized.resetResultsCount = () => {
+    resultsCount = 0
   }
 
   return memoized as Func & Simplify<DefaultMemoizeFields>