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,