fix(vanilla, react): use experimental_use and some refactors

src/react.ts

@@ -1,7 +1,17 @@
-import { useCallback, useDebugValue, useEffect, useMemo, useRef } from 'react'
 import {
-  affectedToPathList,
+  // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+  // @ts-ignore
+  experimental_use,
+  useCallback,
+  useDebugValue,
+  useEffect,
+  useMemo,
+  useRef,
+} from 'react'
+import {
+  // affectedToPathList,
   createProxy as createProxyToCompare,
+  getUntracked,
   isChanged,
 } from 'proxy-compare'
 // import { useSyncExternalStore } from 'use-sync-external-store/shim'
@@ -14,6 +24,38 @@ import type { INTERNAL_Snapshot } from './vanilla'
 
 const { useSyncExternalStore } = useSyncExternalStoreExports
 
+// customized version of affectedToPathList
+// we need to avoid invoking getters
+const affectedToPathList = (
+  obj: unknown,
+  affected: WeakMap<object, unknown>
+) => {
+  const list: (string | symbol)[][] = []
+  const seen = new WeakSet()
+  const walk = (x: unknown, path?: (string | symbol)[]) => {
+    if (seen.has(x as object)) {
+      // for object with cycles
+      return
+    }
+    let used: Set<string | symbol> | undefined
+    if (typeof x === 'object' && x !== null) {
+      seen.add(x)
+      used = affected.get(getUntracked(x) || x) as any
+    }
+    if (used) {
+      used.forEach((key) => {
+        if ('value' in (Object.getOwnPropertyDescriptor(x, key) || {})) {
+          walk((x as any)[key], path ? [...path, key] : [key])
+        }
+      })
+    } else if (path) {
+      list.push(path)
+    }
+  }
+  walk(obj)
+  return list
+}
+
 const useAffectedDebugValue = (
   state: object,
   affected: WeakMap<object, unknown>
@@ -119,7 +161,7 @@ export function useSnapshot<T extends object>(
       [proxyObject, notifyInSync]
     ),
     () => {
-      const nextSnapshot = snapshot(proxyObject)
+      const nextSnapshot = snapshot(proxyObject, experimental_use)
       try {
         if (
           !inRender &&
@@ -140,7 +182,7 @@ export function useSnapshot<T extends object>(
       }
       return nextSnapshot
     },
-    () => snapshot(proxyObject)
+    () => snapshot(proxyObject, experimental_use)
   )
   inRender = false
   const currAffected = new WeakMap()

src/vanilla.ts

@@ -31,11 +31,24 @@ export type INTERNAL_Snapshot<T> = T extends AnyFunction
       readonly [K in keyof T]: INTERNAL_Snapshot<T[K]>
     }
 
+type CreateSnapshot = <T extends object>(
+  target: T,
+  receiver: object,
+  version: number,
+  use?: <V>(promise: Promise<V>) => V
+) => T
+
+type ProxyState = [
+  target: object,
+  receiver: object,
+  version: number,
+  createSnapshot: CreateSnapshot,
+  listeners: Set<Listener>
+]
+
 // shared state
+const PROXY_STATE = __DEV__ ? Symbol('PROXY_STATE') : Symbol()
 const refSet = new WeakSet()
-const VERSION = __DEV__ ? Symbol('VERSION') : Symbol()
-const LISTENERS = __DEV__ ? Symbol('LISTENERS') : Symbol()
-const SNAPSHOT = __DEV__ ? Symbol('SNAPSHOT') : Symbol()
 
 const buildProxyFunction = (
   objectIs = Object.is,
@@ -56,51 +69,60 @@ const buildProxyFunction = (
     !(x instanceof RegExp) &&
     !(x instanceof ArrayBuffer),
 
-  PROMISE_RESULT = __DEV__ ? Symbol('PROMISE_RESULT') : Symbol(),
-  PROMISE_ERROR = __DEV__ ? Symbol('PROMISE_ERROR') : Symbol(),
+  defaultUse = (() => {
+    type PromiseState = { v?: unknown; e?: unknown }
+    const PROMISE_STATE = Symbol()
+    return <V>(promise: Promise<V>): V => {
+      let state: PromiseState | undefined = (promise as any)[PROMISE_STATE]
+      if (!state) {
+        state = {}
+        ;(promise as any)[PROMISE_STATE] = state
+        promise
+          .then((v) => ((state as PromiseState).v = v))
+          .catch((e) => ((state as PromiseState).e = e))
+      }
+      if ('v' in state) {
+        return state.v as V
+      }
+      throw 'e' in state ? state.e : promise
+    }
+  })(),
 
-  snapshotCache = new WeakMap<object, [version: number, snapshot: unknown]>(),
+  snapCache = new WeakMap<object, [version: number, snap: unknown]>(),
 
-  createSnapshot = <T extends object>(
-    version: number,
+  createSnapshot: CreateSnapshot = <T extends object>(
     target: T,
-    receiver: any
+    receiver: object,
+    version: number,
+    use = defaultUse
   ): T => {
-    const cache = snapshotCache.get(receiver)
+    const cache = snapCache.get(receiver)
     if (cache?.[0] === version) {
       return cache[1] as T
     }
-    const snapshot: any = Array.isArray(target)
+    const snap: any = Array.isArray(target)
       ? []
       : Object.create(Object.getPrototypeOf(target))
-    markToTrack(snapshot, true) // mark to track
-    snapshotCache.set(receiver, [version, snapshot])
+    markToTrack(snap, true) // mark to track
+    snapCache.set(receiver, [version, snap])
     Reflect.ownKeys(target).forEach((key) => {
       const value = Reflect.get(target, key, receiver)
       if (refSet.has(value)) {
         markToTrack(value, false) // mark not to track
-        snapshot[key] = value
+        snap[key] = value
       } else if (value instanceof Promise) {
-        if (PROMISE_RESULT in value) {
-          snapshot[key] = (value as any)[PROMISE_RESULT]
-        } else {
-          const errorOrPromise = (value as any)[PROMISE_ERROR] || value
-          Object.defineProperty(snapshot, key, {
-            get() {
-              if (PROMISE_RESULT in value) {
-                return (value as any)[PROMISE_RESULT]
-              }
-              throw errorOrPromise
-            },
-          })
-        }
-      } else if (value?.[LISTENERS]) {
-        snapshot[key] = value[SNAPSHOT]
+        Object.defineProperty(snap, key, {
+          get() {
+            return use(value)
+          },
+        })
+      } else if (value?.[PROXY_STATE]) {
+        snap[key] = snapshot(value, use)
       } else {
-        snapshot[key] = value
+        snap[key] = value
       }
     })
-    return Object.freeze(snapshot)
+    return Object.freeze(snap)
   },
 
   proxyCache = new WeakMap<object, ProxyObject>(),
@@ -145,65 +167,64 @@ const buildProxyFunction = (
       ? []
       : Object.create(Object.getPrototypeOf(initialObject))
     const handler: ProxyHandler<T> = {
-      get(target: T, prop: string | symbol, receiver: any) {
-        if (prop === VERSION) {
-          return version
-        }
-        if (prop === LISTENERS) {
-          return listeners
-        }
-        if (prop === SNAPSHOT) {
-          return createSnapshot(version, target, receiver)
+      get(target: T, prop: string | symbol, receiver: object) {
+        if (prop === PROXY_STATE) {
+          const state: ProxyState = [
+            target,
+            receiver,
+            version,
+            createSnapshot,
+            listeners,
+          ]
+          return state
         }
         return Reflect.get(target, prop, receiver)
       },
       deleteProperty(target: T, prop: string | symbol) {
         const prevValue = Reflect.get(target, prop)
-        const childListeners = prevValue?.[LISTENERS]
+        const childListeners = (
+          (prevValue as any)?.[PROXY_STATE] as ProxyState | undefined
+        )?.[4]
         if (childListeners) {
-          childListeners.delete(popPropListener(prop))
+          childListeners.delete(popPropListener(prop) as Listener)
         }
         const deleted = Reflect.deleteProperty(target, prop)
         if (deleted) {
           notifyUpdate(['delete', [prop], prevValue])
         }
         return deleted
       },
-      set(target: T, prop: string | symbol, value: any, receiver: any) {
+      set(target: T, prop: string | symbol, value: any, receiver: object) {
         const hasPrevValue = Reflect.has(target, prop)
         const prevValue = Reflect.get(target, prop, receiver)
         if (hasPrevValue && objectIs(prevValue, value)) {
           return true
         }
-        const childListeners = prevValue?.[LISTENERS]
+        const childListeners = (
+          (prevValue as any)?.[PROXY_STATE] as ProxyState | undefined
+        )?.[4]
         if (childListeners) {
-          childListeners.delete(popPropListener(prop))
+          childListeners.delete(popPropListener(prop) as Listener)
         }
         if (isObject(value)) {
           value = getUntracked(value) || value
         }
-        let nextValue: any
+        let nextValue = value
         if (Object.getOwnPropertyDescriptor(target, prop)?.set) {
-          nextValue = value
+          // do nothing
         } else if (value instanceof Promise) {
-          nextValue = value
-            .then((v) => {
-              nextValue[PROMISE_RESULT] = v
-              notifyUpdate(['resolve', [prop], v])
-              return v
-            })
-            .catch((e) => {
-              nextValue[PROMISE_ERROR] = e
-              notifyUpdate(['reject', [prop], e])
-            })
-        } else if (value?.[LISTENERS]) {
-          nextValue = value
-          nextValue[LISTENERS].add(getPropListener(prop))
-        } else if (canProxy(value)) {
-          nextValue = proxy(value)
-          nextValue[LISTENERS].add(getPropListener(prop))
+          value
+            .then((v) => notifyUpdate(['resolve', [prop], v]))
+            .catch((e) => notifyUpdate(['reject', [prop], e]))
         } else {
-          nextValue = value
+          if (!value?.[PROXY_STATE] && canProxy(value)) {
+            nextValue = proxy(value)
+          }
+          if (nextValue?.[PROXY_STATE]) {
+            ;(nextValue[PROXY_STATE] as ProxyState)[4].add(
+              getPropListener(prop)
+            )
+          }
         }
         Reflect.set(target, prop, nextValue, receiver)
         notifyUpdate(['set', [prop], value, prevValue])
@@ -230,17 +251,14 @@ const buildProxyFunction = (
     // public functions
     proxyFunction,
     // shared state
+    PROXY_STATE,
     refSet,
-    VERSION,
-    LISTENERS,
-    SNAPSHOT,
     // internal things
     objectIs,
     newProxy,
     canProxy,
-    PROMISE_RESULT,
-    PROMISE_ERROR,
-    snapshotCache,
+    defaultUse,
+    snapCache,
     createSnapshot,
     proxyCache,
     versionHolder,
@@ -253,15 +271,16 @@ export function proxy<T extends object>(initialObject: T = {} as T): T {
 }
 
 export function getVersion(proxyObject: unknown): number | undefined {
-  return isObject(proxyObject) ? (proxyObject as any)[VERSION] : undefined
+  const state: ProxyState | undefined = (proxyObject as any)?.[PROXY_STATE]
+  return state?.[2]
 }
 
 export function subscribe<T extends object>(
   proxyObject: T,
   callback: (ops: Op[]) => void,
   notifyInSync?: boolean
 ) {
-  if (__DEV__ && !(proxyObject as any)?.[LISTENERS]) {
+  if (__DEV__ && !(proxyObject as any)?.[PROXY_STATE]) {
     console.warn('Please use proxy object')
   }
   let promise: Promise<void> | undefined
@@ -279,19 +298,22 @@ export function subscribe<T extends object>(
       })
     }
   }
-  ;(proxyObject as any)[LISTENERS].add(listener)
-  return () => {
-    ;(proxyObject as any)[LISTENERS].delete(listener)
-  }
+  const listeners = ((proxyObject as any)[PROXY_STATE] as ProxyState)[4]
+  listeners.add(listener)
+  return () => listeners.delete(listener)
 }
 
 export function snapshot<T extends object>(
-  proxyObject: T
+  proxyObject: T,
+  use?: <V>(promise: Promise<V>) => V
 ): INTERNAL_Snapshot<T> {
-  if (__DEV__ && !(proxyObject as any)?.[SNAPSHOT]) {
+  if (__DEV__ && !(proxyObject as any)?.[PROXY_STATE]) {
     console.warn('Please use proxy object')
   }
-  return (proxyObject as any)[SNAPSHOT]
+  const [target, receiver, version, createSnapshot] = (proxyObject as any)[
+    PROXY_STATE
+  ] as ProxyState
+  return createSnapshot(target, receiver, version, use) as INTERNAL_Snapshot<T>
 }
 
 export function ref<T extends object>(obj: T): T & AsRef {