From 4751f8ebad06baa79bfb15b538aceca140616bda Mon Sep 17 00:00:00 2001
From: Daishi Kato <dai-shi@users.noreply.github.com>
Date: Sat, 31 Oct 2020 08:13:14 +0900
Subject: [PATCH] support nested memoize (#4)

* support nested memoize

* use deep untracked object as cache key
---
 src/index.ts | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index c073515..1d1da3b 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 => {
@@ -18,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';
@@ -44,14 +51,16 @@ 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];
       }
     }
+    trackMemo(obj);
     const affected = new WeakMap<object, unknown>();
     const proxy = createDeepProxy(obj, affected, proxyCache);
     const result = untrack(fn(proxy), new Set());
@@ -61,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;