From fca2d93a83be3559d1531325cabb282d2cfea8c0 Mon Sep 17 00:00:00 2001
From: Ryan Tsao <ryan.j.tsao@gmail.com>
Date: Wed, 19 Sep 2018 10:14:37 -0700
Subject: [PATCH 1/3] Use WeakMap to memoize rather than mutating ctx

---
 src/memoize.js        | 37 ++++++++++++++++++++++++++++---------
 src/plugins/timing.js |  1 -
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/src/memoize.js b/src/memoize.js
index 14fc8e2d..b550bdb8 100644
--- a/src/memoize.js
+++ b/src/memoize.js
@@ -10,16 +10,35 @@ 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);
+  if (__BROWSER__) {
+    return browserMemoize(fn);
+  }
+
+  const wm = new WeakMap();
+  return ctx => {
+    if (wm.has(ctx)) {
+      return ((wm.get(ctx): any): A); // Refinement doesn't seem to work
+    } else {
+      const result = fn(ctx);
+      wm.set(ctx, result);
+      return result;
+    }
+  };
+}
+
+/**
+ * There is only ever a single ctx object in the browser.
+ * Therefore we can use a simple memoization function.
+ */
+function browserMemoize<A>(fn: MemoizeFn<A>): MemoizeFn<A> {
+  let memoized;
+  let called = false;
+  return ctx => {
+    if (!called) {
+      memoized = fn(ctx);
+      called = true;
     }
-    const result = fn(ctx);
-    ctx.memoized.set(memoizeKey, result);
-    return result;
+    return memoized;
   };
 }
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<TimingPlugin> = createToken('TimingToken');
 
 function middleware(ctx, next) {
-  ctx.memoized = new Map();
   const {start, render, end, downstream, upstream} = timing.from(ctx);
   ctx.timing = {
     start,

From 4b0b7b8d85964be2411f45353a8b1839c9fea150 Mon Sep 17 00:00:00 2001
From: Ryan Tsao <ryan.j.tsao@gmail.com>
Date: Wed, 19 Sep 2018 10:37:07 -0700
Subject: [PATCH 2/3] Remove vestigial property from memoization test

---
 src/__tests__/memoize.js | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/__tests__/memoize.js b/src/__tests__/memoize.js
index 66cc58c0..3857465e 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(() => {

From 612a16de9add5bd4589cdf06a3835689931d5554 Mon Sep 17 00:00:00 2001
From: Ryan Tsao <ryan.j.tsao@gmail.com>
Date: Wed, 19 Sep 2018 16:42:37 -0700
Subject: [PATCH 3/3] Remove browser-specific memoization

---
 src/__tests__/memoize.js | 10 ++++++++++
 src/memoize.js           | 29 ++++-------------------------
 2 files changed, 14 insertions(+), 25 deletions(-)

diff --git a/src/__tests__/memoize.js b/src/__tests__/memoize.js
index 3857465e..cb9636c4 100644
--- a/src/__tests__/memoize.js
+++ b/src/__tests__/memoize.js
@@ -30,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 b550bdb8..61721646 100644
--- a/src/memoize.js
+++ b/src/memoize.js
@@ -11,34 +11,13 @@ import type {Context} from './types.js';
 type MemoizeFn<A> = (ctx: Context) => A;
 
 export function memoize<A>(fn: MemoizeFn<A>): MemoizeFn<A> {
-  if (__BROWSER__) {
-    return browserMemoize(fn);
-  }
-
   const wm = new WeakMap();
   return ctx => {
     if (wm.has(ctx)) {
-      return ((wm.get(ctx): any): A); // Refinement doesn't seem to work
-    } else {
-      const result = fn(ctx);
-      wm.set(ctx, result);
-      return result;
-    }
-  };
-}
-
-/**
- * There is only ever a single ctx object in the browser.
- * Therefore we can use a simple memoization function.
- */
-function browserMemoize<A>(fn: MemoizeFn<A>): MemoizeFn<A> {
-  let memoized;
-  let called = false;
-  return ctx => {
-    if (!called) {
-      memoized = fn(ctx);
-      called = true;
+      return ((wm.get(ctx): any): A); // Refinement with `has` doesn't seem to work
     }
-    return memoized;
+    const result = fn(ctx);
+    wm.set(ctx, result);
+    return result;
   };
 }