diff --git a/src/__tests__/memoize.js b/src/__tests__/memoize.js index 66cc58c0..cb9636c4 100644 --- a/src/__tests__/memoize.js +++ b/src/__tests__/memoize.js @@ -12,10 +12,7 @@ import {memoize} from '../memoize'; import type {Context} from '../types.js'; test('memoize', t => { - // $FlowFixMe - const mockCtx: Context = { - memoized: new Map(), - }; + const mockCtx: Context = ({}: any); let counter = 0; const memoized = memoize(() => { @@ -33,5 +30,15 @@ test('memoize', t => { t.equal(memoizedB(mockCtx), 1, 'memoizes correctly'); t.equal(memoized(mockCtx), 1, 'calls function when it has no value'); t.equal(memoized(mockCtx), 1, 'memoizes correctly'); + + // New context object should cause new calculation + const mockCtx2: Context = ({}: any); + t.equal(memoized(mockCtx2), 2, 'calls function when it has no value'); + t.equal(memoized(mockCtx2), 2, 'memoizes correctly'); + t.equal(memoizedB(mockCtx2), 2, 'calls function when it has no value'); + t.equal(memoizedB(mockCtx2), 2, 'memoizes correctly'); + t.equal(memoized(mockCtx2), 2, 'calls function when it has no value'); + t.equal(memoized(mockCtx2), 2, 'memoizes correctly'); + t.end(); }); diff --git a/src/memoize.js b/src/memoize.js index 14fc8e2d..61721646 100644 --- a/src/memoize.js +++ b/src/memoize.js @@ -10,16 +10,14 @@ import type {Context} from './types.js'; type MemoizeFn = (ctx: Context) => A; -function Container() {} - export function memoize(fn: MemoizeFn): MemoizeFn { - const memoizeKey = __NODE__ ? Symbol('memoize-key') : new Container(); - return function memoized(ctx: Context) { - if (ctx.memoized.has(memoizeKey)) { - return ctx.memoized.get(memoizeKey); + const wm = new WeakMap(); + return ctx => { + if (wm.has(ctx)) { + return ((wm.get(ctx): any): A); // Refinement with `has` doesn't seem to work } const result = fn(ctx); - ctx.memoized.set(memoizeKey, result); + wm.set(ctx, result); return result; }; } diff --git a/src/plugins/timing.js b/src/plugins/timing.js index 8a18d7f0..d7c2154b 100644 --- a/src/plugins/timing.js +++ b/src/plugins/timing.js @@ -43,7 +43,6 @@ const timing: TimingPlugin = { export const TimingToken: Token = createToken('TimingToken'); function middleware(ctx, next) { - ctx.memoized = new Map(); const {start, render, end, downstream, upstream} = timing.from(ctx); ctx.timing = { start,