Refactor Markdown length-limited summary implementation

This commit refactors the implementation of
`markdown_summary_with_limit()`, separating the logic of determining
when the limit has been reached from the actual rendering process.

The main advantage of the new approach is that it guarantees that all
HTML tags are closed, whereas the previous implementation could generate
tags that were never closed. It also ensures that no empty tags are
generated (e.g., `<em></em>`).

The new implementation consists of a general-purpose struct
`HtmlWithLimit` that manages the length-limiting logic and a function
`markdown_summary_with_limit()` that renders Markdown to HTML using the
struct.

4 files changed, 125 insertions, 43 deletions

diff --git a/src/librustdoc/html/length_limit.rs b/src/librustdoc/html/length_limit.rs
new file mode 100644
index 00000000000..42b94b51118
--- /dev/null
+++ b/src/librustdoc/html/length_limit.rs
@@ -0,0 +1,94 @@
+//! See [`HtmlWithLimit`].
+
+use std::fmt::Write;
+use std::ops::ControlFlow;
+
+use crate::html::escape::Escape;
+
+/// A buffer that allows generating HTML with a length limit.
+///
+/// This buffer ensures that:
+///
+/// * all tags are closed,
+/// * only the most recently opened tag is closed,
+/// * no tags are left empty (e.g., `<em></em>`) due to the length limit being reached,
+/// * all text is escaped.
+#[derive(Debug)]
+pub(super) struct HtmlWithLimit {
+    buf: String,
+    len: usize,
+    limit: usize,
+    /// A list of tags that have been requested to be opened via [`Self::open_tag()`]
+    /// but have not actually been pushed to `buf` yet. This ensures that tags are not
+    /// left empty (e.g., `<em></em>`) due to the length limit being reached.
+    queued_tags: Vec<&'static str>,
+    /// A list of all tags that have been opened but not yet closed.
+    unclosed_tags: Vec<&'static str>,
+}
+
+impl HtmlWithLimit {
+    /// Create a new buffer, with a limit of `length_limit`.
+    pub(super) fn new(length_limit: usize) -> Self {
+        Self {
+            buf: String::new(),
+            len: 0,
+            limit: length_limit,
+            unclosed_tags: Vec::new(),
+            queued_tags: Vec::new(),
+        }
+    }
+
+    /// Finish using the buffer and get the written output.
+    /// This function will close all unclosed tags for you.
+    pub(super) fn finish(mut self) -> String {
+        self.close_all_tags();
+        self.buf
+    }
+
+    /// Write some plain text to the buffer, escaping as needed.
+    ///
+    /// This function skips writing the text if the length limit was reached
+    /// and returns [`ControlFlow::Break`].
+    pub(super) fn push(&mut self, text: &str) -> ControlFlow<(), ()> {
+        if self.len + text.len() > self.limit {
+            return ControlFlow::BREAK;
+        }
+
+        self.flush_queue();
+        write!(self.buf, "{}", Escape(text)).unwrap();
+        self.len += text.len();
+
+        ControlFlow::CONTINUE
+    }
+
+    /// Open an HTML tag.
+    pub(super) fn open_tag(&mut self, tag_name: &'static str) {
+        self.queued_tags.push(tag_name);
+    }
+
+    /// Close the most recently opened HTML tag.
+    pub(super) fn close_tag(&mut self) {
+        let tag_name = self.unclosed_tags.pop().unwrap();
+        self.buf.push_str("</");
+        self.buf.push_str(tag_name);
+        self.buf.push('>');
+    }
+
+    /// Write all queued tags and add them to the `unclosed_tags` list.
+    fn flush_queue(&mut self) {
+        for tag_name in self.queued_tags.drain(..) {
+            self.buf.push('<');
+            self.buf.push_str(tag_name);
+            self.buf.push('>');
+
+            self.unclosed_tags.push(tag_name);
+        }
+    }
+
+    /// Close all unclosed tags.
+    fn close_all_tags(&mut self) {
+        while !self.unclosed_tags.is_empty() {
+            self.close_tag();
+        }
+    }
+}
diff --git a/src/librustdoc/html/markdown.rs b/src/librustdoc/html/markdown.rs
index 472323daf30..5a569c690e5 100644
--- a/src/librustdoc/html/markdown.rs
+++ b/src/librustdoc/html/markdown.rs
@@ -23,12 +23,13 @@ use rustc_hir::HirId;
 use rustc_middle::ty::TyCtxt;
 use rustc_span::edition::Edition;
 use rustc_span::Span;
+
 use std::borrow::Cow;
 use std::cell::RefCell;
 use std::collections::VecDeque;
 use std::default::Default;
 use std::fmt::Write;
-use std::ops::Range;
+use std::ops::{ControlFlow, Range};
 use std::str;
 
 use crate::clean::RenderedLink;
@@ -36,6 +37,7 @@ use crate::doctest;
 use crate::html::escape::Escape;
 use crate::html::format::Buffer;
 use crate::html::highlight;
+use crate::html::length_limit::HtmlWithLimit;
 use crate::html::toc::TocBuilder;
 
 use pulldown_cmark::{
@@ -1081,15 +1083,6 @@ fn markdown_summary_with_limit(
         return (String::new(), false);
     }
 
-    let mut s = String::with_capacity(md.len() * 3 / 2);
-    let mut text_length = 0;
-    let mut stopped_early = false;
-
-    fn push(s: &mut String, text_length: &mut usize, text: &str) {
-        write!(s, "{}", Escape(text)).unwrap();
-        *text_length += text.len();
-    }
-
     let mut replacer = |broken_link: BrokenLink<'_>| {
         if let Some(link) =
             link_names.iter().find(|link| &*link.original_text == broken_link.reference)
@@ -1101,56 +1094,49 @@ fn markdown_summary_with_limit(
     };
 
     let p = Parser::new_with_broken_link_callback(md, opts(), Some(&mut replacer));
-    let p = LinkReplacer::new(p, link_names);
+    let mut p = LinkReplacer::new(p, link_names);
 
-    'outer: for event in p {
+    // FIXME: capacity
+    let mut buf = HtmlWithLimit::new(length_limit);
+    let mut stopped_early = false;
+    p.try_for_each(|event| {
         match &event {
             Event::Text(text) => {
-                for word in text.split_inclusive(char::is_whitespace) {
-                    if text_length + word.len() >= length_limit {
-                        stopped_early = true;
-                        break 'outer;
-                    }
-
-                    push(&mut s, &mut text_length, word);
+                let r =
+                    text.split_inclusive(char::is_whitespace).try_for_each(|word| buf.push(word));
+                if r.is_break() {
+                    stopped_early = true;
                 }
+                return r;
             }
             Event::Code(code) => {
-                if text_length + code.len() >= length_limit {
+                buf.open_tag("code");
+                let r = buf.push(code);
+                if r.is_break() {
                     stopped_early = true;
-                    break;
+                } else {
+                    buf.close_tag();
                 }
-
-                s.push_str("<code>");
-                push(&mut s, &mut text_length, code);
-                s.push_str("</code>");
+                return r;
             }
             Event::Start(tag) => match tag {
-                Tag::Emphasis => s.push_str("<em>"),
-                Tag::Strong => s.push_str("<strong>"),
-                Tag::CodeBlock(..) => break,
+                Tag::Emphasis => buf.open_tag("em"),
+                Tag::Strong => buf.open_tag("strong"),
+                Tag::CodeBlock(..) => return ControlFlow::BREAK,
                 _ => {}
             },
             Event::End(tag) => match tag {
-                Tag::Emphasis => s.push_str("</em>"),
-                Tag::Strong => s.push_str("</strong>"),
-                Tag::Paragraph => break,
-                Tag::Heading(..) => break,
+                Tag::Emphasis | Tag::Strong => buf.close_tag(),
+                Tag::Paragraph | Tag::Heading(..) => return ControlFlow::BREAK,
                 _ => {}
             },
-            Event::HardBreak | Event::SoftBreak => {
-                if text_length + 1 >= length_limit {
-                    stopped_early = true;
-                    break;
-                }
-
-                push(&mut s, &mut text_length, " ");
-            }
+            Event::HardBreak | Event::SoftBreak => buf.push(" ")?,
             _ => {}
-        }
-    }
+        };
+        ControlFlow::CONTINUE
+    });
 
-    (s, stopped_early)
+    (buf.finish(), stopped_early)
 }
 
 /// Renders a shortened first paragraph of the given Markdown as a subset of Markdown,
diff --git a/src/librustdoc/html/mod.rs b/src/librustdoc/html/mod.rs
index 60ebdf5690d..109b0a356db 100644
--- a/src/librustdoc/html/mod.rs
+++ b/src/librustdoc/html/mod.rs
@@ -2,6 +2,7 @@ crate mod escape;
 crate mod format;
 crate mod highlight;
 crate mod layout;
+mod length_limit;
 // used by the error-index generator, so it needs to be public
 pub mod markdown;
 crate mod render;
diff --git a/src/librustdoc/lib.rs b/src/librustdoc/lib.rs
index e02d92b11b8..ab94e0d5683 100644
--- a/src/librustdoc/lib.rs
+++ b/src/librustdoc/lib.rs
@@ -5,6 +5,7 @@
 #![feature(rustc_private)]
 #![feature(array_methods)]
 #![feature(box_patterns)]
+#![feature(control_flow_enum)]
 #![feature(in_band_lifetimes)]
 #![feature(nll)]
 #![feature(test)]