From 320fb1043dd1825dd77fd6b27293151151cae57c Mon Sep 17 00:00:00 2001
From: daishi <daishi@axlight.com>
Date: Thu, 29 Oct 2020 10:22:28 +0900
Subject: [PATCH 1/2] support nested memoize

---
 src/index.ts | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/index.ts b/src/index.ts
index c073515..3733e85 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -2,6 +2,7 @@ import {
   createDeepProxy,
   isDeepChanged,
   getUntrackedObject,
+  trackMemo,
 } from 'proxy-compare';
 
 const untrack = <T>(x: T, seen: Set<T>): T => {
@@ -52,6 +53,7 @@ const memoize = <Obj extends object, Result>(
         return memo[RESULT_PROPERTY];
       }
     }
+    trackMemo(obj);
     const affected = new WeakMap<object, unknown>();
     const proxy = createDeepProxy(obj, affected, proxyCache);
     const result = untrack(fn(proxy), new Set());

From 3f58857f0a5c8446697b60e1be6d2b0b68eee6a2 Mon Sep 17 00:00:00 2001
From: daishi <daishi@axlight.com>
Date: Thu, 29 Oct 2020 14:20:57 +0900
Subject: [PATCH 2/2] use deep untracked object as cache key

---
 src/index.ts | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 3733e85..1d1da3b 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -19,6 +19,12 @@ const untrack = <T>(x: T, seen: Set<T>): T => {
   return x;
 };
 
+const getDeepUntrackedObject = <Obj extends object>(obj: Obj): Obj => {
+  const untrackedObj = getUntrackedObject(obj);
+  if (untrackedObj === null) return obj;
+  return getDeepUntrackedObject(untrackedObj);
+};
+
 // properties
 const OBJ_PROPERTY = 'o';
 const RESULT_PROPERTY = 'r';
@@ -45,11 +51,12 @@ const memoize = <Obj extends object, Result>(
   const resultCache = new WeakMap<Obj, Result>();
   const proxyCache = new WeakMap();
   const memoizedFn = (obj: Obj) => {
-    if (resultCache.has(obj)) return resultCache.get(obj) as Result;
+    const cacheKey = getDeepUntrackedObject(obj);
+    if (resultCache.has(cacheKey)) return resultCache.get(cacheKey) as Result;
     for (let i = 0; i < memoList.length; i += 1) {
       const memo = memoList[i];
       if (!isDeepChanged(memo[OBJ_PROPERTY], obj, memo[AFFECTED_PROPERTY], proxyCache)) {
-        resultCache.set(obj, memo[RESULT_PROPERTY]);
+        resultCache.set(cacheKey, memo[RESULT_PROPERTY]);
         return memo[RESULT_PROPERTY];
       }
     }
@@ -63,7 +70,7 @@ const memoize = <Obj extends object, Result>(
       [AFFECTED_PROPERTY]: affected,
     });
     if (memoList.length > size) memoList.pop();
-    resultCache.set(obj, result);
+    resultCache.set(cacheKey, result);
     return result;
   };
   return memoizedFn;