@tanstack/redact 0.0.21 → 0.1.0

package/src/dom/features/suspense/full.ts

@@ -1,5 +1,5 @@
 import { FiberTag, createFiber, type Fiber } from '../../../core'
-import { REACT_SUSPENSE_TYPE } from '../../../react'
+import { REACT_SUSPENSE_TYPE, ReactSharedInternals, startTransition } from '../../../react'
 import {
   registerRenderer,
   registerTypeMatcher,
@@ -15,7 +15,19 @@ import {
   getCurrentRoot,
   withCurrentRoot,
   discardPendingWork,
+  findSuspenseBoundary,
+  RenderSuspenseCapture,
+  activityIsDisconnected,
+  isActivityHidden,
+  attachRef,
+  detachRef,
+  enqueueEffect,
+  flushFiberCommits,
+  scheduleLifecycle,
+  handleCommitError,
 } from '../../reconcile'
+import { cleanupEffect } from '../../dispatcher'
+import { retainedEffects, inSuspensePrimary, setLayoutDisconnected, walkRetained } from '../../retained-effects'
 import {
   HydrationCursor,
   setHydrationCursor,
@@ -24,15 +36,25 @@ import {
   tryConsumeBoundary,
   isHydrationBailout,
 } from '../hydration'
+import { hasContextChanged, renderContextConsumers } from '../context'
+import { checkpointCommit, rewindCommit, queueMutation, onCommitFailure } from '../../commit'
 
 let suspendHandler: ((t: Promise<any>) => void) | null = null
 
+function scheduleRetry(fiber: Fiber): void {
+  const previous = ReactSharedInternals.T
+  ReactSharedInternals.T = null
+  try { startTransition(() => scheduleUpdate(fiber)) } finally { ReactSharedInternals.T = previous }
+}
+
 function realHandleSuspended(fiber: Fiber, thenable: Promise<any>): void {
   if (suspendHandler) return suspendHandler(thenable)
+  const boundary = findSuspenseBoundary(fiber)
+  if (boundary) throw new RenderSuspenseCapture(boundary)
   // Fallback: schedule re-render when promise settles
   thenable.then(
-    () => scheduleUpdate(fiber),
-    () => scheduleUpdate(fiber),
+    () => scheduleRetry(fiber),
+    () => scheduleRetry(fiber),
   )
 }
 
@@ -46,14 +68,15 @@ function realHandleSuspended(fiber: Fiber, thenable: Promise<any>): void {
 // their original `display` so we can restore it) and the fallback as a
 // detached Fragment fiber in `state.f` (deliberately kept OUT of
 // `fiber.child` so reconciles against `props.children` don't trip on it).
-// First-mount suspensions have no committed DOM worth preserving, so they
-// keep the original unmount-and-render-fallback behavior.
+// First-mount suspensions discard the failed primary but retain a separate
+// fallback fiber so unsuccessful retries preserve its state and DOM.
 interface SuspenseState {
   p?: Promise<any> | null
   b?: Comment
   e?: Comment
   r?: any
   a?: boolean
+  u?: boolean
   // Re-suspend preservation:
   d?: Array<[HTMLElement, string]> | null
   f?: Fiber | null
@@ -61,11 +84,12 @@ interface SuspenseState {
 
 function renderSuspense(fiber: Fiber, domParent: Node, anchor: Node | null): void {
   const props = fiber.pp ?? {}
-  const state = (fiber.ms ??= {}) as SuspenseState
+  let state = (fiber.ms ??= {}) as SuspenseState
 
   // Streaming hydration: if the next DOM node is a server-emitted boundary
   // marker, route through the boundary-aware hydration path.
   const root = getCurrentRoot()
+  if (root) root.sp = true
   if (root?.h && !state.b) {
     const boundary = tryConsumeBoundary(fiber.parent!)
     if (boundary) {
@@ -85,14 +109,11 @@ function renderSuspense(fiber: Fiber, domParent: Node, anchor: Node | null): voi
     return
   }
 
-  // We were already in the suspended-with-preserved-primary state. Don't
-  // re-attempt primary children (would re-throw and churn the tree). Just
-  // refresh the fallback in case its JSX changed, and wait for the pending
-  // promise to fire scheduleUpdate.
-  if (state.p) {
+  // New primary children or context can unblock a boundary before its old
+  // promise settles. A fallback-only update can keep the pending primary.
+  if (state.p && !state.u && props.children === fiber.mp?.children && !hasContextChanged(fiber, false)) {
     if (state.f) {
-      state.f.pp = { children: props.fallback }
-      renderFiber(state.f, domParent, anchor)
+      renderFallback(state.f, props.fallback, domParent, anchor)
     } else {
       // Initial-mount suspended path: no committed primary to preserve.
       reconcileChildren(fiber, childrenToArray(props.fallback), domParent, anchor)
@@ -105,8 +126,11 @@ function renderSuspense(fiber: Fiber, domParent: Node, anchor: Node | null): voi
   // attempting the new render. If the new attempt suspends and we did have a
   // committed primary, we keep it (hidden) rather than destroying it.
   const hadCommittedPrimary = fiber.mp && fiber.child
+  state.u = false
+  const checkpoint = checkpointCommit(fiber)
 
   const prevHandler = suspendHandler
+  fiber.cx?.clear()
   let pendingThenable: any
   suspendHandler = (thenable) => {
     pendingThenable = thenable
@@ -118,63 +142,142 @@ function renderSuspense(fiber: Fiber, domParent: Node, anchor: Node | null): voi
   }
 
   if (pendingThenable) {
+    rewindCommit(checkpoint)
+    state = (fiber.ms ??= {}) as SuspenseState
+    const previousThenable = state.p
     state.p = pendingThenable
     const onSettle = () => {
-      state.p = null
-      scheduleUpdate(fiber)
+      const current = fiber.ms as SuspenseState | null
+      if (!current || current.p !== pendingThenable) return
+      current.p = null
+      scheduleRetry(fiber)
     }
-    pendingThenable.then(onSettle, onSettle)
+    if (previousThenable !== pendingThenable) pendingThenable.then(onSettle, onSettle)
 
-    if (hadCommittedPrimary) {
+    // Rewind restored the committed tree, including its existing hidden state.
+    // Only the first hide needs update-target, lifecycle and DOM work.
+    if (hadCommittedPrimary && !state.d) {
       // Hide the primary subtree's root host doms so the fallback is the only
       // thing visible, but the underlying nodes (and their scroll/state/focus)
       // survive. Save original `display` for the resume path.
       const hidden: Array<[HTMLElement, string]> = []
       let c: Fiber | null = hadCommittedPrimary
+      while (c) {
+        setUpdateBoundary(c, fiber)
+        walkRetained(c, disconnectLayout)
+        c = c.sibling
+      }
+      c = hadCommittedPrimary
       while (c) {
         hideRootHostDoms(c, hidden)
         c = c.sibling
       }
       state.d = hidden
 
-      // Mount fallback in a detached Fragment fiber. Kept off `fiber.child`
-      // so reconciles of primary don't see it as a stale match candidate.
-      if (!state.f) {
-        state.f = createFiber(FiberTag.Fragment, null, null)
-        state.f.parent = fiber
-      }
-      state.f.pp = { children: props.fallback }
-      renderFiber(state.f, domParent, anchor)
-    } else {
+    } else if (!hadCommittedPrimary) {
       // First-mount suspension — nothing to preserve.
-      unmountAllChildren(fiber)
-      reconcileChildren(fiber, childrenToArray(props.fallback), domParent, anchor)
+      unmountAllChildren(fiber, domParent)
+    }
+    // Keep fallback identity across unsuccessful retries of an initial mount.
+    if (!state.f) {
+      state.f = createFiber(FiberTag.Fragment, null, null)
+      state.f.parent = fiber
     }
+    renderFallback(state.f, props.fallback, domParent, anchor)
   } else {
+    state.p = null
+    fiber.cx?.clear()
+    const wasHidden = !!state.d
     // Render succeeded. Clean up any preserved-suspend state from a prior
     // suspension cycle: unhide primary, unmount the orphan fallback fiber.
     if (state.d) {
-      for (const [el, origDisplay] of state.d) {
-        el.style.display = origDisplay
-      }
+      for (let child = fiber.child; child; child = child.sibling) setUpdateBoundary(child, fiber.su)
+      const hidden = state.d
+      queueMutation(() => {
+        for (const [el, origDisplay] of hidden) el.style.display = origDisplay
+      })
       state.d = null
     }
     if (state.f) {
-      unmountFiber(state.f)
+      unmountFiber(state.f, domParent)
       state.f = null
     }
+    if (wasHidden) for (let child = fiber.child; child; child = child.sibling) walkRetained(child, reconnectLayout, true)
   }
   fiber.mp = props
 }
 
+function renderFallback(fiber: Fiber, children: any, domParent: Node, anchor: Node | null): void {
+  // Local updates remain independently queued, unless an outer hidden primary
+  // owns them. Context must still cross an otherwise unchanged fallback.
+  if (fiber.mp && fiber.mp.children === children && !fiber.su) renderContextConsumers(fiber)
+  else {
+    fiber.pp = { children }
+    renderFiber(fiber, domParent, anchor)
+  }
+}
+
+function disconnectLayout(fiber: Fiber): void {
+  if (setLayoutDisconnected(fiber, true) || activityIsDisconnected(fiber)) return
+  queueMutation(() => {
+    for (const [hook, effect] of retainedEffects.get(fiber) ?? []) {
+      if (effect.t === 1 && !effect.s) cleanupEffect(hook, fiber)
+    }
+    const cleanup = fiber.rc
+    fiber.rc = null
+    if (cleanup) {
+      const index = fiber.cu?.indexOf(cleanup) ?? -1
+      if (index >= 0) fiber.cu!.splice(index, 1)
+      try { cleanup() } catch (error) { handleCommitError(fiber, error) }
+    }
+    detachRef(fiber.cr ?? fiber.ref, fiber)
+    fiber.cr = null
+    const instance = fiber.sn?._fiber === fiber ? fiber.sn : null
+    if (instance?.componentWillUnmount) {
+      const props = instance.props, state = instance.state, committed = fiber.ms?.c
+      if (committed) { instance.props = committed.props; instance.state = committed.state }
+      try { instance.componentWillUnmount() } catch (error) { handleCommitError(fiber, error) }
+      finally { instance.props = props; instance.state = state }
+    }
+  })
+}
+
+function reconnectLayout(fiber: Fiber): void {
+  if (!fiber.ld || isActivityHidden(fiber) || activityIsDisconnected(fiber)) return
+  if (inSuspensePrimary(fiber)) return
+  setLayoutDisconnected(fiber, false)
+  const instance = fiber.sn?._fiber === fiber ? fiber.sn : null
+  if (fiber.tag === FiberTag.Host || fiber.tag === FiberTag.Resource) attachRef(fiber, fiber.dom)
+  else if (instance || fiber.tag === FiberTag.Fragment) attachRef(fiber, fiber.sn)
+  for (const effect of retainedEffects.get(fiber)?.values() ?? []) {
+    if (effect.t === 1 && !effect.s) enqueueEffect(fiber, effect)
+  }
+  if (instance?.componentDidMount) scheduleLifecycle(fiber, () => instance.componentDidMount())
+  flushFiberCommits(fiber)
+}
+
+function setUpdateBoundary(fiber: Fiber, boundary: Fiber | null | undefined): void {
+  fiber.su = boundary ?? null
+  const primary = fiber.tag === FiberTag.Suspense && fiber.ms?.p ? fiber : boundary
+  for (let child = fiber.child; child; child = child.sibling) setUpdateBoundary(child, primary)
+  if (fiber.tag === FiberTag.Suspense && fiber.ms?.f) setUpdateBoundary(fiber.ms.f, boundary)
+}
+
 // Walk a fiber subtree collecting host/text DOM nodes that sit at the root
 // of the subtree (do not descend through their children — display:none on
 // the root hides the whole element). Used by the hide-on-suspend path.
 function hideRootHostDoms(fiber: Fiber, out: Array<[HTMLElement, string]>): void {
   if (fiber.tag === FiberTag.Host) {
     const el = fiber.dom as HTMLElement
-    out.push([el, el.style.display])
-    el.style.display = 'none'
+    const existing = out.find(([node]) => node === el)
+    const entry: [HTMLElement, string] = existing ?? [el, el.style.display]
+    if (!existing) out.push(entry)
+    const display = fiber.pp?.style?.display
+    queueMutation(() => {
+      // Earlier Activity or host updates in this commit can change display.
+      if (!existing || el.style.display !== 'none' || display === 'none') entry[1] = el.style.display
+      el.style.display = 'none'
+    })
     return
   }
   if (fiber.tag === FiberTag.Portal) return
@@ -188,7 +291,7 @@ function hideRootHostDoms(fiber: Fiber, out: Array<[HTMLElement, string]>): void
 function hydrateSuspenseBoundary(
   fiber: Fiber,
   props: any,
-  boundary: [0 | 1, number, Comment, Comment],
+  boundary: [0 | 1 | 2 | 3, number, Comment, Comment],
   domParent: Node,
   anchor: Node | null,
 ): void {
@@ -200,6 +303,13 @@ function hydrateSuspenseBoundary(
     r: props.children,
   }
 
+  if (pendingBoundary >= 2) {
+    if (pendingBoundary === 3) reportServerFailure(fiber)
+    recoverBoundaryHydration(fiber, props.children, startMark.parentNode!, startMark, endMark, true)
+    fiber.mp = props
+    return
+  }
+
   if (!pendingBoundary) {
     // Real DOM is inline between startMark and endMark. Hydrate into it.
     const parent = startMark.parentNode!
@@ -209,8 +319,10 @@ function hydrateSuspenseBoundary(
       advanceCursorPast(fiber.parent!, endMark)
     } catch (e) {
       if (!isHydrationBailout(e)) {
-        clearBoundaryRange(startMark, endMark)
-        unmountAllChildren(fiber)
+        onCommitFailure(() => {
+          clearBoundaryRange(startMark, endMark)
+          unmountAllChildren(fiber, parent)
+        })
         throw e
       }
       recoverBoundaryHydration(fiber, props.children, parent, startMark, endMark)
@@ -253,19 +365,25 @@ function recoverBoundaryHydration(
   parent: Node,
   startMark: Comment,
   endMark: Comment,
+  retryBoundary = false,
 ): void {
   const root = findRoot(fiber)!
   const prevHydrating = root.h
-  discardPendingWork(root)
+  if (!retryBoundary) discardPendingWork(root)
   clearBoundaryRange(startMark, endMark)
-  unmountAllChildren(fiber)
+  unmountAllChildren(fiber, parent)
   root.h = false
   try {
-    reconcileChildren(fiber, childrenToArray(children), parent, endMark)
+    if (retryBoundary) {
+      fiber.ms = {}
+      renderSuspense(fiber, parent, endMark)
+    } else reconcileChildren(fiber, childrenToArray(children), parent, endMark)
     advanceCursorPast(fiber.parent!, endMark)
   } catch (clientError) {
-    clearBoundaryRange(startMark, endMark)
-    unmountAllChildren(fiber)
+    onCommitFailure(() => {
+      clearBoundaryRange(startMark, endMark)
+      unmountAllChildren(fiber, parent)
+    })
     throw clientError
   } finally {
     root.h = prevHydrating
@@ -276,14 +394,16 @@ function recoverFallbackHydration(fiber: Fiber, fallback: any, parent: HTMLEleme
   const root = findRoot(fiber)!
   const prevHydrating = root.h
   discardPendingWork(root)
-  parent.textContent = ''
-  unmountAllChildren(fiber)
+  queueMutation(() => { parent.textContent = '' })
+  unmountAllChildren(fiber, parent)
   root.h = false
   try {
     reconcileChildren(fiber, childrenToArray(fallback), parent, null)
   } catch (clientError) {
-    parent.textContent = ''
-    unmountAllChildren(fiber)
+    onCommitFailure(() => {
+      parent.textContent = ''
+      unmountAllChildren(fiber, parent)
+    })
     throw clientError
   } finally {
     root.h = prevHydrating
@@ -302,8 +422,14 @@ function rehydrateBoundary(fiber: Fiber): void {
   // (or at least its container); unmounting here cleans up fibers + fx.
   withCurrentRoot(root, () => {
     const prevHydrating = root.h
+    if (state.b.data.startsWith('$!') || state.b.data.startsWith('$E')) {
+      if (state.b.data.startsWith('$E')) reportServerFailure(fiber)
+      recoverBoundaryHydration(fiber, state.r, parent, state.b, state.e, true)
+      runEffects(root)
+      return
+    }
     try {
-      unmountAllChildren(fiber)
+      unmountAllChildren(fiber, parent)
 
       // Re-hydrate with real children against the now-real DOM range.
       root.h = true
@@ -311,8 +437,10 @@ function rehydrateBoundary(fiber: Fiber): void {
       reconcileChildren(fiber, childrenToArray(state.r), parent, null)
     } catch (e) {
       if (!isHydrationBailout(e)) {
-        clearBoundaryRange(state.b, state.e)
-        unmountAllChildren(fiber)
+        onCommitFailure(() => {
+          clearBoundaryRange(state.b, state.e)
+          unmountAllChildren(fiber, parent)
+        })
         throw e
       }
 
@@ -325,13 +453,22 @@ function rehydrateBoundary(fiber: Fiber): void {
   })
 }
 
+function reportServerFailure(fiber: Fiber): void {
+  const error = new Error('Switched to client rendering because the server render did not complete.')
+  const handler = findRoot(fiber)?.re
+  if (handler) handler(error, { componentStack: null })
+  else console.error(error)
+}
+
 function clearBoundaryRange(startMark: Comment, endMark: Comment): void {
-  let node = startMark.nextSibling
-  while (node && node !== endMark) {
-    const next = node.nextSibling
-    node.parentNode?.removeChild(node)
-    node = next
-  }
+  queueMutation(() => {
+    let node = startMark.nextSibling
+    while (node && node !== endMark) {
+      const next = node.nextSibling
+      node.parentNode?.removeChild(node)
+      node = next
+    }
+  })
 }
 
 registerTypeMatcher((type) => (type === REACT_SUSPENSE_TYPE ? FiberTag.Suspense : null))