feat(Store): Add support for generating custom createSelector functio…

…ns (#734) Closes #478, #724
ngrx · Jan 22, 2018 · cb0d185 · cb0d185
1 parent b82c35d
commit cb0d185
Show file tree

Hide file tree

Showing 3 changed files with 136 additions and 31 deletions.
diff --git a/modules/store/spec/selector.spec.ts b/modules/store/spec/selector.spec.ts
@@ -1,7 +1,12 @@
 import 'rxjs/add/operator/distinctUntilChanged';
 import 'rxjs/add/operator/map';
 import { cold } from 'jasmine-marbles';
-import { createSelector, createFeatureSelector } from '../';
+import {
+  createSelector,
+  createFeatureSelector,
+  defaultMemoize,
+  createSelectorFactory,
+} from '../';
 
 describe('Selectors', () => {
   let countOne: number;
@@ -229,4 +234,53 @@ describe('Selectors', () => {
       expect(featureState$).toBeObservable(expected$);
     });
   });
+
+  describe('createSelectorFactory', () => {
+    it('should return a selector creator function', () => {
+      const projectFn = jasmine.createSpy('projectionFn');
+      const selectorFunc = createSelectorFactory(defaultMemoize);
+
+      const selector = selectorFunc(incrementOne, incrementTwo, projectFn)({});
+
+      expect(projectFn).toHaveBeenCalledWith(countOne, countTwo);
+    });
+
+    it('should allow a custom memoization function', () => {
+      const projectFn = jasmine.createSpy('projectionFn');
+      const anyFn = jasmine.createSpy('t').and.callFake(() => true);
+      const equalFn = jasmine.createSpy('isEqual').and.callFake(() => true);
+      const customMemoizer = (aFn: any = anyFn, eFn: any = equalFn) =>
+        defaultMemoize(anyFn, equalFn);
+      const customSelector = createSelectorFactory(customMemoizer);
+
+      const selector = customSelector(incrementOne, incrementTwo, projectFn);
+      selector(1);
+      selector(2);
+
+      expect(anyFn.calls.count()).toEqual(1);
+    });
+
+    it('should allow a custom state memoization function', () => {
+      const projectFn = jasmine.createSpy('projectionFn');
+      const stateFn = jasmine.createSpy('stateFn');
+      const selectorFunc = createSelectorFactory(defaultMemoize, { stateFn });
+
+      const selector = selectorFunc(incrementOne, incrementTwo, projectFn)({});
+
+      expect(stateFn).toHaveBeenCalled();
+    });
+  });
+
+  describe('defaultMemoize', () => {
+    it('should allow a custom equality function', () => {
+      const anyFn = jasmine.createSpy('t').and.callFake(() => true);
+      const equalFn = jasmine.createSpy('isEqual').and.callFake(() => true);
+      const memoizer = defaultMemoize(anyFn, equalFn);
+
+      memoizer.memoized(1, 2, 3);
+      memoizer.memoized(1, 2);
+
+      expect(anyFn.calls.count()).toEqual(1);
+    });
+  });
 });
diff --git a/modules/store/src/index.ts b/modules/store/src/index.ts
@@ -18,7 +18,12 @@ export {
 export { ScannedActionsSubject } from './scanned_actions_subject';
 export {
   createSelector,
+  createSelectorFactory,
   createFeatureSelector,
+  defaultMemoize,
+  defaultStateFn,
+  MemoizeFn,
+  MemoizedProjection,
   MemoizedSelector,
 } from './selector';
 export { State, StateObservable, reduceState } from './state';

diff --git a/modules/store/src/selector.ts b/modules/store/src/selector.ts
@@ -2,13 +2,24 @@ import { Selector } from './models';
 
 export type AnyFn = (...args: any[]) => any;
 
+export type MemoizedProjection = { memoized: AnyFn; reset: () => void };
+
+export type MemoizeFn = (t: AnyFn) => MemoizedProjection;
+
 export interface MemoizedSelector<State, Result>
   extends Selector<State, Result> {
   release(): void;
   projector: AnyFn;
 }
 
-export function memoize(t: AnyFn): { memoized: AnyFn; reset: () => void } {
+export function isEqualCheck(a: any, b: any): boolean {
+  return a === b;
+}
+
+export function defaultMemoize(
+  t: AnyFn,
+  isEqual = isEqualCheck
+): MemoizedProjection {
   let lastArguments: null | IArguments = null;
   let lastResult: any = null;
 
@@ -24,8 +35,9 @@ export function memoize(t: AnyFn): { memoized: AnyFn; reset: () => void } {
 
       return lastResult;
     }
+
     for (let i = 0; i < arguments.length; i++) {
-      if (arguments[i] !== lastArguments[i]) {
+      if (!isEqual(arguments[i], lastArguments[i])) {
         lastResult = t.apply(null, arguments);
         lastArguments = arguments;
 
@@ -184,41 +196,75 @@ export function createSelector<State, S1, S2, S3, S4, S5, S6, S7, S8, Result>(
     s8: S8
   ) => Result
 ): MemoizedSelector<State, Result>;
-export function createSelector(...input: any[]): Selector<any, any> {
-  let args = input;
-  if (Array.isArray(args[0])) {
-    const [head, ...tail] = args;
-    args = [...head, ...tail];
-  }
+export function createSelector(...input: any[]) {
+  return createSelectorFactory(defaultMemoize)(...input);
+}
 
-  const selectors = args.slice(0, args.length - 1);
-  const projector = args[args.length - 1];
-  const memoizedSelectors = selectors.filter(
-    (selector: any) =>
-      selector.release && typeof selector.release === 'function'
-  );
+export function defaultStateFn(
+  state: any,
+  selectors: Selector<any, any>[],
+  memoizedProjector: MemoizedProjection
+): any {
+  const args = selectors.map(fn => fn(state));
 
-  const memoizedProjector = memoize(function(...selectors: any[]) {
-    return projector.apply(null, selectors);
-  });
+  return memoizedProjector.memoized.apply(null, args);
+}
 
-  const memoizedState = memoize(function(state: any) {
-    const args = selectors.map(fn => fn(state));
+export type SelectorFactoryConfig<T = any, V = any> = {
+  stateFn: (
+    state: T,
+    selectors: Selector<any, any>[],
+    memoizedProjector: MemoizedProjection
+  ) => V;
+};
 
-    return memoizedProjector.memoized.apply(null, args);
-  });
+export function createSelectorFactory<T = any, V = any>(
+  memoize: MemoizeFn
+): (...input: any[]) => Selector<T, V>;
+export function createSelectorFactory<T = any, V = any>(
+  memoize: MemoizeFn,
+  options: SelectorFactoryConfig<T, V>
+): (...input: any[]) => Selector<T, V>;
+export function createSelectorFactory(
+  memoize: MemoizeFn,
+  options: SelectorFactoryConfig<any, any> = {
+    stateFn: defaultStateFn,
+  }
+) {
+  return function(...input: any[]): Selector<any, any> {
+    let args = input;
+    if (Array.isArray(args[0])) {
+      const [head, ...tail] = args;
+      args = [...head, ...tail];
+    }
 
-  function release() {
-    memoizedState.reset();
-    memoizedProjector.reset();
+    const selectors = args.slice(0, args.length - 1);
+    const projector = args[args.length - 1];
+    const memoizedSelectors = selectors.filter(
+      (selector: any) =>
+        selector.release && typeof selector.release === 'function'
+    );
 
-    memoizedSelectors.forEach(selector => selector.release());
-  }
+    const memoizedProjector = memoize(function(...selectors: any[]) {
+      return projector.apply(null, selectors);
+    });
+
+    const memoizedState = defaultMemoize(function(state: any) {
+      return options.stateFn.apply(null, [state, selectors, memoizedProjector]);
+    });
+
+    function release() {
+      memoizedState.reset();
+      memoizedProjector.reset();
+
+      memoizedSelectors.forEach(selector => selector.release());
+    }
 
-  return Object.assign(memoizedState.memoized, {
-    release,
-    projector: memoizedProjector.memoized,
-  });
+    return Object.assign(memoizedState.memoized, {
+      release,
+      projector: memoizedProjector.memoized,
+    });
+  };
 }
 
 export function createFeatureSelector<T>(