Add several missing features to treeSelect

client/lib/tree-select/index.js

@@ -5,15 +5,20 @@
  */
 import { isObject, some, isFunction } from 'lodash';
 
+const defaultGetCacheKey = ( ...args ) => args.join();
+
 /**
  * Returns a selector that caches values.
  *
  * @param  {Function} getDependents A Function describing the dependent(s) of the selector.
  *                                    Must return an array which gets passed as the first arg to the selector
  * @param  {Function} selector      A standard selector for calculating cached result
+ * @param  {Object}   options       Options bag with additional arguments
+ * @param  {Function} options.getCacheKey
+ *                                  Custom way to compute the cache key from the `args` list
  * @return {Function}               Cached selector
  */
-export default function treeSelect( getDependents, selector ) {
+export default function treeSelect( getDependents, selector, options = {} ) {
 	if ( process.env.NODE_ENV !== 'production' ) {
 		if ( ! isFunction( getDependents ) || ! isFunction( selector ) ) {
 			throw new TypeError(
@@ -22,13 +27,15 @@ export default function treeSelect( getDependents, selector ) {
 		}
 	}
 
-	const cache = new WeakMap();
+	let cache = new WeakMap();
+
+	const { getCacheKey = defaultGetCacheKey } = options;
 
-	return function( state, ...args ) {
+	const cachedSelector = function( state, ...args ) {
 		const dependents = getDependents( state, ...args );
 
 		if ( process.env.NODE_ENV !== 'production' ) {
-			if ( some( args, isObject ) ) {
+			if ( getCacheKey === defaultGetCacheKey && some( args, isObject ) ) {
 				throw new Error( 'Do not pass objects as arguments to a treeSelector' );
 			}
 		}
@@ -38,7 +45,7 @@ export default function treeSelect( getDependents, selector ) {
 		// garbage collect any values that are based on outdated dependents
 		const leafCache = dependents.reduce( insertDependentKey, cache );
 
-		const key = args.join();
+		const key = getCacheKey( ...args );
 		if ( leafCache.has( key ) ) {
 			return leafCache.get( key );
 		}
@@ -47,8 +54,20 @@ export default function treeSelect( getDependents, selector ) {
 		leafCache.set( key, value );
 		return value;
 	};
+
+	cachedSelector.clearCache = () => {
+		// WeakMap doesn't have `clear` method, so we need to recreate it
+		cache = new WeakMap();
+	};
+
+	return cachedSelector;
 }
 
+/*
+ * This object will be used as a WeakMap key if a dependency is a falsy value (null, undefined, ...)
+ */
+const STATIC_FALSY_KEY = {};
+
 /*
  * First tries to get the value for the key.
  * If the key is not present, then inserts a new map and returns it
@@ -57,11 +76,14 @@ export default function treeSelect( getDependents, selector ) {
  * The last map is a regular one because the the key for the last map is the string results of args.join().
  */
 function insertDependentKey( map, key, currentIndex, arr ) {
-	if ( map.has( key ) ) {
-		return map.get( key );
+	const weakMapKey = key || STATIC_FALSY_KEY;
+
+	const existingMap = map.get( weakMapKey );
+	if ( existingMap ) {
+		return existingMap;
 	}
 
 	const newMap = currentIndex === arr.length - 1 ? new Map() : new WeakMap();
-	map.set( key, newMap );
+	map.set( weakMapKey, newMap );
 	return newMap;
 }

client/lib/tree-select/test/index.js

@@ -197,5 +197,59 @@ describe( 'index', () => {
 
 			expect( getDependents ).toHaveBeenCalledWith( [ state.posts ], 1, 2, 3 );
 		} );
+
+		test( 'should bust the cache when clearCache() method is called', () => {
+			const reduxState = {
+				posts: {
+					[ post1.id ]: post1,
+					[ post2.id ]: post2,
+					[ post3.id ]: post3,
+				},
+			};
+
+			const firstResult = getSitePosts( reduxState, 'site1' );
+			const memoizedResult = getSitePosts( reduxState, 'site1' );
+
+			// Repeated call returns identical object
+			expect( memoizedResult ).toBe( firstResult );
+
+			getSitePosts.clearCache();
+			const afterClearResult = getSitePosts( reduxState, 2916284 );
+
+			// Is forced to compute a new result after clearCache()
+			expect( afterClearResult ).not.toBe( firstResult );
+		} );
+
+		test( 'should memoize a falsy value returned by getDependents', () => {
+			const memoizedSelector = treeSelect( () => [ null ], () => [] );
+			const state = {};
+
+			const firstResult = memoizedSelector( state );
+			const secondResult = memoizedSelector( state );
+			expect( firstResult ).toBe( secondResult );
+		} );
+
+		test( 'accepts a getCacheKey option that enables object arguments', () => {
+			const reduxState = {
+				posts: {
+					[ post1.id ]: post1,
+					[ post2.id ]: post2,
+					[ post3.id ]: post3,
+				},
+			};
+
+			const memoizedSelector = treeSelect(
+				state => [ state.posts ],
+				( [ posts ], query ) => filter( posts, { siteId: query.siteId } ),
+				{ getCacheKey: query => `key:${ query.siteId }` }
+			);
+
+			// The arguments are objects, they are not identical, but generated keys are.
+			// Therefore, the second call returns the memoized result.
+			const firstResult = memoizedSelector( reduxState, { siteId: 'site1', foo: 'bar' } );
+			const secondResult = memoizedSelector( reduxState, { siteId: 'site1', foo: 'baz' } );
+			expect( firstResult ).toEqual( [ post1, post2 ] );
+			expect( firstResult ).toBe( secondResult );
+		} );
 	} );
 } );