Core Data: Improve blocks cache in useEntityBlockEditor

packages/core-data/src/hooks/use-entity-block-editor.js

@@ -13,7 +13,7 @@ import useEntityId from './use-entity-id';
 import { updateFootnotesFromMeta } from '../footnotes';
 
 const EMPTY_ARRAY = [];
-const parsedBlocksCache = new WeakMap();
+const parsedBlocksCache = new Map();
 
 /**
  * Hook that returns block content getters and setters for
@@ -36,7 +36,6 @@ const parsedBlocksCache = new WeakMap();
 export default function useEntityBlockEditor( kind, name, { id: _id } = {} ) {
 	const providerId = useEntityId( kind, name );
 	const id = _id ?? providerId;
-	const { getEntityRecord, getEntityRecordEdits } = useSelect( STORE_NAME );
 	const { content, editedBlocks, meta } = useSelect(
 		( select ) => {
 			if ( ! id ) {
@@ -68,28 +67,21 @@ export default function useEntityBlockEditor( kind, name, { id: _id } = {} ) {
 			return EMPTY_ARRAY;
 		}
 
-		// If there's an edit, cache the parsed blocks by the edit.
-		// If not, cache by the original entity record.
-		const edits = getEntityRecordEdits( kind, name, id );
-		const isUnedited = ! edits || ! Object.keys( edits ).length;
-		const cackeKey = isUnedited ? getEntityRecord( kind, name, id ) : edits;
-		let _blocks = parsedBlocksCache.get( cackeKey );
+		// Cache parsed blocks by entity identity. Store the content
+		// alongside the blocks so we can validate it hasn't changed.
+		const cacheKey = `${ kind }:${ name }:${ id }`;
+		const cached = parsedBlocksCache.get( cacheKey );
+		let _blocks;
 
-		if ( ! _blocks ) {
+		if ( cached && cached.content === content ) {
+			_blocks = cached.blocks;
+		} else {
 			_blocks = parse( content );
-			parsedBlocksCache.set( cackeKey, _blocks );
+			parsedBlocksCache.set( cacheKey, { content, blocks: _blocks } );
 		}
 
 		return _blocks;
-	}, [
-		kind,
-		name,
-		id,
-		editedBlocks,
-		content,
-		getEntityRecord,
-		getEntityRecordEdits,
-	] );
+	}, [ kind, name, id, editedBlocks, content ] );
 
 	const onChange = useCallback(
 		( newBlocks, options ) => {

packages/core-data/src/test/entity-provider.js

@@ -276,4 +276,78 @@ describe( 'useEntityBlockEditor', () => {
 			'A paragraph<sup data-fn="abcd" class="fn"><a href="#abcd" id="abcd-link">2</a></sup>'
 		);
 	} );
+
+	it( 'preserves block clientIds across unmount and remount when content is unchanged', () => {
+		let blocks;
+		const TestComponent = () => {
+			[ blocks ] = useEntityBlockEditor( 'postType', 'post', {
+				id: 1,
+			} );
+			return <div />;
+		};
+
+		const { unmount } = render(
+			<RegistryProvider value={ registry }>
+				<TestComponent />
+			</RegistryProvider>
+		);
+
+		const firstClientIds = blocks.map( ( b ) => b.clientId );
+		expect( firstClientIds ).toHaveLength( 2 );
+
+		// Simulate navigating away.
+		unmount();
+
+		// Simulate navigating back — same entity, same content.
+		render(
+			<RegistryProvider value={ registry }>
+				<TestComponent />
+			</RegistryProvider>
+		);
+
+		// The cache should return the same block objects with the same clientIds.
+		expect( blocks.map( ( b ) => b.clientId ) ).toEqual( firstClientIds );
+	} );
+
+	it( 'returns new blocks when content changes', () => {
+		let blocks;
+		const TestComponent = () => {
+			[ blocks ] = useEntityBlockEditor( 'postType', 'post', {
+				id: 1,
+			} );
+			return <div />;
+		};
+
+		render(
+			<RegistryProvider value={ registry }>
+				<TestComponent />
+			</RegistryProvider>
+		);
+
+		const firstClientIds = blocks.map( ( b ) => b.clientId );
+
+		// Receive a new entity record with different content.
+		act( () => {
+			registry
+				.dispatch( coreDataStore )
+				.receiveEntityRecords( 'postType', 'post', [
+					{
+						id: 1,
+						type: 'post',
+						content: {
+							raw: '<!-- wp:test-block --><p>Different content</p><!-- /wp:test-block -->',
+							rendered: '<p>Different content</p>',
+						},
+						meta: { footnotes: '[]' },
+					},
+				] );
+		} );
+
+		// Blocks should be new objects with new clientIds.
+		expect( blocks ).toHaveLength( 1 );
+		expect( blocks[ 0 ].attributes.content ).toEqual( 'Different content' );
+		expect( blocks.map( ( b ) => b.clientId ) ).not.toEqual(
+			firstClientIds
+		);
+	} );
 } );