support nested memoize (2nd) (#5)

* support nested memoize (2nd) * we need to touch affected for the first time too * only copy affected if necessary * fix tracking object to memo * update README
dai-shi · Nov 1, 2020 · a66cbc2 · a66cbc2
1 parent fec8831
commit a66cbc2
Show file tree

Hide file tree

Showing 2 changed files with 88 additions and 17 deletions.
diff --git a/README.md b/README.md
@@ -72,18 +72,21 @@ Instead of [reselect](https://github.com/reduxjs/reselect).
 ```js
 import { useSelector } from 'react-redux';
 
-const selectTotal = memoize(({ state, id }) => ({
-  total: state.a + state.b,
-  title: state.titles[id]
-}))
+const getScore = memoize(state => ({
+  score: heavyComputation(state.a + state.b),
+  createdAt: Date.now(),
+}));
 
 const Component = ({ id }) => {
-  const { total, title } = useSelector(state => selectTotal({ state, id }));
-  return <div>{total} {title}</div>;
+  const { score, title } = useSelector(useCallback(memoize(state => ({
+    score: getScore(state),
+    title: state.titles[id],
+  })), [id]));
+  return <div>{score.score} {score.createdAt} {title}</div>;
 };
 ```
 
-[CodeSandbox](https://codesandbox.io/s/proxy-memoize-demo-c1021)
+- [CodeSandbox 1](https://codesandbox.io/s/proxy-memoize-demo-c1021)
 
 ## Usage with Zustand
 

diff --git a/src/index.ts b/src/index.ts
@@ -2,11 +2,32 @@ import {
   createDeepProxy,
   isDeepChanged,
   getUntrackedObject,
-  trackMemo,
 } from 'proxy-compare';
 
+type Affected = WeakMap<object, Set<string | number | symbol>>;
+
+const isObject = (x: unknown): x is object => typeof x === 'object' && x !== null;
+
+/*
+const affectedToPathList = (rootObj: object, affected: Affected) => {
+  const list: (string | number | symbol)[][] = [];
+  const walk = (obj: object, path?: (string | number | symbol)[]) => {
+    const used = affected.get(obj);
+    if (used) {
+      used.forEach((key) => {
+        walk(obj[key as keyof typeof obj], path ? [...path, key] : [key]);
+      });
+    } else if (path) {
+      list.push(path);
+    }
+  };
+  walk(rootObj);
+  return list;
+};
+*/
+
 const untrack = <T>(x: T, seen: Set<T>): T => {
-  if (typeof x !== 'object' || x === null) return x;
+  if (!isObject(x)) return x;
   const untrackedObj = getUntrackedObject(x);
   if (untrackedObj !== null) return untrackedObj;
   if (!seen.has(x)) {
@@ -25,6 +46,33 @@ const getDeepUntrackedObject = <Obj extends object>(obj: Obj): Obj => {
   return getDeepUntrackedObject(untrackedObj);
 };
 
+const copyAffected = (orig: unknown, x: unknown, affected: Affected) => {
+  if (!isObject(orig) || !isObject(x)) return;
+  const used = affected.get(x);
+  if (!used) return;
+  affected.set(orig, used);
+  used.forEach((key) => {
+    copyAffected(
+      orig[key as keyof typeof orig],
+      x[key as keyof typeof x],
+      affected,
+    );
+  });
+};
+
+const touchAffected = (x: unknown, orig: unknown, affected: Affected) => {
+  if (!isObject(x) || !isObject(orig)) return;
+  const used = affected.get(orig);
+  if (!used) return;
+  used.forEach((key) => {
+    touchAffected(
+      x[key as keyof typeof x],
+      orig[key as keyof typeof orig],
+      affected,
+    );
+  });
+};
+
 // properties
 const OBJ_PROPERTY = 'o';
 const RESULT_PROPERTY = 'r';
@@ -46,31 +94,51 @@ const memoize = <Obj extends object, Result>(
   const memoList: {
     [OBJ_PROPERTY]: Obj;
     [RESULT_PROPERTY]: Result;
-    [AFFECTED_PROPERTY]: WeakMap<object, unknown>;
+    [AFFECTED_PROPERTY]: Affected;
   }[] = [];
-  const resultCache = new WeakMap<Obj, Result>();
+  const resultCache = new WeakMap<Obj, {
+    [RESULT_PROPERTY]: Result;
+    [AFFECTED_PROPERTY]: Affected;
+  }>();
   const proxyCache = new WeakMap();
   const memoizedFn = (obj: Obj) => {
     const cacheKey = getDeepUntrackedObject(obj);
-    if (resultCache.has(cacheKey)) return resultCache.get(cacheKey) as Result;
+    const cache = resultCache.get(cacheKey);
+    if (cache) {
+      touchAffected(obj, cacheKey, cache[AFFECTED_PROPERTY]);
+      return cache[RESULT_PROPERTY];
+    }
     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(cacheKey, memo[RESULT_PROPERTY]);
+        resultCache.set(cacheKey, {
+          [RESULT_PROPERTY]: memo[RESULT_PROPERTY],
+          [AFFECTED_PROPERTY]: memo[AFFECTED_PROPERTY],
+        });
+        touchAffected(obj, cacheKey, memo[AFFECTED_PROPERTY]);
         return memo[RESULT_PROPERTY];
       }
     }
-    trackMemo(obj);
-    const affected = new WeakMap<object, unknown>();
+    const affected: Affected = new WeakMap();
     const proxy = createDeepProxy(obj, affected, proxyCache);
     const result = untrack(fn(proxy), new Set());
+    const origObj = getUntrackedObject(obj);
+    if (obj !== cacheKey) {
+      if (cacheKey !== origObj) {
+        copyAffected(cacheKey, origObj, affected);
+      }
+      touchAffected(obj, cacheKey, affected);
+    }
     memoList.unshift({
-      [OBJ_PROPERTY]: obj,
+      [OBJ_PROPERTY]: origObj || obj,
       [RESULT_PROPERTY]: result,
       [AFFECTED_PROPERTY]: affected,
     });
     if (memoList.length > size) memoList.pop();
-    resultCache.set(cacheKey, result);
+    resultCache.set(cacheKey, {
+      [RESULT_PROPERTY]: result,
+      [AFFECTED_PROPERTY]: affected,
+    });
     return result;
   };
   return memoizedFn;