Rewrite defaultMemoize to use an LRU cache with configurable size

Implementation based on Erik Rasmussen's `lru-memoize`:
https://github.com/erikras/lru-memoize

Author: markerikson

src/defaultMemoize.ts

@@ -1,31 +1,98 @@
 import type { EqualityFn } from './types'
 
-export const defaultEqualityCheck: EqualityFn = (a, b): boolean => {
-  return a === b
+// Cache implementation based on Erik Rasmussen's `lru-memoize`:
+// https://github.com/erikras/lru-memoize
+
+interface Entry {
+  key: any
+  value: any
+}
+
+interface Cache {
+  get(key: any): Entry | undefined
+  put(key: any, value: any): void
 }
 
-function areArgumentsShallowlyEqual(
-  equalityCheck: EqualityFn,
-  prev: unknown[] | IArguments | null,
-  next: unknown[] | IArguments | null
-): boolean {
-  if (prev === null || next === null || prev.length !== next.length) {
-    return false
+function createSingletonCache(equals: EqualityFn): Cache {
+  let entry: Entry
+  return {
+    get(key: any) {
+      if (entry && equals(entry.key, key)) {
+        return entry.value
+      }
+    },
+
+    put(key: any, value: any) {
+      entry = { key, value }
+    }
   }
+}
 
-  // Do this in a for loop (and not a `forEach` or an `every`) so we can determine equality as fast as possible.
-  const length = prev.length
-  for (let i = 0; i < length; i++) {
-    if (!equalityCheck(prev[i], next[i])) {
-      return false
+function createLruCache(maxSize: number, equals: EqualityFn): Cache {
+  const entries: Entry[] = []
+
+  function get(key: any) {
+    const cacheIndex = entries.findIndex(entry => equals(key, entry.key))
+
+    // We found a cached entry
+    if (cacheIndex > -1) {
+      const entry = entries[cacheIndex]
+
+      // Cached entry not at top of cache, move it to the top
+      if (cacheIndex > 0) {
+        entries.splice(cacheIndex, 1)
+        entries.unshift(entry)
+      }
+
+      return entry.value
+    }
+
+    // No entry found in cache, return null
+    return undefined
+  }
+
+  function put(key: any, value: any) {
+    if (!get(key)) {
+      // TODO Is unshift slow?
+      entries.unshift({ key, value })
+      if (entries.length > maxSize) {
+        entries.pop()
+      }
     }
   }
 
-  return true
+  return { get, put }
+}
+
+export const defaultEqualityCheck: EqualityFn = (a, b): boolean => {
+  return a === b
+}
+
+function createCacheKeyComparator(equalityCheck: EqualityFn) {
+  return function areArgumentsShallowlyEqual(
+    prev: unknown[] | IArguments | null,
+    next: unknown[] | IArguments | null
+  ): boolean {
+    if (prev === null || next === null || prev.length !== next.length) {
+      return false
+    }
+
+    // Do this in a for loop (and not a `forEach` or an `every`) so we can determine equality as fast as possible.
+    const length = prev.length
+    for (let i = 0; i < length; i++) {
+      if (!equalityCheck(prev[i], next[i])) {
+        return false
+      }
+    }
+
+    return true
+  }
 }
 
 export interface DefaultMemoizeOptions {
   equalityCheck?: EqualityFn
+  resultEqualityCheck?: EqualityFn
+  maxSize?: number
 }
 
 // defaultMemoize now supports a configurable cache size and comparison of the result value.
@@ -34,25 +101,35 @@ export function defaultMemoize<F extends (...args: any[]) => any>(
   func: F,
   equalityCheckOrOptions?: EqualityFn | DefaultMemoizeOptions
 ): F {
-  let lastArgs: any = null
-  let lastResult: any = null
-
   const providedOptions =
     typeof equalityCheckOrOptions === 'object'
       ? equalityCheckOrOptions
       : { equalityCheck: equalityCheckOrOptions }
 
-  const { equalityCheck = defaultEqualityCheck } = providedOptions
+  const {
+    equalityCheck = defaultEqualityCheck,
+    maxSize = 1,
+    resultEqualityCheck
+  } = providedOptions
+
+  const comparator = createCacheKeyComparator(equalityCheck)
+  let resultComparator = resultEqualityCheck
+    ? createCacheKeyComparator(resultEqualityCheck)
+    : undefined
+
+  const cache =
+    maxSize === 1
+      ? createSingletonCache(comparator)
+      : createLruCache(maxSize, comparator)
 
   // we reference arguments instead of spreading them for performance reasons
   return function () {
-    if (!areArgumentsShallowlyEqual(equalityCheck, lastArgs, arguments)) {
-      // apply arguments instead of spreading for performance.
+    let value = cache.get(arguments)
+    if (value === undefined) {
       // @ts-ignore
-      lastResult = func.apply(null, arguments)
+      value = func.apply(null, arguments)
+      cache.put(arguments, value)
     }
-
-    lastArgs = arguments
-    return lastResult
+    return value
   } as F
 }

test/test_selector.ts

@@ -447,13 +447,72 @@ describe('defaultMemoize', () => {
     // call with different object does shallow compare
     expect(fallthroughs).toBe(1)
 
+    /*
+    This test was useful when we had a cache size of 1 previously, and always saved `lastArgs`.
+    But, with the new implementation, this doesn't make sense any more.
+
+
     // the third call does not fall through because `defaultMemoize` passes `anotherObject` as
     // both the `newVal` and `oldVal` params. This allows `shallowEqual` to be much more performant
     // than if it had passed `someObject` as `oldVal`, even though `someObject` and `anotherObject`
     // are shallowly equal
     memoized(anotherObject)
     // call with same object as previous call does not shallow compare
     expect(fallthroughs).toBe(1)
+
+    */
+  })
+
+  test('Accepts a max size greater than 1 with LRU cache behavior', () => {
+    let funcCalls = 0
+
+    const memoizer = defaultMemoize(
+      (state: any) => {
+        funcCalls++
+        return state
+      },
+      {
+        maxSize: 3
+      }
+    )
+
+    // Initial call
+    memoizer('a') // ['a']
+    expect(funcCalls).toBe(1)
+
+    // In cache - memoized
+    memoizer('a') // ['a']
+    expect(funcCalls).toBe(1)
+
+    // Added
+    memoizer('b') // ['b', 'a']
+    expect(funcCalls).toBe(2)
+
+    // Added
+    memoizer('c') // ['c', 'b', 'a']
+    expect(funcCalls).toBe(3)
+
+    // Added, removes 'a'
+    memoizer('d') // ['d', 'c', 'b']
+    expect(funcCalls).toBe(4)
+
+    // No longer in cache, re-added, removes 'b'
+    memoizer('a') // ['a', 'd', 'c']
+    expect(funcCalls).toBe(5)
+
+    // In cache, moved to front
+    memoizer('c') // ['c', 'a', 'd']
+    expect(funcCalls).toBe(5)
+
+    // Added, removes 'd'
+    memoizer('e') // ['e', 'c', 'a']
+    expect(funcCalls).toBe(6)
+
+    // No longer in cache, re-added, removes 'a'
+    memoizer('d') // ['d', 'e', 'c']
+    expect(funcCalls).toBe(7)
+  })
+
   })
 
   test('Accepts an options object as an arg', () => {