diff --git a/client/lib/tree-select/index.js b/client/lib/tree-select/index.js index bff8ac2d4783..082e3f514b02 100644 --- a/client/lib/tree-select/index.js +++ b/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; } diff --git a/client/lib/tree-select/test/index.js b/client/lib/tree-select/test/index.js index c78dc117a95b..1e27b92e34ee 100644 --- a/client/lib/tree-select/test/index.js +++ b/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 ); + } ); } ); } );