fix(store): Compare results in addition to arguments change in memoizer (#1175)

Author: alex-okrushko

modules/store/spec/selector.spec.ts

@@ -4,6 +4,8 @@ import {
   createFeatureSelector,
   defaultMemoize,
   createSelectorFactory,
+  resultMemoize,
+  MemoizedProjection,
 } from '@ngrx/store';
 import { map, distinctUntilChanged } from 'rxjs/operators';
 
@@ -285,4 +287,72 @@ describe('Selectors', () => {
       expect(anyFn.calls.count()).toEqual(1);
     });
   });
+
+  describe('resultMemoize', () => {
+    let projectionFnSpy: jasmine.Spy;
+    const ARRAY = ['a', 'ab', 'b'];
+    const ARRAY_CHANGED = [...ARRAY, 'bc'];
+    const A_FILTER: { by: string } = { by: 'a' };
+    const B_FILTER: { by: string } = { by: 'b' };
+
+    let arrayMemoizer: MemoizedProjection;
+
+    // Compare a and b on equality. If a and b are Arrays then compare them
+    // on their content.
+    function isResultEqual(a: any, b: any) {
+      if (a instanceof Array) {
+        return a.length === b.length && a.every(fromA => b.includes(fromA));
+      }
+      // Default comparison
+      return a === b;
+    }
+
+    beforeEach(() => {
+      projectionFnSpy = jasmine
+        .createSpy('projectionFn')
+        .and.callFake((arr: string[], filter: { by: string }) =>
+          arr.filter(item => item.startsWith(filter.by))
+        );
+
+      arrayMemoizer = resultMemoize(projectionFnSpy, isResultEqual);
+    });
+
+    it('should not rerun projector function when arguments stayed the same', () => {
+      arrayMemoizer.memoized(ARRAY, A_FILTER);
+      arrayMemoizer.memoized(ARRAY, A_FILTER);
+
+      expect(projectionFnSpy.calls.count()).toBe(1);
+    });
+
+    it('should rerun projector function when arguments changed', () => {
+      arrayMemoizer.memoized(ARRAY, A_FILTER);
+      arrayMemoizer.memoized(ARRAY_CHANGED, A_FILTER);
+
+      expect(projectionFnSpy.calls.count()).toBe(2);
+    });
+
+    it('should return the same instance of results when projector function produces the same results array', () => {
+      const result1 = arrayMemoizer.memoized(ARRAY, A_FILTER);
+      const result2 = arrayMemoizer.memoized(ARRAY, A_FILTER);
+
+      expect(result1).toBe(result2);
+    });
+
+    it('should return the same instance of results when projector function produces similar results array', () => {
+      const result1 = arrayMemoizer.memoized(ARRAY, A_FILTER);
+      const result2 = arrayMemoizer.memoized(ARRAY_CHANGED, A_FILTER);
+
+      expect(result1).toBe(result2);
+    });
+
+    it('should return the new instance of results when projector function produces different result', () => {
+      const result1 = arrayMemoizer.memoized(ARRAY, A_FILTER);
+      const result2 = arrayMemoizer.memoized(ARRAY_CHANGED, B_FILTER);
+
+      expect(result1).toBeDefined();
+      expect(result2).toBeDefined();
+      expect(result1).not.toBe(result2);
+      expect(result1).not.toEqual(result2);
+    });
+  });
 });

modules/store/src/index.ts

@@ -25,6 +25,7 @@ export {
   MemoizeFn,
   MemoizedProjection,
   MemoizedSelector,
+  resultMemoize,
 } from './selector';
 export { State, StateObservable, reduceState } from './state';
 export {

modules/store/src/selector.ts

@@ -6,6 +6,8 @@ export type MemoizedProjection = { memoized: AnyFn; reset: () => void };
 
 export type MemoizeFn = (t: AnyFn) => MemoizedProjection;
 
+export type ComparatorFn = (a: any, b: any) => boolean;
+
 export interface MemoizedSelector<State, Result>
   extends Selector<State, Result> {
   release(): void;
@@ -16,36 +18,61 @@ export function isEqualCheck(a: any, b: any): boolean {
   return a === b;
 }
 
+function isArgumentsChanged(
+  args: IArguments,
+  lastArguments: IArguments,
+  comparator: ComparatorFn
+) {
+  for (let i = 0; i < args.length; i++) {
+    if (!comparator(args[i], lastArguments[i])) {
+      return true;
+    }
+  }
+  return false;
+}
+
+export function resultMemoize(
+  projectionFn: AnyFn,
+  isResultEqual: ComparatorFn
+) {
+  return defaultMemoize(projectionFn, isEqualCheck, isResultEqual);
+}
+
 export function defaultMemoize(
-  t: AnyFn,
-  isEqual = isEqualCheck
+  projectionFn: AnyFn,
+  isArgumentsEqual = isEqualCheck,
+  isResultEqual = isEqualCheck
 ): MemoizedProjection {
   let lastArguments: null | IArguments = null;
+  // tslint:disable-next-line:no-any anything could be the result.
   let lastResult: any = null;
 
   function reset() {
     lastArguments = null;
     lastResult = null;
   }
 
+  // tslint:disable-next-line:no-any anything could be the result.
   function memoized(): any {
     if (!lastArguments) {
-      lastResult = t.apply(null, arguments);
+      lastResult = projectionFn.apply(null, arguments);
       lastArguments = arguments;
-
       return lastResult;
     }
 
-    for (let i = 0; i < arguments.length; i++) {
-      if (!isEqual(arguments[i], lastArguments[i])) {
-        lastResult = t.apply(null, arguments);
-        lastArguments = arguments;
+    if (!isArgumentsChanged(arguments, lastArguments, isArgumentsEqual)) {
+      return lastResult;
+    }
 
-        return lastResult;
-      }
+    const newResult = projectionFn.apply(null, arguments);
+    if (isResultEqual(lastResult, newResult)) {
+      return lastResult;
     }
 
-    return lastResult;
+    lastResult = newResult;
+    lastArguments = arguments;
+
+    return newResult;
   }
 
   return { memoized, reset };