Use a WeakMap to store memoized values

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();
 });

src/memoize.js

@@ -10,16 +10,14 @@ import type {Context} from './types.js';
 
 type MemoizeFn<A> = (ctx: Context) => A;
 
-function Container() {}
-
 export function memoize<A>(fn: MemoizeFn<A>): MemoizeFn<A> {
-  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;
   };
 }

src/plugins/timing.js

@@ -43,7 +43,6 @@ const timing: TimingPlugin = {
 export const TimingToken: Token<TimingPlugin> = createToken('TimingToken');
 
 function middleware(ctx, next) {
-  ctx.memoized = new Map();
   const {start, render, end, downstream, upstream} = timing.from(ctx);
   ctx.timing = {
     start,