From 98fc99226897d070dd885119a24c8795e598edff Mon Sep 17 00:00:00 2001
From: Dennis Snell <dennis.snell@automattic.com>
Date: Tue, 24 Oct 2023 10:25:41 +0200
Subject: [PATCH] WIP: HTML API: Add function to balance tags ala
 `force_balance_tags()`

```php
$balanced = WP_HTML_Processor::balance_tags( '<p><div><button>First<button><b here>Second' );
$balanced === '<p></p><div><button>First</button><button><b here>Second</b></button></div>';
```
---
 .../html-api/class-wp-html-open-elements.php  | 11 ++++
 .../html-api/class-wp-html-processor.php      | 54 +++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/src/wp-includes/html-api/class-wp-html-open-elements.php b/src/wp-includes/html-api/class-wp-html-open-elements.php
index fe5625545b0ac..c5f7aa69dedc4 100644
--- a/src/wp-includes/html-api/class-wp-html-open-elements.php
+++ b/src/wp-includes/html-api/class-wp-html-open-elements.php
@@ -51,6 +51,13 @@ class WP_HTML_Open_Elements {
 	 */
 	private $has_p_in_button_scope = false;
 
+	/**
+	 * Called when an element is popped from the stack of open elements.
+	 *
+	 * @var callable|null
+	 */
+	public $on_pop = null;
+
 	/**
 	 * Reports if a specific node is in the stack of open elements.
 	 *
@@ -428,5 +435,9 @@ public function after_element_pop( $item ) {
 				$this->has_p_in_button_scope = $this->has_element_in_button_scope( 'P' );
 				break;
 		}
+
+		if ( $this->on_pop ) {
+			call_user_func( $this->on_pop, $item );
+		}
 	}
 }
diff --git a/src/wp-includes/html-api/class-wp-html-processor.php b/src/wp-includes/html-api/class-wp-html-processor.php
index f27f83b028cd2..3f29b4918eeea 100644
--- a/src/wp-includes/html-api/class-wp-html-processor.php
+++ b/src/wp-includes/html-api/class-wp-html-processor.php
@@ -473,6 +473,60 @@ public function matches_breadcrumbs( $breadcrumbs ) {
 		return false;
 	}
 
+	/**
+	 * Balances tags
+	 *
+	 * @throws Exception When bookmarks can't be created.
+	 *
+	 * @param string $html HTML to balance.
+	 *
+	 * @return string
+	 */
+	public static function balance_tags( $html ) {
+		$processor = self::create_fragment( $html );
+		$output    = '';
+		$at        = 0;
+
+		/**
+		 * Adds closing tags.
+		 *
+		 * @param WP_HTML_Token $item Item popped off of stack.
+		 *
+		 * @return void
+		 */
+		$close_tag = function ( $item ) use ( &$at, $html, &$output, $processor ) {
+			$token    = $processor->bookmarks[ $processor->state->current_token->bookmark_name ];
+			$output  .= substr( $html, $at, $token->start - $at );
+			$tag_name = substr( $html, $processor->bookmarks[ $item->bookmark_name ]->start + 1, strlen( $item->node_name ) );
+			if ( null !== $processor->get_last_error() ) {
+				$output  .= "</{$tag_name}>";
+			}
+			$at       = $token->start;
+		};
+
+		$processor->state->stack_of_open_elements->on_pop = $close_tag;
+		while ( $processor->next_tag() ) {
+			continue;
+		}
+		if ( null !== $processor->get_last_error() ) {
+			echo "\e[34mError: \e[32m{$processor->get_last_error()}\e[34m at \e[32m{$processor->state->current_token->node_name}\e[m\n";
+		}
+
+		$output .= substr( $html, $at );
+
+		foreach ( $processor->state->stack_of_open_elements->walk_up() as $item ) {
+			if ( 'context-node' === $item->bookmark_name ) {
+				break;
+			}
+			$tag_name = substr( $html, $processor->bookmarks[ $item->bookmark_name ]->start + 1, strlen( $item->node_name ) );
+			if ( null !== $processor->get_last_error() ) {
+				$output  .= "</{$tag_name}>";
+			}
+		}
+
+		return $output;
+	}
+
 	/**
 	 * Steps through the HTML document and stop at the next tag, if any.
 	 *