Improve handling of get_updated_html

plugins/optimization-detective/class-od-html-tag-processor.php

@@ -178,6 +178,15 @@ final class OD_HTML_Tag_Processor extends WP_HTML_Tag_Processor {
 	 */
 	private $buffered_text_replacements = array();
 
+	/**
+	 * Whether the end of the document was reached.
+	 *
+	 * @since n.e.x.t
+	 * @see self::next_token()
+	 * @var bool
+	 */
+	private $reached_end_of_document = false;
+
 	/**
 	 * Count for the number of times that the cursor was moved.
 	 *
@@ -263,6 +272,9 @@ public function next_token(): bool {
 		if ( ! parent::next_token() ) {
 			$this->open_stack_tags    = array();
 			$this->open_stack_indices = array();
+
+			// Mark that the end of the document was reached, meaning that get_modified_html() can should now be able to append markup to the HEAD and the BODY.
+			$this->reached_end_of_document = true;
 			return false;
 		}
 
@@ -557,16 +569,27 @@ public function append_body_html( string $html ): void {
 	}
 
 	/**
-	 * Gets the final updated HTML.
+	 * Returns the string representation of the HTML Tag Processor.
 	 *
-	 * This should only be called after the closing HTML tag has been reached and just before
-	 * calling {@see WP_HTML_Tag_Processor::get_updated_html()} to send the document back in the response.
+	 * This can only be called once the end of the document has been reached. It is responsible for adding the pending
+	 * markup to append to the HEAD and the BODY. Originally this was done in an overridden get_updated_html() method
+	 * before calling the parent method. However, every time that seek() is called it the HTML Processor will flush any
+	 * pending updates to the document. This means that if there is any pending markup to append to the end of the BODY
+	 * then the insertion will fail because the closing tag for the BODY has not been encountered yet. Additionally, by
+	 * not processing the buffered text replacements in get_updated_html() then we avoid trying to insert them every
+	 * time that seek() is called which is wasteful as they are only needed once finishing iterating over the document.
 	 *
-	 * @since n.e.x.t
+	 * @since 0.4.0
+	 * @see WP_HTML_Tag_Processor::get_updated_html()
+	 * @see WP_HTML_Tag_Processor::seek()
 	 *
-	 * @return string Final updated HTML.
+	 * @return string The processed HTML.
 	 */
-	public function get_final_updated_html(): string {
+	public function get_updated_html(): string {
+		if ( ! $this->reached_end_of_document ) {
+			return parent::get_updated_html();
+		}
+
 		foreach ( array_keys( $this->buffered_text_replacements ) as $bookmark ) {
 			$html_strings = $this->buffered_text_replacements[ $bookmark ];
 			if ( count( $html_strings ) === 0 ) {

plugins/optimization-detective/optimization.php

@@ -250,5 +250,5 @@ function od_optimize_template_output_buffer( string $buffer ): string {
 		$processor->append_body_html( od_get_detection_script( $slug, $group_collection ) );
 	}
 
-	return $processor->get_final_updated_html();
+	return $processor->get_updated_html();
 }

plugins/optimization-detective/tests/test-class-od-html-tag-processor.php

@@ -343,7 +343,7 @@ public function test_next_tag_with_query(): void {
 	 *
 	 * @covers ::append_head_html
 	 * @covers ::append_body_html
-	 * @covers ::get_final_updated_html
+	 * @covers ::get_updated_html
 	 */
 	public function test_append_head_and_body_html(): void {
 		$html                = '
@@ -370,16 +370,29 @@ public function test_append_head_and_body_html(): void {
 
 		$saw_head = false;
 		$saw_body = false;
+		$did_seek = false;
 		while ( $processor->next_open_tag() ) {
+			$this->assertStringNotContainsString( $head_injected, $processor->get_updated_html(), 'Only expecting end-of-head injection once document was finalized.' );
+			$this->assertStringNotContainsString( $body_injected, $processor->get_updated_html(), 'Only expecting end-of-body injection once document was finalized.' );
 			$tag = $processor->get_tag();
 			if ( 'HEAD' === $tag ) {
 				$saw_head = true;
 			} elseif ( 'BODY' === $tag ) {
 				$saw_body = true;
+				$this->assertTrue( $processor->set_bookmark( 'cuerpo' ) );
+			}
+			if ( ! $did_seek && 'H1' === $tag ) {
+				$processor->append_head_html( '<!--H1 appends to HEAD-->' );
+				$processor->append_body_html( '<!--H1 appends to BODY-->' );
+				$this->assertTrue( $processor->seek( 'cuerpo' ) );
+				$did_seek = true;
 			}
 		}
+		$this->assertTrue( $did_seek );
 		$this->assertTrue( $saw_head );
 		$this->assertTrue( $saw_body );
+		$this->assertStringContainsString( $head_injected, $processor->get_updated_html(), 'Only expecting end-of-head injection once document was finalized.' );
+		$this->assertStringContainsString( $body_injected, $processor->get_updated_html(), 'Only expecting end-of-body injection once document was finalized.' );
 
 		$processor->append_head_html( $later_head_injected );
 
@@ -388,16 +401,16 @@ public function test_append_head_and_body_html(): void {
 				<head>
 					<meta charset=utf-8>
 					<!-- </head> -->
-				{$head_injected}{$later_head_injected}</head>
+				{$head_injected}<!--H1 appends to HEAD-->{$later_head_injected}</head>
 				<!--</HEAD>-->
 				<body>
 					<h1>Hello World</h1>
 					<!-- </body> -->
-				{$body_injected}</body>
+				{$body_injected}<!--H1 appends to BODY--></body>
 				<!--</BODY>-->
 			</html>
 		";
-		$this->assertSame( $expected, $processor->get_final_updated_html() );
+		$this->assertSame( $expected, $processor->get_updated_html() );
 	}
 
 	/**