Print block styles after wp-block-library and everything else at end of HEAD

Author: westonruter

src/wp-includes/script-loader.php

@@ -3634,7 +3634,7 @@ function wp_hoist_late_printed_styles() {
 	 * enhancement output buffer, essentially no style can be enqueued too late, because an output buffer filter can
 	 * always hoist it to the HEAD.
 	 */
-	add_filter( 'print_late_styles', '__return_true', PHP_INT_MAX );
+	add_filter( 'print_late_styles', '__return_true', PHP_INT_MAX ); // TODO: Remove.
 
 	/*
 	 * Print a placeholder comment where the late styles can be hoisted from the footer to be printed in the header
@@ -3645,11 +3645,44 @@ function wp_hoist_late_printed_styles() {
 	wp_add_inline_style( 'wp-block-library', $placeholder );
 
 	// Wrap print_late_styles() with a closure that captures the late-printed styles.
-	$printed_late_styles = '';
-	$capture_late_styles = static function () use ( &$printed_late_styles ) {
+	$printed_block_styles = '';
+	$printed_late_styles  = '';
+	$capture_late_styles  = static function () use ( &$printed_block_styles, &$printed_late_styles ) {
+		global $concatenate_scripts;
+		script_concat_settings();
+
+		// Print the styles related to on-demand block enqueues.
+		$all_block_style_handles = array();
+		foreach ( WP_Block_Type_Registry::get_instance()->get_all_registered() as $block_type ) {
+			foreach ( $block_type->style_handles as $style_handle ) {
+				$all_block_style_handles[] = $style_handle;
+			}
+		}
+		$all_block_style_handles = array_merge(
+			$all_block_style_handles,
+			array( 'global-styles', 'core-block-supports', 'block-style-variation-styles' ) // TODO: What else?
+		);
+
+		$enqueued_block_styles = array_values( array_intersect( $all_block_style_handles, wp_styles()->queue ) );
+		if ( count( $enqueued_block_styles ) > 0 ) {
+			wp_styles()->do_concat = $concatenate_scripts;
+			ob_start();
+			wp_styles()->do_items( $enqueued_block_styles );
+			_print_styles();
+			$printed_block_styles = ob_get_clean();
+			wp_styles()->reset();
+		}
+
+		/*
+		 * Print remaining styles not related to blocks. This is the same logic as in print_late_styles(), but without
+		 * the filter to control whether late styles are printed (since they are being hoisted anyway).
+		 */
+		wp_styles()->do_concat = $concatenate_scripts;
 		ob_start();
-		print_late_styles();
+		wp_styles()->do_footer_items();
+		_print_styles();
 		$printed_late_styles = ob_get_clean();
+		wp_styles()->reset();
 	};
 
 	/*
@@ -3680,7 +3713,7 @@ static function () use ( $capture_late_styles ) {
 	// Replace placeholder with the captured late styles.
 	add_filter(
 		'wp_template_enhancement_output_buffer',
-		function ( $buffer ) use ( $placeholder, &$printed_late_styles ) {
+		function ( $buffer ) use ( $placeholder, &$printed_block_styles, &$printed_late_styles ) {
 
 			// Anonymous subclass of WP_HTML_Tag_Processor which exposes underlying bookmark spans.
 			$processor = new class( $buffer ) extends WP_HTML_Tag_Processor {
@@ -3690,19 +3723,18 @@ public function get_span(): WP_HTML_Span {
 				}
 			};
 
-			// Loop over STYLE tags.
-			while ( $processor->next_tag( array( 'tag_closers' => 'visit' ) ) ) {
-
-				// We've encountered the inline style for the 'wp-block-library' stylesheet which probably has the placeholder comment.
+			// TODO: If there are no block styles to print, it would be nice to not have to replace the placehoolder comment.
+			// Locate the inline style for the 'wp-block-library' stylesheet which probably has the placeholder comment.
+			while ( $processor->next_tag( array( 'tag_name' => 'STYLE' ) ) ) {
 				if (
 					! $processor->is_tag_closer() &&
 					'STYLE' === $processor->get_tag() &&
 					'wp-block-library-inline-css' === $processor->get_attribute( 'id' )
 				) {
-					// If the inline style lacks the placeholder comment, then we have to continue until we get to </HEAD> to append the styles there.
+					// If the inline style lacks the placeholder comment, then all the styles will be inserted below before </HEAD>.
 					$css_text = $processor->get_modifiable_text();
 					if ( ! str_contains( $css_text, $placeholder ) ) {
-						continue;
+						break;
 					}
 
 					// Remove the placeholder now that we've located the inline style.
@@ -3715,25 +3747,47 @@ public function get_span(): WP_HTML_Span {
 						'',
 						array(
 							substr( $buffer, 0, $span->start + $span->length ),
-							$printed_late_styles,
+							$printed_block_styles,
 							substr( $buffer, $span->start + $span->length ),
 						)
 					);
+
+					// Prevent printing them again.
+					$printed_block_styles = '';
 					break;
 				}
+			}
+
+			// If there are remaining styles to hoist, either because
+			if ( $printed_block_styles || $printed_late_styles ) {
+
+				// Anonymous subclass of WP_HTML_Tag_Processor which exposes underlying bookmark spans.
+				$processor = new class( $buffer ) extends WP_HTML_Tag_Processor {
+					public function get_span(): WP_HTML_Span {
+						$this->set_bookmark( 'here' );
+						return $this->bookmarks['here'];
+					}
+				};
 
 				// As a fallback, append the hoisted late styles to the end of the HEAD.
-				if ( $processor->is_tag_closer() && 'HEAD' === $processor->get_tag() ) {
-					$span   = $processor->get_span();
-					$buffer = implode(
-						'',
-						array(
-							substr( $buffer, 0, $span->start ),
-							$printed_late_styles,
-							substr( $buffer, $span->start ),
-						)
-					);
-					break;
+				while ( $processor->next_tag(
+					array(
+						'tag_name'    => 'HEAD',
+						'tag_closers' => 'visit',
+					)
+				) ) {
+					if ( $processor->is_tag_closer() ) {
+						$span   = $processor->get_span();
+						$buffer = implode(
+							'',
+							array(
+								substr( $buffer, 0, $span->start ),
+								$printed_block_styles . $printed_late_styles,
+								substr( $buffer, $span->start ),
+							)
+						);
+						break;
+					}
 				}
 			}
 

tests/phpunit/tests/template.php

@@ -1075,9 +1075,7 @@ public function data_wp_hoist_late_printed_styles(): array {
 			'early-css',
 			'early-inline-css',
 			'wp-emoji-styles-inline-css',
-			'wp-block-library-css',
 			'wp-block-library-inline-css',
-
 			'wp-block-separator-inline-css',
 			'global-styles-inline-css',
 			'core-block-supports-inline-css',
@@ -1109,16 +1107,14 @@ public function data_wp_hoist_late_printed_styles(): array {
 					'wp-emoji-styles-inline-css',
 					'wp-block-library-css',
 					'wp-block-library-inline-css',
-					'late-css',
-					'late-inline-css',
-
-					// TODO: The following three are enqueued in a different order compared to $expected_head_styles above. Why?
 					'core-block-supports-inline-css',
 					'classic-theme-styles-inline-css',
 					'global-styles-inline-css',
-
 					'normal-css',
 					'normal-inline-css',
+					'wp-custom-css',
+					'late-css',
+					'late-inline-css',
 				),
 			),
 			'disabled_printing_late_styles'   => array(
@@ -1183,6 +1179,8 @@ public function data_wp_hoist_late_printed_styles(): array {
 	 */
 	public function test_wp_hoist_late_printed_styles( ?Closure $set_up, array $theme_supports, array $expected ): void {
 		switch_theme( 'default' );
+		global $wp_styles;
+		$wp_styles = null;
 
 		foreach ( $theme_supports as $theme_support ) {
 			add_theme_support( $theme_support );
@@ -1205,6 +1203,7 @@ static function () {
 		foreach ( array_keys( $block_registry->get_all_registered() ) as $block_name ) {
 			$block_registry->unregister( $block_name );
 		}
+		register_core_block_style_handles();
 		register_core_block_types_from_metadata();
 
 		$this->assertFalse( wp_is_block_theme(), 'Test is only relevant to block themes.' );
@@ -1252,7 +1251,7 @@ static function () {
 
 		$this->assertStringContainsString( '</head>', $buffer, 'Expected the closing HEAD tag to be in the response.' );
 
-		$this->assertDoesNotMatchRegularExpression( '#/\*wp_late_styles_placeholder:[a-f0-9-]+\*/#', $filtered_buffer, 'Expected the placeholder to be removed.' );
+		$this->assertDoesNotMatchRegularExpression( '#/\*wp_late_styles_placeholder:[a-f0-9]+\*/#', $filtered_buffer, 'Expected the placeholder to be removed.' );
 		$found_styles = array(
 			'HEAD' => array(),
 			'BODY' => array(),