Interactivity Router: Fix initial page cache

packages/block-library/src/query/view.js

@@ -23,7 +23,9 @@ store( 'core/query', {
 		*navigate( event ) {
 			const ctx = getContext();
 			const { ref } = getElement();
-			const { queryRef } = ctx;
+			const queryRef = ref.closest(
+				'.wp-block-query[data-wp-router-region]'
+			);
 			const isDisabled = queryRef?.dataset.wpNavigationDisabled;
 
 			if ( isValidLink( ref ) && isValidEvent( event ) && ! isDisabled ) {
@@ -41,8 +43,10 @@ store( 'core/query', {
 			}
 		},
 		*prefetch() {
-			const { queryRef } = getContext();
 			const { ref } = getElement();
+			const queryRef = ref.closest(
+				'.wp-block-query[data-wp-router-region]'
+			);
 			const isDisabled = queryRef?.dataset.wpNavigationDisabled;
 			if ( isValidLink( ref ) && ! isDisabled ) {
 				const { actions } = yield import(
@@ -63,10 +67,5 @@ store( 'core/query', {
 				yield actions.prefetch( ref.href );
 			}
 		},
-		setQueryRef() {
-			const ctx = getContext();
-			const { ref } = getElement();
-			ctx.queryRef = ref;
-		},
 	},
 } );

packages/interactivity-router/src/index.js

@@ -1,13 +1,12 @@
 /**
  * WordPress dependencies
  */
-import {
-	render,
-	directivePrefix,
-	toVdom,
-	getRegionRootFragment,
-	store,
-} from '@wordpress/interactivity';
+import { render, store, privateApis } from '@wordpress/interactivity';
+
+const { directivePrefix, getRegionRootFragment, initialVdom, toVdom } =
+	privateApis(
+		'I acknowledge that using private APIs means my theme or plugin will inevitably break in the next version of WordPress.'
+	);
 
 // The cache of visited and prefetched pages.
 const pages = new Map();
@@ -36,12 +35,14 @@ const fetchPage = async ( url, { html } ) => {
 
 // Return an object with VDOM trees of those HTML regions marked with a
 // `router-region` directive.
-const regionsToVdom = ( dom ) => {
+const regionsToVdom = ( dom, { vdom } = {} ) => {
 	const regions = {};
 	const attrName = `data-${ directivePrefix }-router-region`;
 	dom.querySelectorAll( `[${ attrName }]` ).forEach( ( region ) => {
 		const id = region.getAttribute( attrName );
-		regions[ id ] = toVdom( region );
+		regions[ id ] = vdom?.has( region )
+			? vdom.get( region )
+			: toVdom( region );
 	} );
 	const title = dom.querySelector( 'title' )?.innerText;
 	return { regions, title };
@@ -74,10 +75,10 @@ window.addEventListener( 'popstate', async () => {
 	}
 } );
 
-// Cache the current regions.
+// Cache the initial page using the intially parsed vDOM.
 pages.set(
 	getPagePath( window.location ),
-	Promise.resolve( regionsToVdom( document ) )
+	Promise.resolve( regionsToVdom( document, { vdom: initialVdom } ) )
 );
 
 // Variable to store the current navigation.

packages/interactivity/src/index.js

@@ -2,7 +2,9 @@
  * Internal dependencies
  */
 import registerDirectives from './directives';
-import { init } from './init';
+import { init, getRegionRootFragment, initialVdom } from './init';
+import { directivePrefix } from './constants';
+import { toVdom } from './vdom';
 
 export { store } from './store';
 export { directive, getContext, getElement, getNamespace } from './hooks';
@@ -15,14 +17,27 @@ export {
 	useCallback,
 	useMemo,
 } from './utils';
-export { directivePrefix } from './constants';
-export { toVdom } from './vdom';
-export { getRegionRootFragment } from './init';
 
 export { h as createElement, cloneElement, render } from 'preact';
 export { useContext, useState, useRef } from 'preact/hooks';
 export { deepSignal } from 'deepsignal';
 
+const requiredConsent =
+	'I acknowledge that using private APIs means my theme or plugin will inevitably break in the next version of WordPress.';
+
+export const privateApis = ( lock ) => {
+	if ( lock === requiredConsent ) {
+		return {
+			directivePrefix,
+			getRegionRootFragment,
+			initialVdom,
+			toVdom,
+		};
+	}
+
+	throw new Error( 'Forbidden access.' );
+};
+
 document.addEventListener( 'DOMContentLoaded', async () => {
 	registerDirectives();
 	await init();

packages/interactivity/src/init.js

@@ -28,6 +28,9 @@ function yieldToMain() {
 	} );
 }
 
+// Initial vDOM regions associated with its DOM element.
+export const initialVdom = new WeakMap();
+
 // Initialize the router with the initial DOM.
 export const init = async () => {
 	const nodes = document.querySelectorAll(
@@ -39,6 +42,7 @@ export const init = async () => {
 			await yieldToMain();
 			const fragment = getRegionRootFragment( node );
 			const vdom = toVdom( node );
+			initialVdom.set( node, vdom );
 			await yieldToMain();
 			hydrate( vdom, fragment );
 		}

packages/interactivity/src/vdom.js

@@ -35,7 +35,12 @@ const nsPathRegExp = /^([\w-_\/]+)::(.+)$/;
 
 export const hydratedIslands = new WeakSet();
 
-// Recursive function that transforms a DOM tree into vDOM.
+/**
+ * Recursive function that transforms a DOM tree into vDOM.
+ *
+ * @param {Node} root The root element or node to start traversing on.
+ * @return {import('preact').VNode[]} The resulting vDOM tree.
+ */
 export function toVdom( root ) {
 	const treeWalker = document.createTreeWalker(
 		root,