refactor(core): avoid onCancel hack (#3019)

* wip: refactor: avoid onCancel hack

* fix typo

* fix useAtomValue

* refactor internals

* refactor agian to keep the behavior

* small refactor

* refactor a bit

* recover a line

* another refactor

Author: dai-shi

src/react/useAtomValue.ts

@@ -1,6 +1,7 @@
 /// <reference types="react/experimental" />
 
 import ReactExports, { useDebugValue, useEffect, useReducer } from 'react'
+import { INTERNAL_registerAbortHandler as registerAbortHandler } from '../vanilla/internals.ts'
 import type { Atom, ExtractAtomValue } from '../vanilla.ts'
 import { useStore } from './Provider.ts'
 
@@ -55,7 +56,10 @@ const continuablePromiseMap = new WeakMap<
   Promise<unknown>
 >()
 
-const createContinuablePromise = <T>(promise: PromiseLike<T>) => {
+const createContinuablePromise = <T>(
+  promise: PromiseLike<T>,
+  getValue: () => PromiseLike<T> | T,
+) => {
   let continuablePromise = continuablePromiseMap.get(promise)
   if (!continuablePromise) {
     continuablePromise = new Promise<T>((resolve, reject) => {
@@ -70,25 +74,23 @@ const createContinuablePromise = <T>(promise: PromiseLike<T>) => {
           reject(e)
         }
       }
-      const registerCancelHandler = (p: PromiseLike<T>) => {
-        if ('onCancel' in p && typeof p.onCancel === 'function') {
-          p.onCancel((nextValue: PromiseLike<T> | T) => {
-            if (import.meta.env?.MODE !== 'production' && nextValue === p) {
-              throw new Error('[Bug] p is not updated even after cancelation')
-            }
-            if (isPromiseLike(nextValue)) {
-              continuablePromiseMap.set(nextValue, continuablePromise!)
-              curr = nextValue
-              nextValue.then(onFulfilled(nextValue), onRejected(nextValue))
-              registerCancelHandler(nextValue)
-            } else {
-              resolve(nextValue)
-            }
-          })
+      const onAbort = () => {
+        try {
+          const nextValue = getValue()
+          if (isPromiseLike(nextValue)) {
+            continuablePromiseMap.set(nextValue, continuablePromise!)
+            curr = nextValue
+            nextValue.then(onFulfilled(nextValue), onRejected(nextValue))
+            registerAbortHandler(nextValue, onAbort)
+          } else {
+            resolve(nextValue)
+          }
+        } catch (e) {
+          reject(e)
         }
       }
       promise.then(onFulfilled(promise), onRejected(promise))
-      registerCancelHandler(promise)
+      registerAbortHandler(promise, onAbort)
     })
     continuablePromiseMap.set(promise, continuablePromise)
   }
@@ -141,7 +143,9 @@ export function useAtomValue<Value>(atom: Atom<Value>, options?: Options) {
       if (typeof delay === 'number') {
         const value = store.get(atom)
         if (isPromiseLike(value)) {
-          attachPromiseMeta(createContinuablePromise(value))
+          attachPromiseMeta(
+            createContinuablePromise(value, () => store.get(atom)),
+          )
         }
         // delay rerendering to wait a promise possibly to resolve
         setTimeout(rerender, delay)
@@ -157,7 +161,7 @@ export function useAtomValue<Value>(atom: Atom<Value>, options?: Options) {
   // The use of isPromiseLike is to be consistent with `use` type.
   // `instanceof Promise` actually works fine in this case.
   if (isPromiseLike(value)) {
-    const promise = createContinuablePromise(value)
+    const promise = createContinuablePromise(value, () => store.get(atom))
     return use(promise)
   }
   return value as Awaited<Value>

src/vanilla/internals.ts

@@ -176,54 +176,42 @@ const returnAtomValue = <Value>(atomState: AtomState<Value>): Value => {
 }
 
 //
-// Cancelable Promise
-// TODO(daishi): revisit this implementation
+// Abortable Promise
 //
 
-type CancelHandler = (nextValue: unknown) => void
-type PromiseState = [cancelHandlers: Set<CancelHandler>, settled: boolean]
-
-const PROMISE_STATE = Symbol()
-
-const getPromiseState = <T>(
-  promise: PromiseLike<T>,
-): PromiseState | undefined => (promise as any)[PROMISE_STATE]
+const promiseStateMap: WeakMap<
+  PromiseLike<unknown>,
+  [pending: boolean, abortHandlers: Set<() => void>]
+> = new WeakMap()
 
 const isPendingPromise = (value: unknown): value is PromiseLike<unknown> =>
-  isPromiseLike(value) && !getPromiseState(value)?.[1]
+  isPromiseLike(value) && !!promiseStateMap.get(value as never)?.[0]
 
-const cancelPromise = <T>(
-  promise: PromiseLike<T>,
-  nextValue: unknown,
-): void => {
-  const promiseState = getPromiseState(promise)
-  if (promiseState) {
-    promiseState[1] = true
-    promiseState[0].forEach((fn) => fn(nextValue))
-  } else if (import.meta.env?.MODE !== 'production') {
-    throw new Error('[Bug] cancelable promise not found')
+const abortPromise = <T>(promise: PromiseLike<T>): void => {
+  const promiseState = promiseStateMap.get(promise)
+  if (promiseState?.[0]) {
+    promiseState[0] = false
+    promiseState[1].forEach((fn) => fn())
   }
 }
 
-const patchPromiseForCancelability = <T>(promise: PromiseLike<T>): void => {
-  if (getPromiseState(promise)) {
-    // already patched
-    return
-  }
-  const promiseState: PromiseState = [new Set(), false]
-  ;(promise as any)[PROMISE_STATE] = promiseState
-  const settle = () => {
-    promiseState[1] = true
-  }
-  promise.then(settle, settle)
-  ;(promise as { onCancel?: (fn: CancelHandler) => void }).onCancel = (fn) => {
-    promiseState[0].add(fn)
+const registerAbortHandler = <T>(
+  promise: PromiseLike<T>,
+  abortHandler: () => void,
+): void => {
+  let promiseState = promiseStateMap.get(promise)
+  if (!promiseState) {
+    promiseState = [true, new Set()]
+    promiseStateMap.set(promise, promiseState)
+    const settle = () => {
+      promiseState![0] = false
+    }
+    promise.then(settle, settle)
   }
+  promiseState[1].add(abortHandler)
 }
 
-const isPromiseLike = (
-  p: unknown,
-): p is PromiseLike<unknown> & { onCancel?: (fn: CancelHandler) => void } =>
+const isPromiseLike = (p: unknown): p is PromiseLike<unknown> =>
   typeof (p as any)?.then === 'function'
 
 const addPendingPromiseToDependency = (
@@ -253,9 +241,7 @@ const setAtomStateValueOrPromise = (
   const atomState = ensureAtomState(atom)
   const hasPrevValue = 'v' in atomState
   const prevValue = atomState.v
-  const pendingPromise = isPendingPromise(atomState.v) ? atomState.v : null
   if (isPromiseLike(valueOrPromise)) {
-    patchPromiseForCancelability(valueOrPromise)
     for (const a of atomState.d.keys()) {
       addPendingPromiseToDependency(atom, valueOrPromise, ensureAtomState(a))
     }
@@ -264,8 +250,8 @@ const setAtomStateValueOrPromise = (
   delete atomState.e
   if (!hasPrevValue || !Object.is(prevValue, atomState.v)) {
     ++atomState.n
-    if (pendingPromise) {
-      cancelPromise(pendingPromise, valueOrPromise)
+    if (isPromiseLike(prevValue)) {
+      abortPromise(prevValue)
     }
   }
 }
@@ -662,7 +648,7 @@ const buildStore = (
         const valueOrPromise = atomRead(atom, getter, options as never)
         setAtomStateValueOrPromise(atom, valueOrPromise, ensureAtomState!)
         if (isPromiseLike(valueOrPromise)) {
-          valueOrPromise.onCancel?.(() => controller?.abort())
+          registerAbortHandler(valueOrPromise, () => controller?.abort())
           valueOrPromise.then(
             mountDependenciesIfAsync,
             mountDependenciesIfAsync,
@@ -931,12 +917,12 @@ export const INTERNAL_isActuallyWritableAtom: typeof isActuallyWritableAtom =
 export const INTERNAL_isAtomStateInitialized: typeof isAtomStateInitialized =
   isAtomStateInitialized
 export const INTERNAL_returnAtomValue: typeof returnAtomValue = returnAtomValue
-export const INTERNAL_getPromiseState: typeof getPromiseState = getPromiseState
+export const INTERNAL_promiseStateMap: typeof promiseStateMap = promiseStateMap
 export const INTERNAL_isPendingPromise: typeof isPendingPromise =
   isPendingPromise
-export const INTERNAL_cancelPromise: typeof cancelPromise = cancelPromise
-export const INTERNAL_patchPromiseForCancelability: typeof patchPromiseForCancelability =
-  patchPromiseForCancelability
+export const INTERNAL_abortPromise: typeof abortPromise = abortPromise
+export const INTERNAL_registerAbortHandler: typeof registerAbortHandler =
+  registerAbortHandler
 export const INTERNAL_isPromiseLike: typeof isPromiseLike = isPromiseLike
 export const INTERNAL_addPendingPromiseToDependency: typeof addPendingPromiseToDependency =
   addPendingPromiseToDependency