From 11583e157c0e58c4332d8e6a00259fb08e8efe44 Mon Sep 17 00:00:00 2001
From: Ed Page <eopage@gmail.com>
Date: Thu, 16 Mar 2023 10:17:56 -0500
Subject: [PATCH] refactor(help): Track via style via ANSI codes

---
 src/builder/arg.rs          |   2 +-
 src/builder/styled_str.rs   | 183 ++++++++++--------------------------
 src/error/format.rs         |  10 +-
 src/error/mod.rs            |   2 +-
 src/output/fmt.rs           |   8 +-
 src/output/help.rs          |  12 +--
 src/output/help_template.rs |  37 +++-----
 src/output/usage.rs         |  10 +-
 8 files changed, 82 insertions(+), 182 deletions(-)

diff --git a/src/builder/arg.rs b/src/builder/arg.rs
index bf898cbfd642..107676a9c5c1 100644
--- a/src/builder/arg.rs
+++ b/src/builder/arg.rs
@@ -4149,7 +4149,7 @@ impl Arg {
             styled.literal("-");
             styled.literal(s);
         }
-        styled.extend(self.stylize_arg_suffix(required).into_iter());
+        styled.push_styled(&self.stylize_arg_suffix(required));
         styled
     }
 
diff --git a/src/builder/styled_str.rs b/src/builder/styled_str.rs
index dde3057b3845..fac37e661407 100644
--- a/src/builder/styled_str.rs
+++ b/src/builder/styled_str.rs
@@ -2,27 +2,13 @@
 ///
 /// For now, this is the same as a [`Str`][crate::builder::Str].  This exists to reserve space in
 /// the API for exposing terminal styling.
-#[derive(Clone, Default, Debug, PartialEq, Eq)]
-pub struct StyledStr {
-    #[cfg(feature = "color")]
-    pieces: Vec<(Option<Style>, String)>,
-    #[cfg(not(feature = "color"))]
-    pieces: String,
-}
+#[derive(Clone, Default, Debug, PartialEq, Eq, PartialOrd, Ord)]
+pub struct StyledStr(String);
 
 impl StyledStr {
     /// Create an empty buffer
-    #[cfg(feature = "color")]
-    pub const fn new() -> Self {
-        Self { pieces: Vec::new() }
-    }
-
-    /// Create an empty buffer
-    #[cfg(not(feature = "color"))]
     pub const fn new() -> Self {
-        Self {
-            pieces: String::new(),
-        }
+        Self(String::new())
     }
 
     /// Display using [ANSI Escape Code](https://en.wikipedia.org/wiki/ANSI_escape_code) styling
@@ -66,43 +52,22 @@ impl StyledStr {
         self.stylize_(None, msg.into());
     }
 
-    pub(crate) fn stylize(&mut self, style: impl Into<Option<Style>>, msg: impl Into<String>) {
-        self.stylize_(style.into(), msg.into());
-    }
-
     pub(crate) fn trim(&mut self) {
-        self.trim_start();
-        self.trim_end();
+        self.0 = self.0.trim().to_owned()
     }
 
-    pub(crate) fn trim_start(&mut self) {
-        if let Some((_, item)) = self.iter_mut().next() {
-            *item = item.trim_start().to_owned();
-        }
-    }
-
-    #[cfg(feature = "color")]
-    pub(crate) fn trim_end(&mut self) {
-        if let Some((_, item)) = self.pieces.last_mut() {
-            *item = item.trim_end().to_owned();
-        }
-    }
-
-    #[cfg(not(feature = "color"))]
-    pub(crate) fn trim_end(&mut self) {
-        self.pieces = self.pieces.trim_end().to_owned();
+    #[cfg(feature = "help")]
+    pub(crate) fn replace_newline_var(&mut self) {
+        self.0 = self.0.replace("{n}", "\n");
     }
 
     #[cfg(feature = "help")]
     pub(crate) fn indent(&mut self, initial: &str, trailing: &str) {
-        if let Some((_, first)) = self.iter_mut().next() {
-            first.insert_str(0, initial);
-        }
+        self.0.insert_str(0, initial);
+
         let mut line_sep = "\n".to_owned();
         line_sep.push_str(trailing);
-        for (_, content) in self.iter_mut() {
-            *content = content.replace('\n', &line_sep);
-        }
+        self.0 = self.0.replace('\n', &line_sep);
     }
 
     #[cfg(all(not(feature = "wrap_help"), feature = "help"))]
@@ -110,42 +75,57 @@ impl StyledStr {
 
     #[cfg(feature = "wrap_help")]
     pub(crate) fn wrap(&mut self, hard_width: usize) {
+        let mut new = String::with_capacity(self.0.len());
+
+        let mut last = 0;
         let mut wrapper = crate::output::textwrap::wrap_algorithms::LineWrapper::new(hard_width);
-        for (_, content) in self.iter_mut() {
-            let mut total = Vec::new();
+        for content in self.iter_text() {
+            // Preserve styling
+            let current = content.as_ptr() as usize - self.0.as_str().as_ptr() as usize;
+            if last != current {
+                new.push_str(&self.0.as_str()[last..current]);
+            }
+            last = current + content.len();
+
             for (i, line) in content.split_inclusive('\n').enumerate() {
                 if 0 < i {
-                    // start of a section does not imply newline
+                    // reset char count on newline, skipping the start as we might have carried
+                    // over from a prior block of styled text
                     wrapper.reset();
                 }
                 let line = crate::output::textwrap::word_separators::find_words_ascii_space(line)
                     .collect::<Vec<_>>();
-                total.extend(wrapper.wrap(line));
+                new.extend(wrapper.wrap(line));
             }
-            let total = total.join("");
-            *content = total;
         }
+        if last != self.0.len() {
+            new.push_str(&self.0.as_str()[last..]);
+        }
+        new = new.trim_end().to_owned();
 
-        self.trim_end();
+        self.0 = new;
     }
 
     #[cfg(feature = "color")]
     fn stylize_(&mut self, style: Option<Style>, msg: String) {
         if !msg.is_empty() {
-            self.pieces.push((style, msg));
+            use std::fmt::Write as _;
+
+            let style = style.map(|s| s.as_style()).unwrap_or_default();
+            let _ = write!(self.0, "{}{}{}", style.render(), msg, style.render_reset());
         }
     }
 
     #[cfg(not(feature = "color"))]
     fn stylize_(&mut self, _style: Option<Style>, msg: String) {
-        self.pieces.push_str(&msg);
+        self.0.push_str(&msg);
     }
 
     #[inline(never)]
     #[cfg(feature = "help")]
     pub(crate) fn display_width(&self) -> usize {
         let mut width = 0;
-        for (_, c) in self.iter() {
+        for c in self.iter_text() {
             width += crate::output::display_width(c);
         }
         width
@@ -153,56 +133,30 @@ impl StyledStr {
 
     #[cfg(feature = "help")]
     pub(crate) fn is_empty(&self) -> bool {
-        self.pieces.is_empty()
+        self.0.is_empty()
     }
 
-    #[cfg(feature = "color")]
-    pub(crate) fn iter(&self) -> impl Iterator<Item = (Option<Style>, &str)> {
-        self.pieces.iter().map(|(s, c)| (*s, c.as_str()))
-    }
-
-    #[cfg(not(feature = "color"))]
-    pub(crate) fn iter(&self) -> impl Iterator<Item = (Option<Style>, &str)> {
-        [(None, self.pieces.as_str())].into_iter()
-    }
-
-    #[cfg(feature = "color")]
-    pub(crate) fn iter_mut(&mut self) -> impl Iterator<Item = (Option<Style>, &mut String)> {
-        self.pieces.iter_mut().map(|(s, c)| (*s, c))
-    }
-
-    #[cfg(not(feature = "color"))]
-    pub(crate) fn iter_mut(&mut self) -> impl Iterator<Item = (Option<Style>, &mut String)> {
-        [(None, &mut self.pieces)].into_iter()
+    #[cfg(feature = "help")]
+    pub(crate) fn as_styled_str(&self) -> &str {
+        &self.0
     }
 
     #[cfg(feature = "color")]
-    pub(crate) fn into_iter(self) -> impl Iterator<Item = (Option<Style>, String)> {
-        self.pieces.into_iter()
+    pub(crate) fn iter_text(&self) -> impl Iterator<Item = &str> {
+        anstream::adapter::strip_str(&self.0)
     }
 
     #[cfg(not(feature = "color"))]
-    pub(crate) fn into_iter(self) -> impl Iterator<Item = (Option<Style>, String)> {
-        [(None, self.pieces)].into_iter()
+    pub(crate) fn iter_text(&self) -> impl Iterator<Item = &str> {
+        [self.0.as_str()].into_iter()
     }
 
-    pub(crate) fn extend(
-        &mut self,
-        other: impl IntoIterator<Item = (impl Into<Option<Style>>, impl Into<String>)>,
-    ) {
-        for (style, content) in other {
-            self.stylize(style.into(), content.into());
-        }
+    pub(crate) fn push_styled(&mut self, other: &Self) {
+        self.0.push_str(&other.0);
     }
 
-    #[cfg(feature = "color")]
-    pub(crate) fn write_colored(&self, buffer: &mut dyn std::io::Write) -> std::io::Result<()> {
-        for (style, content) in &self.pieces {
-            let style = style.map(|s| s.as_style()).unwrap_or_default();
-            ok!(style.write_to(buffer));
-            ok!(buffer.write_all(content.as_bytes()));
-            ok!(style.write_reset_to(buffer));
-        }
+    pub(crate) fn write_to(&self, buffer: &mut dyn std::io::Write) -> std::io::Result<()> {
+        ok!(buffer.write_all(self.0.as_bytes()));
 
         Ok(())
     }
@@ -245,29 +199,11 @@ impl From<&'_ &'static str> for StyledStr {
     }
 }
 
-impl PartialOrd for StyledStr {
-    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
-        Some(self.cmp(other))
-    }
-}
-
-impl Ord for StyledStr {
-    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
-        self.iter().map(cmp_key).cmp(other.iter().map(cmp_key))
-    }
-}
-
-fn cmp_key(c: (Option<Style>, &str)) -> (Option<usize>, &str) {
-    let style = c.0.map(|s| s.as_usize());
-    let content = c.1;
-    (style, content)
-}
-
 /// Color-unaware printing. Never uses coloring.
 impl std::fmt::Display for StyledStr {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
-        for (_, content) in self.iter() {
-            ok!(std::fmt::Display::fmt(content, f));
+        for part in self.iter_text() {
+            part.fmt(f)?;
         }
 
         Ok(())
@@ -282,14 +218,7 @@ struct AnsiDisplay<'s> {
 #[cfg(feature = "color")]
 impl std::fmt::Display for AnsiDisplay<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
-        for (style, content) in &self.styled.pieces {
-            let style = style.map(|s| s.as_style()).unwrap_or_default();
-            ok!(style.render().fmt(f));
-            ok!(content.fmt(f));
-            ok!(style.render_reset().fmt(f));
-        }
-
-        Ok(())
+        self.styled.0.fmt(f)
     }
 }
 
@@ -305,18 +234,6 @@ pub(crate) enum Style {
 }
 
 impl Style {
-    fn as_usize(&self) -> usize {
-        match self {
-            Self::Header => 0,
-            Self::Literal => 1,
-            Self::Placeholder => 2,
-            Self::Good => 3,
-            Self::Warning => 4,
-            Self::Error => 5,
-            Self::Hint => 6,
-        }
-    }
-
     #[cfg(feature = "color")]
     fn as_style(&self) -> anstyle::Style {
         match self {
diff --git a/src/error/format.rs b/src/error/format.rs
index d3ce6c64c19e..422f306c9b0c 100644
--- a/src/error/format.rs
+++ b/src/error/format.rs
@@ -100,13 +100,13 @@ impl ErrorFormatter for RichFormatter {
                 styled.none("\n");
                 styled.none(TAB);
                 styled.good("note: ");
-                styled.extend(suggestion.iter());
+                styled.push_styled(suggestion);
             }
         }
 
         let usage = error.get(ContextKind::Usage);
         if let Some(ContextValue::StyledStr(usage)) = usage {
-            put_usage(&mut styled, usage.clone());
+            put_usage(&mut styled, usage);
         }
 
         try_help(&mut styled, error.inner.help_flag);
@@ -377,7 +377,7 @@ fn write_dynamic_context(error: &crate::error::Error, styled: &mut StyledStr) ->
 pub(crate) fn format_error_message(
     message: &str,
     cmd: Option<&Command>,
-    usage: Option<StyledStr>,
+    usage: Option<&StyledStr>,
 ) -> StyledStr {
     let mut styled = StyledStr::new();
     start_error(&mut styled);
@@ -400,9 +400,9 @@ fn singular_or_plural(n: usize) -> &'static str {
     }
 }
 
-fn put_usage(styled: &mut StyledStr, usage: StyledStr) {
+fn put_usage(styled: &mut StyledStr, usage: &StyledStr) {
     styled.none("\n\n");
-    styled.extend(usage.into_iter());
+    styled.push_styled(usage);
 }
 
 pub(crate) fn get_help_flag(cmd: &Command) -> Option<&'static str> {
diff --git a/src/error/mod.rs b/src/error/mod.rs
index 2d8bf62e47c7..964ea645cda0 100644
--- a/src/error/mod.rs
+++ b/src/error/mod.rs
@@ -801,7 +801,7 @@ impl Message {
                 let mut message = String::new();
                 std::mem::swap(s, &mut message);
 
-                let styled = format::format_error_message(&message, Some(cmd), usage);
+                let styled = format::format_error_message(&message, Some(cmd), usage.as_ref());
 
                 *self = Self::Formatted(styled);
             }
diff --git a/src/output/fmt.rs b/src/output/fmt.rs
index 93137adb2488..0c9a24f33a22 100644
--- a/src/output/fmt.rs
+++ b/src/output/fmt.rs
@@ -53,25 +53,23 @@ impl Colorizer {
             }
         };
 
-        self.content.write_colored(writer)
+        self.content.write_to(writer)
     }
 
     #[cfg(not(feature = "color"))]
     pub(crate) fn print(&self) -> std::io::Result<()> {
-        use std::io::Write;
-
         // [e]println can't be used here because it panics
         // if something went wrong. We don't want that.
         match self.stream {
             Stream::Stdout => {
                 let stdout = std::io::stdout();
                 let mut stdout = stdout.lock();
-                write!(stdout, "{}", self)
+                self.content.write_to(&mut stdout)
             }
             Stream::Stderr => {
                 let stderr = std::io::stderr();
                 let mut stderr = stderr.lock();
-                write!(stderr, "{}", self)
+                self.content.write_to(&mut stderr)
             }
         }
     }
diff --git a/src/output/help.rs b/src/output/help.rs
index 4e4fdbfdcaad..4921f5f71f3e 100644
--- a/src/output/help.rs
+++ b/src/output/help.rs
@@ -10,21 +10,15 @@ pub(crate) fn write_help(writer: &mut StyledStr, cmd: &Command, usage: &Usage<'_
     debug!("write_help");
 
     if let Some(h) = cmd.get_override_help() {
-        writer.extend(h.iter());
+        writer.push_styled(h);
     } else {
         #[cfg(feature = "help")]
         {
             use super::AutoHelp;
             use super::HelpTemplate;
             if let Some(tmpl) = cmd.get_help_template() {
-                for (style, content) in tmpl.iter() {
-                    if style.is_none() {
-                        HelpTemplate::new(writer, cmd, usage, use_long)
-                            .write_templated_help(content);
-                    } else {
-                        writer.stylize(style, content);
-                    }
-                }
+                HelpTemplate::new(writer, cmd, usage, use_long)
+                    .write_templated_help(tmpl.as_styled_str());
             } else {
                 AutoHelp::new(writer, cmd, usage, use_long).write_help();
             }
diff --git a/src/output/help_template.rs b/src/output/help_template.rs
index 4c7adda9046b..b8f2705404ea 100644
--- a/src/output/help_template.rs
+++ b/src/output/help_template.rs
@@ -166,11 +166,8 @@ impl<'cmd, 'writer> HelpTemplate<'cmd, 'writer> {
                         self.header("Usage:");
                     }
                     "usage" => {
-                        self.writer.extend(
-                            self.usage
-                                .create_usage_no_title(&[])
-                                .unwrap_or_default()
-                                .into_iter(),
+                        self.writer.push_styled(
+                            &self.usage.create_usage_no_title(&[]).unwrap_or_default(),
                         );
                     }
                     "all-args" => {
@@ -283,9 +280,9 @@ impl<'cmd, 'writer> HelpTemplate<'cmd, 'writer> {
                 self.none("\n");
             }
             let mut output = output.clone();
-            replace_newline_var(&mut output);
+            output.replace_newline_var();
             output.wrap(self.term_w);
-            self.writer.extend(output.into_iter());
+            self.writer.push_styled(&output);
             if after_new_line {
                 self.none("\n");
             }
@@ -303,9 +300,9 @@ impl<'cmd, 'writer> HelpTemplate<'cmd, 'writer> {
         };
         if let Some(output) = before_help {
             let mut output = output.clone();
-            replace_newline_var(&mut output);
+            output.replace_newline_var();
             output.wrap(self.term_w);
-            self.writer.extend(output.into_iter());
+            self.writer.push_styled(&output);
             self.none("\n\n");
         }
     }
@@ -322,9 +319,9 @@ impl<'cmd, 'writer> HelpTemplate<'cmd, 'writer> {
         if let Some(output) = after_help {
             self.none("\n\n");
             let mut output = output.clone();
-            replace_newline_var(&mut output);
+            output.replace_newline_var();
             output.wrap(self.term_w);
-            self.writer.extend(output.into_iter());
+            self.writer.push_styled(&output);
         }
     }
 }
@@ -464,7 +461,7 @@ impl<'cmd, 'writer> HelpTemplate<'cmd, 'writer> {
         self.none(TAB);
         self.short(arg);
         self.long(arg);
-        self.writer.extend(arg.stylize_arg_suffix(None).into_iter());
+        self.writer.push_styled(&arg.stylize_arg_suffix(None));
         self.align_to_about(arg, next_line_help, longest);
 
         let about = if self.use_long {
@@ -574,7 +571,7 @@ impl<'cmd, 'writer> HelpTemplate<'cmd, 'writer> {
         let trailing_indent = self.get_spaces(trailing_indent);
 
         let mut help = about.clone();
-        replace_newline_var(&mut help);
+        help.replace_newline_var();
         if !spec_vals.is_empty() {
             if !help.is_empty() {
                 let sep = if self.use_long && arg.is_some() {
@@ -596,7 +593,7 @@ impl<'cmd, 'writer> HelpTemplate<'cmd, 'writer> {
         help.wrap(avail_chars);
         help.indent("", &trailing_indent);
         let help_is_empty = help.is_empty();
-        self.writer.extend(help.into_iter());
+        self.writer.push_styled(&help);
         if let Some(arg) = arg {
             const DASH_SPACE: usize = "- ".len();
             const COLON_SPACE: usize = ": ".len();
@@ -662,10 +659,10 @@ impl<'cmd, 'writer> HelpTemplate<'cmd, 'writer> {
                         };
 
                         let mut help = help.clone();
-                        replace_newline_var(&mut help);
+                        help.replace_newline_var();
                         help.wrap(avail_chars);
                         help.indent("", &trailing_indent);
-                        self.writer.extend(help.into_iter());
+                        self.writer.push_styled(&help);
                     }
                 }
             }
@@ -942,7 +939,7 @@ impl<'cmd, 'writer> HelpTemplate<'cmd, 'writer> {
         let width = sc_str.display_width();
 
         self.none(TAB);
-        self.writer.extend(sc_str.into_iter());
+        self.writer.push_styled(&sc_str);
         if !next_line_help {
             self.spaces(longest + TAB_WIDTH - width);
         }
@@ -1015,12 +1012,6 @@ fn should_show_subcommand(subcommand: &Command) -> bool {
     !subcommand.is_hide_set()
 }
 
-fn replace_newline_var(styled: &mut StyledStr) {
-    for (_, content) in styled.iter_mut() {
-        *content = content.replace("{n}", "\n");
-    }
-}
-
 fn longest_filter(arg: &Arg) -> bool {
     arg.is_takes_value_set() || arg.get_long().is_some() || arg.get_short().is_none()
 }
diff --git a/src/output/usage.rs b/src/output/usage.rs
index 145fce12d084..c49ca2cb0639 100644
--- a/src/output/usage.rs
+++ b/src/output/usage.rs
@@ -40,7 +40,7 @@ impl<'cmd> Usage<'cmd> {
         let mut styled = StyledStr::new();
         styled.header("Usage:");
         styled.none(" ");
-        styled.extend(usage.into_iter());
+        styled.push_styled(&usage);
         Some(styled)
     }
 
@@ -103,7 +103,7 @@ impl<'cmd> Usage<'cmd> {
                     // Short-circuit full usage creation since no args will be relevant
                     styled.literal(name);
                 } else {
-                    styled.extend(self.create_help_usage(false).into_iter());
+                    styled.push_styled(&self.create_help_usage(false));
                 }
                 styled.placeholder(" <");
                 styled.placeholder(placeholder);
@@ -190,7 +190,7 @@ impl<'cmd> Usage<'cmd> {
     pub(crate) fn write_args(&self, incls: &[Id], force_optional: bool, styled: &mut StyledStr) {
         for required in self.get_args(incls, force_optional) {
             styled.none(" ");
-            styled.extend(required.into_iter());
+            styled.push_styled(&required);
         }
     }
 
@@ -280,7 +280,7 @@ impl<'cmd> Usage<'cmd> {
                     let styled = required_positionals[index].take().unwrap();
                     let mut new = StyledStr::new();
                     new.literal("-- ");
-                    new.extend(styled.into_iter());
+                    new.push_styled(&styled);
                     required_positionals[index] = Some(new);
                 }
             } else {
@@ -288,7 +288,7 @@ impl<'cmd> Usage<'cmd> {
                 if pos.is_last_set() {
                     styled = StyledStr::new();
                     styled.literal("[-- ");
-                    styled.extend(pos.stylized(Some(true)).into_iter());
+                    styled.push_styled(&pos.stylized(Some(true)));
                     styled.literal("]");
                 } else {
                     styled = pos.stylized(Some(false));