@remix-run/ui 0.4.0 → 0.6.0

package/dist/runtime/diff-dom.js

@@ -1,4 +1,6 @@
-import { invariant } from "./invariant.js";
+import { invariant } from './invariant.js';
+import { disposeClientEntryBoundary, getClientEntryBoundaryOwner } from './client-entry-boundary.js';
+const REMIX_PRESERVE_DOM_ATTRIBUTE = 'rmx-preserve-dom';
 export function diffNodes(curr, next, context) {
     let parent = curr[0]?.parentNode ?? context.regionParent ?? null;
     invariant(parent, 'Parent node not found');
@@ -6,60 +8,7 @@ export function diffNodes(curr, next, context) {
     // nodes before the region tail ref rather than appending to the parent.
     let regionTailRef = context.regionTailRef ??
         (curr.length > 0 ? curr[curr.length - 1].nextSibling : null);
-    let currentIndex = 0;
-    let nextIndex = 0;
-    while (currentIndex < curr.length || nextIndex < next.length) {
-        let c = curr[currentIndex];
-        let n = next[nextIndex];
-        if (!c && n) {
-            if (regionTailRef) {
-                parent.insertBefore(n, regionTailRef);
-            }
-            else {
-                parent.appendChild(n);
-            }
-            nextIndex++;
-        }
-        else if (c && !n) {
-            removeNode(c, parent, context);
-            currentIndex++;
-        }
-        else if (c && n) {
-            let replacement = getCommentMarkerRangeReplacement(c, n, curr, next, currentIndex, nextIndex, context);
-            if (replacement) {
-                replaceCommentMarkerRange(replacement, parent, context);
-                currentIndex = replacement.currentEndIndex + 1;
-                nextIndex = replacement.nextEndIndex + 1;
-                continue;
-            }
-            // Skip hydrated client-entry marker ranges; hydration pass re-renders
-            // roots with new props from incoming payload
-            if (isVirtualRootStartMarker(c) && isVirtualRootStartMarker(n)) {
-                let currentEnd = findHydrationEndMarker(c);
-                let nextEnd = findHydrationEndMarker(n);
-                let nextData = n.data;
-                if (c.data !== nextData)
-                    c.data = nextData;
-                let currentEndIndex = curr.indexOf(currentEnd);
-                let nextEndIndex = next.indexOf(nextEnd);
-                currentIndex = currentEndIndex + 1;
-                nextIndex = nextEndIndex + 1;
-                continue;
-            }
-            let cursor = diffNode(c, n, context);
-            if (cursor) {
-                let nextEndIndex = next.indexOf(cursor);
-                let currentEndIndex = isFrameStartMarker(c) && isFrameEndMarker(cursor)
-                    ? findFrameEndIndex(curr, currentIndex)
-                    : currentIndex;
-                currentIndex = currentEndIndex + 1;
-                nextIndex = nextEndIndex + 1;
-                continue;
-            }
-            currentIndex++;
-            nextIndex++;
-        }
-    }
+    diffSiblingUnits(curr, next, parent, regionTailRef, context);
 }
 function diffNode(current, next, context) {
     // Text -> Text
@@ -80,40 +29,41 @@ function diffNode(current, next, context) {
         return end;
     }
     // Comment -> Comment
-    if (isCommentNode(current) && isCommentNode(next)) {
+    if (isCommentNode(current) && isCommentNode(next) && markerKindsMatch(current, next)) {
         let newData = next.data;
-        if (current.data !== newData) {
-            let updated = false;
-            if (isFrameStartMarker(current)) {
-                if (shouldPreserveFrameStartMarker(current, next, context)) {
+        let updated = false;
+        if (isFrameStartMarker(current)) {
+            if (shouldPreserveFrameStartMarker(current, next, context)) {
+                if (current.data !== newData) {
                     current.data = newData;
-                    updated = true;
-                    let frame = context.frameInstances.get(current);
-                    let nextMarkerData = getFrameMarkerData(next, context);
-                    if (frame && nextMarkerData) {
-                        if (nextMarkerData.status === 'resolved') {
-                            let nextEnd = findFrameEndMarker(next);
-                            let nextContent = collectFrameContentFragment(current.ownerDocument, next, nextEnd);
-                            void frame.renderMarkerContent({ ...nextMarkerData, id: getFrameId(next) }, nextContent, {
-                                signal: context.signal,
-                            });
-                            return nextEnd;
-                        }
-                        if (frame.isDisplayingResolvedContent()) {
-                            return findFrameEndMarker(next);
-                        }
-                    }
                 }
-                else {
-                    disposeFrameStartMarker(current, context);
-                    current.data = newData;
-                    updated = true;
+                updated = true;
+                let frame = context.frameInstances.get(current);
+                let nextMarkerData = getFrameMarkerData(next, context);
+                if (frame && nextMarkerData) {
+                    if (nextMarkerData.status === 'resolved') {
+                        let nextEnd = findFrameEndMarker(next);
+                        let nextContent = collectFrameContentFragment(current.ownerDocument, next, nextEnd);
+                        void frame.renderMarkerContent({ ...nextMarkerData, id: getFrameId(next) }, nextContent, {
+                            data: context.data,
+                            signal: context.signal,
+                        });
+                        return nextEnd;
+                    }
+                    if (frame.isDisplayingResolvedContent()) {
+                        return findFrameEndMarker(next);
+                    }
                 }
             }
-            if (!updated) {
+            else if (current.data !== newData) {
+                disposeFrameStartMarker(current, context);
                 current.data = newData;
+                updated = true;
             }
         }
+        if (!updated && current.data !== newData) {
+            current.data = newData;
+        }
         return;
     }
     // Element -> Element
@@ -128,6 +78,8 @@ function diffNode(current, next, context) {
             return;
         }
         // Same tag: update attributes then children
+        if (shouldPreserveDomElement(current, next))
+            return;
         diffElementAttributes(current, next);
         if (shouldPreserveElementChildren(current, next))
             return;
@@ -164,6 +116,14 @@ function diffElementAttributes(current, next) {
         }
     }
 }
+function shouldPreserveDomElement(current, next) {
+    if (!next.hasAttribute(REMIX_PRESERVE_DOM_ATTRIBUTE))
+        return false;
+    if (!current.hasAttribute(REMIX_PRESERVE_DOM_ATTRIBUTE)) {
+        current.setAttribute(REMIX_PRESERVE_DOM_ATTRIBUTE, '');
+    }
+    return true;
+}
 function shouldPreserveLiveAttribute(current, next, name) {
     if (name === 'open') {
         if (current instanceof HTMLDetailsElement && next instanceof HTMLDetailsElement) {
@@ -219,164 +179,253 @@ function isPopoverOpen(element) {
     }
 }
 function diffElementChildren(current, next, context) {
-    let currentChildren = Array.from(current.childNodes);
-    let nextChildren = Array.from(next.childNodes);
-    // Keyed map by data-key for current children
-    let keyToIndex = new Map();
-    for (let i = 0; i < currentChildren.length; i++) {
-        let node = currentChildren[i];
-        if (isElement(node)) {
-            let key = node.getAttribute('data-key');
-            if (key != null)
-                keyToIndex.set(key, i);
+    let currentChildren;
+    // Allow actively managed preload link tags in the head to stay in the document
+    // during diffing rather than removing them which aborts the preload in Safari
+    if (context.isActiveModulePreload && current === current.ownerDocument.head) {
+        currentChildren = [];
+        for (let node of current.childNodes) {
+            if (!context.isActiveModulePreload(node))
+                currentChildren.push(node);
         }
     }
-    let used = new Array(currentChildren.length).fill(false);
-    let matchIndexForNext = new Array(nextChildren.length).fill(-1);
-    for (let i = 0; i < nextChildren.length; i++) {
-        let nextChild = nextChildren[i];
+    else {
+        currentChildren = Array.from(current.childNodes);
+    }
+    let nextChildren = Array.from(next.childNodes);
+    diffSiblingUnits(currentChildren, nextChildren, current, null, context);
+}
+function diffSiblingUnits(currentNodes, nextNodes, parent, regionTailRef, context) {
+    let currentUnits = parseSiblingUnits(currentNodes);
+    let nextUnits = parseSiblingUnits(nextNodes);
+    let keyToIndex = new Map();
+    for (let i = 0; i < currentUnits.length; i++) {
+        let key = getSiblingUnitKey(currentUnits[i]);
+        if (key !== undefined)
+            keyToIndex.set(key, i);
+    }
+    let used = new Array(currentUnits.length).fill(false);
+    let matchIndexForNext = new Array(nextUnits.length).fill(-1);
+    // Reserve globally matched keyed elements and semantic boundaries before
+    // positionally pairing the remaining ordinary siblings.
+    for (let i = 0; i < nextUnits.length; i++) {
+        let nextUnit = nextUnits[i];
         let matchIndex = -1;
-        if (isFrameEndMarker(nextChild)) {
-            for (let j = 0; j < currentChildren.length; j++) {
-                if (!used[j] && isFrameEndMarker(currentChildren[j])) {
+        let key = getSiblingUnitKey(nextUnit);
+        if (key !== undefined) {
+            let keyedIndex = keyToIndex.get(key);
+            if (keyedIndex !== undefined &&
+                !used[keyedIndex] &&
+                siblingUnitsComparable(currentUnits[keyedIndex], nextUnit)) {
+                matchIndex = keyedIndex;
+            }
+        }
+        // Boundary identities behave like implicit keys. Repeated boundaries with
+        // the same identity are paired in source order because server markers do
+        // not carry a separate application key.
+        if (matchIndex === -1 && nextUnit.kind !== 'node') {
+            for (let j = 0; j < currentUnits.length; j++) {
+                let currentUnit = currentUnits[j];
+                if (used[j] || currentUnit.kind !== nextUnit.kind)
+                    continue;
+                if (shouldPreserveBoundaryUnit(currentUnit, nextUnit, context)) {
                     matchIndex = j;
                     break;
                 }
             }
         }
-        else if (isElement(nextChild)) {
-            let key = nextChild.getAttribute('data-key');
-            if (key != null && keyToIndex.has(key)) {
-                let idx = keyToIndex.get(key);
-                if (!used[idx])
-                    matchIndex = idx;
-            }
-        }
-        if (matchIndex === -1) {
-            let candidateIndex = i;
-            if (candidateIndex < currentChildren.length &&
-                !used[candidateIndex] &&
-                nodeTypesComparable(currentChildren[candidateIndex], nextChild)) {
-                matchIndex = candidateIndex;
-            }
-        }
         if (matchIndex !== -1)
             used[matchIndex] = true;
         matchIndexForNext[i] = matchIndex;
     }
-    // Forward pass: update matched, collect committed
-    let committed = new Array(nextChildren.length);
-    for (let i = 0; i < nextChildren.length; i++) {
-        let mi = matchIndexForNext[i];
-        if (mi !== -1) {
-            let curChild = currentChildren[mi];
-            let nextChild = nextChildren[i];
-            let replacement = getCommentMarkerRangeReplacement(curChild, nextChild, currentChildren, nextChildren, mi, i, context);
-            if (replacement) {
-                replaceCommentMarkerRange(replacement, current, context);
-                for (let k = mi; k <= replacement.currentEndIndex; k++)
-                    used[k] = true;
-                committed[i] = replacement.nextStart;
-                committed[replacement.nextEndIndex] = nextChildren[replacement.nextEndIndex];
-                for (let j = i + 1; j < replacement.nextEndIndex; j++)
-                    committed[j] = undefined;
-                i = replacement.nextEndIndex;
-                continue;
-            }
-            let cursor = diffNode(curChild, nextChild, context);
-            if (cursor) {
-                // `diffNode` can preserve hydration and frame marker ranges by returning
-                // the next end marker, so skip the owned interior nodes here.
-                let nextEndIdx = nextChildren.indexOf(cursor);
-                let currEndIdx = isFrameStartMarker(curChild) && isFrameEndMarker(cursor)
-                    ? findFrameEndIndex(currentChildren, mi)
-                    : findHydrationEndIndex(currentChildren, mi);
-                // Adjacent marker ranges can pre-match into the next range and leave an
-                // orphaned end marker behind. Clear those stale matches first.
-                for (let j = i + 1; j <= nextEndIdx; j++) {
-                    let matchedIndex = matchIndexForNext[j];
-                    if (matchedIndex > currEndIdx) {
-                        used[matchedIndex] = false;
-                    }
-                    matchIndexForNext[j] = -1;
-                }
-                // Mark the entire current region as used to avoid removals.
-                for (let k = mi; k <= currEndIdx; k++)
-                    used[k] = true;
-                // Preserve both comment markers in committed; skip interior in reorder pass.
-                committed[i] = curChild; // start marker
-                committed[nextEndIdx] = currentChildren[currEndIdx]; // end marker
-                for (let j = i + 1; j < nextEndIdx; j++)
-                    committed[j] = undefined;
-                // Jump to end of region.
-                i = nextEndIdx;
-                continue;
-            }
-            committed[i] = curChild;
-        }
-        else {
-            committed[i] = nextChildren[i];
+    let remainingCurrentIndexes = [];
+    for (let i = 0; i < currentUnits.length; i++) {
+        if (!used[i])
+            remainingCurrentIndexes.push(i);
+    }
+    let remainingNextIndexes = [];
+    for (let i = 0; i < nextUnits.length; i++) {
+        if (matchIndexForNext[i] === -1)
+            remainingNextIndexes.push(i);
+    }
+    let remainingLength = Math.min(remainingCurrentIndexes.length, remainingNextIndexes.length);
+    for (let i = 0; i < remainingLength; i++) {
+        let currentIndex = remainingCurrentIndexes[i];
+        let nextIndex = remainingNextIndexes[i];
+        let currentUnit = currentUnits[currentIndex];
+        let nextUnit = nextUnits[nextIndex];
+        if (currentUnit.kind !== 'node' ||
+            nextUnit.kind !== 'node' ||
+            getSiblingUnitKey(currentUnit) !== undefined ||
+            getSiblingUnitKey(nextUnit) !== undefined ||
+            !siblingUnitsComparable(currentUnit, nextUnit)) {
+            continue;
         }
+        used[currentIndex] = true;
+        matchIndexForNext[nextIndex] = currentIndex;
     }
-    // Backward pass: reorder via inserts while avoiding redundant moves
-    let anchor = undefined;
-    for (let i = committed.length - 1; i >= 0; i--) {
-        let node = committed[i];
-        if (!node)
+    let committed = new Array(nextUnits.length);
+    for (let i = 0; i < nextUnits.length; i++) {
+        let matchIndex = matchIndexForNext[i];
+        let nextUnit = nextUnits[i];
+        if (matchIndex === -1) {
+            committed[i] = nextUnit;
             continue;
-        // Use only an anchor that is actually a child of the current parent
-        let ref = anchor && anchor.parentNode === current ? anchor : null;
-        // Existing comment markers delimit live ranges, so do not reorder them
-        // independently from their contents. New markers still need to be inserted
-        // before they can be used as anchors.
-        if (isVirtualRootStartMarker(node) ||
-            isVirtualRootEndMarker(node) ||
-            isFrameStartMarker(node) ||
-            isFrameEndMarker(node)) {
-            if (node.parentNode !== current) {
-                current.insertBefore(node, ref);
-            }
-            anchor = node;
+        }
+        let currentUnit = currentUnits[matchIndex];
+        if (currentUnit.kind === 'node' && nextUnit.kind === 'node') {
+            diffNode(currentUnit.node, nextUnit.node, context);
+            committed[i] = currentUnit;
             continue;
         }
-        if (node.parentNode === current) {
-            // Node already in parent: move only if its nextSibling is not the desired ref.
-            let targetNext = ref;
-            let alreadyInPlace = (targetNext === null && node.nextSibling === null) || node.nextSibling === targetNext;
-            if (!alreadyInPlace) {
-                current.insertBefore(node, targetNext);
-            }
+        invariant(currentUnit.kind !== 'node' && nextUnit.kind !== 'node', 'Expected boundaries');
+        let replacement = getCommentMarkerRangeReplacement(currentUnit.start, nextUnit.start, currentNodes, nextNodes, currentUnit.startIndex, nextUnit.startIndex, context);
+        if (replacement) {
+            replaceCommentMarkerRange(replacement, parent, context);
+            committed[i] = nextUnit;
+            continue;
+        }
+        let cursor = diffNode(currentUnit.start, nextUnit.start, context);
+        if (!cursor && currentUnit.kind === 'frame' && nextUnit.kind === 'frame') {
+            diffNodes(collectNodesBetween(currentUnit.start, currentUnit.end), nextNodes.slice(nextUnit.startIndex + 1, nextUnit.endIndex), {
+                ...context,
+                regionParent: parent,
+                regionTailRef: currentUnit.end,
+            });
         }
-        else {
-            // New node: insert relative to a valid ref or append
-            current.insertBefore(node, ref);
+        committed[i] = currentUnit;
+    }
+    let anchor = regionTailRef;
+    for (let i = committed.length - 1; i >= 0; i--) {
+        let unit = committed[i];
+        let first = getSiblingUnitFirstNode(unit);
+        let ref = anchor?.parentNode === parent ? anchor : null;
+        if (unit.kind === 'node' &&
+            isPreservedDomElement(unit.node) &&
+            unit.node.parentNode === parent) {
+            anchor = unit.node;
+            continue;
         }
-        // Advance anchor only after the node is placed in the correct parent
-        if (node.parentNode === current) {
-            anchor = node;
+        placeSiblingUnitBefore(unit, parent, ref);
+        if (first.parentNode === parent)
+            anchor = first;
+    }
+    for (let i = 0; i < currentUnits.length; i++) {
+        if (!used[i])
+            removeSiblingUnit(currentUnits[i], parent, context);
+    }
+}
+function parseSiblingUnits(nodes) {
+    let units = [];
+    for (let i = 0; i < nodes.length; i++) {
+        let node = nodes[i];
+        if (isVirtualRootStartMarker(node)) {
+            let endIndex = findHydrationEndIndex(nodes, i);
+            invariant(endIndex > i, 'Hydration end marker not found');
+            let end = nodes[endIndex];
+            invariant(isVirtualRootEndMarker(end), 'Expected hydration end marker');
+            units.push({ kind: 'hydration', start: node, end, startIndex: i, endIndex });
+            i = endIndex;
+            continue;
         }
+        if (isFrameStartMarker(node)) {
+            let endIndex = findFrameEndIndex(nodes, i);
+            invariant(endIndex > i, 'Frame end marker not found');
+            let end = nodes[endIndex];
+            invariant(isFrameEndMarker(end), 'Expected frame end marker');
+            units.push({ kind: 'frame', start: node, end, startIndex: i, endIndex });
+            i = endIndex;
+            continue;
+        }
+        invariant(!isVirtualRootEndMarker(node), 'Unexpected hydration end marker');
+        invariant(!isFrameEndMarker(node), 'Unexpected frame end marker');
+        units.push({ kind: 'node', node, startIndex: i, endIndex: i });
     }
-    // Remove any current children not used
-    for (let i = 0; i < currentChildren.length; i++) {
-        if (!used[i]) {
-            let nodeToRemove = currentChildren[i];
-            removeNode(nodeToRemove, current, context);
+    return units;
+}
+function getSiblingUnitKey(unit) {
+    if (unit.kind !== 'node' || !isElement(unit.node))
+        return;
+    return unit.node.getAttribute('data-key') ?? undefined;
+}
+function siblingUnitsComparable(current, next) {
+    if (current.kind !== next.kind)
+        return false;
+    if (current.kind !== 'node' || next.kind !== 'node')
+        return true;
+    return nodeTypesComparable(current.node, next.node);
+}
+function shouldPreserveBoundaryUnit(current, next, context) {
+    if (current.kind !== next.kind)
+        return false;
+    return current.kind === 'frame'
+        ? shouldPreserveFrameStartMarker(current.start, next.start, context)
+        : shouldPreserveHydrationStartMarker(current.start, next.start, context);
+}
+function getSiblingUnitFirstNode(unit) {
+    return unit.kind === 'node' ? unit.node : unit.start;
+}
+function placeSiblingUnitBefore(unit, parent, ref) {
+    if (unit.kind === 'node') {
+        if (unit.node.parentNode !== parent || unit.node.nextSibling !== ref) {
+            parent.insertBefore(unit.node, ref);
         }
+        return;
+    }
+    let nodes = collectNodeRange(unit.start, unit.end);
+    if (unit.start.parentNode === parent && unit.end.nextSibling === ref)
+        return;
+    let fragment = unit.start.ownerDocument.createDocumentFragment();
+    for (let node of nodes)
+        fragment.appendChild(node);
+    parent.insertBefore(fragment, ref);
+}
+function removeSiblingUnit(unit, parent, context) {
+    if (unit.kind === 'node') {
+        removeNode(unit.node, parent, context);
+        return;
+    }
+    for (let node of collectNodeRange(unit.start, unit.end)) {
+        removeNode(node, parent, context);
+    }
+}
+function collectNodesBetween(start, end) {
+    let nodes = [];
+    let node = start.nextSibling;
+    while (node && node !== end) {
+        nodes.push(node);
+        node = node.nextSibling;
     }
+    return nodes;
+}
+function isPreservedDomElement(node) {
+    return isElement(node) && node.hasAttribute(REMIX_PRESERVE_DOM_ATTRIBUTE);
 }
 function nodeTypesComparable(a, b) {
     if (isTextNode(a) && isTextNode(b))
         return true;
     if (isElement(a) && isElement(b))
         return a.tagName === b.tagName;
-    if (isVirtualRootStartMarker(a) && isVirtualRootStartMarker(b))
-        return true;
-    if (isVirtualRootEndMarker(a) && isVirtualRootEndMarker(b))
-        return true;
     if (isCommentNode(a) && isCommentNode(b))
-        return true;
+        return markerKindsMatch(a, b);
     return false;
 }
+function getMarkerKind(node) {
+    if (!isCommentNode(node))
+        return undefined;
+    if (isFrameStartMarker(node))
+        return 'frame-start';
+    if (isFrameEndMarker(node))
+        return 'frame-end';
+    if (isVirtualRootStartMarker(node))
+        return 'virtual-root-start';
+    if (isVirtualRootEndMarker(node))
+        return 'virtual-root-end';
+    return undefined;
+}
+function markerKindsMatch(a, b) {
+    return getMarkerKind(a) === getMarkerKind(b);
+}
 function isHydrationEndComment(node) {
     return isCommentNode(node) && node.data.trim() === '/rmx:h';
 }
@@ -398,9 +447,16 @@ function findHydrationEndMarker(start) {
     throw new Error('Hydration end marker not found');
 }
 function findHydrationEndIndex(nodes, startIdx) {
+    let depth = 1;
     for (let j = startIdx + 1; j < nodes.length; j++) {
-        if (isHydrationEndComment(nodes[j]))
-            return j;
+        let node = nodes[j];
+        if (isVirtualRootStartMarker(node))
+            depth++;
+        if (isHydrationEndComment(node)) {
+            depth--;
+            if (depth === 0)
+                return j;
+        }
     }
     return startIdx;
 }
@@ -451,12 +507,21 @@ function isFrameEndMarker(node) {
 function shouldPreserveFrameStartMarker(current, next, context) {
     if (!isFrameStartMarker(next))
         return false;
-    let currentData = getFrameMarkerData(current, context);
     let nextData = getFrameMarkerData(next, context);
-    return (currentData !== undefined &&
+    let currentFrame = context.frameInstances.get(current);
+    return (currentFrame !== undefined &&
         nextData !== undefined &&
-        currentData.src === nextData.src &&
-        currentData.name === nextData.name);
+        currentFrame.matchesIdentity(nextData.src, nextData.name));
+}
+function shouldPreserveHydrationStartMarker(current, next, context) {
+    if (!isVirtualRootStartMarker(next))
+        return false;
+    let currentOwner = getClientEntryBoundaryOwner(current);
+    let nextData = getHydrationMarkerData(next, context);
+    return (currentOwner !== undefined &&
+        nextData !== undefined &&
+        currentOwner.identity.moduleUrl === nextData.moduleUrl &&
+        currentOwner.identity.exportName === nextData.exportName);
 }
 function getCommentMarkerRangeReplacement(current, next, currentNodes, nextNodes, currentIndex, nextIndex, context) {
     // Comment markers can represent owned DOM ranges, not standalone comments. If
@@ -472,6 +537,25 @@ function getCommentMarkerRangeReplacement(current, next, currentNodes, nextNodes
             nextEndIndex: findFrameEndIndex(nextNodes, nextIndex),
         };
     }
+    if (isVirtualRootStartMarker(current) &&
+        isVirtualRootStartMarker(next) &&
+        !shouldPreserveHydrationStartMarker(current, next, context)) {
+        return {
+            currentStart: current,
+            nextStart: next,
+            currentEndIndex: findHydrationEndIndex(currentNodes, currentIndex),
+            nextEndIndex: findHydrationEndIndex(nextNodes, nextIndex),
+        };
+    }
+}
+function getHydrationMarkerData(marker, context) {
+    let id = getHydrationId(marker);
+    return context.data.h?.[id];
+}
+function getHydrationId(marker) {
+    let trimmed = marker.data.trim();
+    invariant(trimmed.startsWith('rmx:h:'), 'Invalid hydration start marker');
+    return trimmed.slice('rmx:h:'.length);
 }
 function getFrameMarkerData(marker, context) {
     let id = getFrameId(marker);
@@ -483,8 +567,8 @@ function getFrameId(marker) {
     return trimmed.slice('rmx:f:'.length);
 }
 function replaceCommentMarkerRange(replacement, parent, context) {
-    let currentEnd = findFrameEndMarker(replacement.currentStart);
-    let nextEnd = findFrameEndMarker(replacement.nextStart);
+    let currentEnd = findCommentMarkerRangeEnd(replacement.currentStart);
+    let nextEnd = findCommentMarkerRangeEnd(replacement.nextStart);
     let nextNodes = collectNodeRange(replacement.nextStart, nextEnd);
     let currentNodes = collectNodeRange(replacement.currentStart, currentEnd);
     for (let node of nextNodes) {
@@ -494,6 +578,13 @@ function replaceCommentMarkerRange(replacement, parent, context) {
         removeNode(node, parent, context);
     }
 }
+function findCommentMarkerRangeEnd(start) {
+    if (isFrameStartMarker(start))
+        return findFrameEndMarker(start);
+    if (isVirtualRootStartMarker(start))
+        return findHydrationEndMarker(start);
+    throw new Error('Comment marker range start not found');
+}
 function collectNodeRange(start, end) {
     let nodes = [];
     let node = start;
@@ -528,8 +619,7 @@ function disposeRemovedVirtualRoots(node) {
         let next = stack.pop();
         if (!next)
             continue;
-        if (isHydratedVirtualRootStartMarker(next)) {
-            next.$rmx.dispose();
+        if (isVirtualRootStartMarker(next) && disposeClientEntryBoundary(next)) {
             continue;
         }
         for (let child of Array.from(next.childNodes)) {
@@ -561,9 +651,6 @@ function disposeFrameStartMarker(marker, context) {
 function isVirtualRootStartMarker(node) {
     return isCommentNode(node) && node.data.trim().startsWith('rmx:h:');
 }
-function isHydratedVirtualRootStartMarker(node) {
-    return isVirtualRootStartMarker(node) && '$rmx' in node;
-}
 function isVirtualRootEndMarker(node) {
     return isCommentNode(node) && node.data.trim() === '/rmx:h';
 }