Repository: rust-lang/annotate-snippets-rs
Pull request #30 (2020-04-01 22:33:17)

Fix to no copy when format with yansi_term
This PR eliminates the need to reallocate while formatting. 
For this I have forked the `ansi_term` library to refactor to `fmt::Display` allowing to introduce flow with a closure that takes the current formatter as argument. 
It will depend on `yansi-term` instead of` ansi_term`, which I will keep.

zbraniecki: @botika what's your take on `yansi-term` vs `termcolor`? @brendanzab says that `termcolor` may be a better long term bet.

I'm happy to land this PR, but I'm also curious on your opinion on that :)
botika: The library is very simple, just enter a prefix and a suffix. So it does not need much maintenance nor can it innovate significantly. With the fork in the `ansi_term` I have left the minimum to fulfill its function as quickly as possible

The reason for the PR is to remove all the unnecessary `mallocs` in the formatter.
zbraniecki: I understand. I'm just curious on your take on long-term replacing `ansi_term` or `yansi_term` with `termcolor`? 
botika: It is more optimal implement `fmt::Display` than `io::Write`. It forces you to use his buffer or stream output. I think the only thing this library should do is put a prefix and a suffix around the format flow, and let the user use whatever writer they want. So you couldn't use `fmt::Formatter` for output stream. And I am unaware of compatibility outside of unix systems 
botika: I've been looking at `termcolor` in more depth, just as I expected, it's forced to reallocate, [here allocate to stack](https://github.com/BurntSushi/termcolor/blob/master/src/lib.rs#L1372-L1373) and [here allocate to writer]( https://github.com/BurntSushi/termcolor/blob/master/src/lib.rs#L1398) (wherever the writer writes). 


So, the main problems of `termcolor` are unnecessary use of memory and lack of abstraction of the writer

Patch:
diff --git a/Cargo.toml b/Cargo.toml
index 9bb8f7c..086db9d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -16,14 +16,14 @@ coveralls = { repository = "rust-lang/annotate-snippets-rs", branch = "master",
 maintenance = { status = "actively-developed" }
 
 [dependencies]
-ansi_term = { version = "0.12", optional = true }
+yansi-term = { version = "0.1", optional = true }
 
 [dev-dependencies]
 glob = "0.3"
 serde_yaml = "0.8"
 serde = { version = "1.0", features = ["derive"] }
 difference = "2.0"
-ansi_term = "0.12"
+yansi-term = "0.1"
 criterion = "0.3"
 
 [[bench]]
@@ -32,4 +32,4 @@ harness = false
 
 [features]
 default = []
-color = ["ansi_term"]
+color = ["yansi-term"]
diff --git a/src/formatter/mod.rs b/src/formatter/mod.rs
index 7812a4d..f90d583 100644
--- a/src/formatter/mod.rs
+++ b/src/formatter/mod.rs
@@ -1,7 +1,6 @@
 use std::{
-    cell::Cell,
     cmp,
-    fmt::{self, Display, Formatter, Write},
+    fmt::{self, Display, Write},
 };
 
 pub mod style;
@@ -12,28 +11,6 @@ use self::style::{Style, StyleClass, Stylesheet};
 use crate::stylesheets::color::AnsiTermStylesheet;
 use crate::{display_list::*, stylesheets::no_color::NoColorStylesheet};
 
-pub struct DisplayFn<F: FnOnce(&mut Formatter<'_>) -> fmt::Result>(std::cell::Cell<Option<F>>);
-
-impl<F: FnOnce(&mut Formatter<'_>) -> fmt::Result> DisplayFn<F> {
-    pub fn new(f: F) -> Self {
-        Self(Cell::new(Some(f)))
-    }
-}
-
-impl<F: FnOnce(&mut Formatter<'_>) -> fmt::Result> Display for DisplayFn<F> {
-    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
-        self.0.take().ok_or(fmt::Error).and_then(|cl| cl(f))
-    }
-}
-
-fn repeat_char(c: char, n: usize) -> String {
-    let mut s = String::with_capacity(c.len_utf8() * n);
-    for _ in 0..n {
-        s.push(c);
-    }
-    s
-}
-
 fn format_repeat_char(c: char, n: usize, f: &mut fmt::Formatter<'_>) -> fmt::Result {
     for _ in 0..n {
         f.write_char(c)?;
@@ -72,14 +49,18 @@ impl fmt::Display for DisplayList {
                 lineno: Some(lineno),
                 ..
             } => {
-                if self.anonymized_line_numbers {
-                    Self::ANONYMIZED_LINE_NUM.len()
-                } else {
-                    cmp::max(lineno.to_string().len(), max)
-                }
+                // The largest line is the largest width.
+                cmp::max(*lineno, max)
             }
             _ => max,
         });
+        let lineno_width = if lineno_width == 0 {
+            lineno_width
+        } else if self.anonymized_line_numbers {
+            Self::ANONYMIZED_LINE_NUM.len()
+        } else {
+            ((lineno_width as f64).log10().floor() as usize) + 1
+        };
         let inline_marks_width = self.body.iter().fold(0, |max, line| match line {
             DisplayLine::Source { inline_marks, .. } => cmp::max(inline_marks.len(), max),
             _ => max,
@@ -97,22 +78,38 @@ impl fmt::Display for DisplayList {
 
 impl DisplayList {
     const ANONYMIZED_LINE_NUM: &'static str = "LL";
+    const ERROR_TXT: &'static str = "error";
+    const HELP_TXT: &'static str = "help";
+    const INFO_TXT: &'static str = "info";
+    const NOTE_TXT: &'static str = "note";
+    const WARNING_TXT: &'static str = "warning";
 
+    #[inline]
     fn format_annotation_type(
-        &self,
         annotation_type: &DisplayAnnotationType,
         f: &mut fmt::Formatter<'_>,
     ) -> fmt::Result {
         match annotation_type {
-            DisplayAnnotationType::Error => f.write_str("error"),
-            DisplayAnnotationType::Warning => f.write_str("warning"),
-            DisplayAnnotationType::Info => f.write_str("info"),
-            DisplayAnnotationType::Note => f.write_str("note"),
-            DisplayAnnotationType::Help => f.write_str("help"),
+            DisplayAnnotationType::Error => f.write_str(Self::ERROR_TXT),
+            DisplayAnnotationType::Help => f.write_str(Self::HELP_TXT),
+            DisplayAnnotationType::Info => f.write_str(Self::INFO_TXT),
+            DisplayAnnotationType::Note => f.write_str(Self::NOTE_TXT),
+            DisplayAnnotationType::Warning => f.write_str(Self::WARNING_TXT),
             DisplayAnnotationType::None => Ok(()),
         }
     }
 
+    fn annotation_type_len(annotation_type: &DisplayAnnotationType) -> usize {
+        match annotation_type {
+            DisplayAnnotationType::Error => Self::ERROR_TXT.len(),
+            DisplayAnnotationType::Help => Self::HELP_TXT.len(),
+            DisplayAnnotationType::Info => Self::INFO_TXT.len(),
+            DisplayAnnotationType::Note => Self::NOTE_TXT.len(),
+            DisplayAnnotationType::Warning => Self::WARNING_TXT.len(),
+            DisplayAnnotationType::None => 0,
+        }
+    }
+
     fn get_annotation_style(&self, annotation_type: &DisplayAnnotationType) -> Box<dyn Style> {
         self.stylesheet.get_style(match annotation_type {
             DisplayAnnotationType::Error => StyleClass::Error,
@@ -148,37 +145,38 @@ impl DisplayList {
         f: &mut fmt::Formatter<'_>,
     ) -> fmt::Result {
         let color = self.get_annotation_style(&annotation.annotation_type);
-
-        let formatted_type = if let Some(id) = &annotation.id {
-            DisplayFn::new(|f| {
-                self.format_annotation_type(&annotation.annotation_type, f)?;
-                f.write_char('[')?;
-                f.write_str(id)?;
-                f.write_char(']')
-            })
-            .to_string()
+        let formatted_len = if let Some(id) = &annotation.id {
+            2 + id.len() + Self::annotation_type_len(&annotation.annotation_type)
         } else {
-            DisplayFn::new(|f| self.format_annotation_type(&annotation.annotation_type, f))
-                .to_string()
+            Self::annotation_type_len(&annotation.annotation_type)
         };
 
         if continuation {
-            let indent = formatted_type.len() + 2;
-            format_repeat_char(' ', indent, f)?;
+            format_repeat_char(' ', formatted_len + 2, f)?;
             return self.format_label(&annotation.label, f);
         }
-        if formatted_type.is_empty() {
+        if formatted_len == 0 {
             self.format_label(&annotation.label, f)
         } else {
-            color.paint(&formatted_type, f)?;
+            color.paint_fn(
+                Box::new(|f| {
+                    Self::format_annotation_type(&annotation.annotation_type, f)?;
+                    if let Some(id) = &annotation.id {
+                        f.write_char('[')?;
+                        f.write_str(id)?;
+                        f.write_char(']')?;
+                    }
+                    Ok(())
+                }),
+                f,
+            )?;
             if !is_annotation_empty(annotation) {
                 if in_source {
-                    color.paint(
-                        &DisplayFn::new(|f| {
+                    color.paint_fn(
+                        Box::new(|f| {
                             f.write_str(": ")?;
                             self.format_label(&annotation.label, f)
-                        })
-                        .to_string(),
+                        }),
                         f,
                     )?;
                 } else {
@@ -230,21 +228,25 @@ impl DisplayList {
                     _ => range.0,
                 };
 
-                color.paint(&repeat_char(indent_char, indent_length + 1), f)?;
-                color.paint(&repeat_char(mark, range.1 - indent_length), f)?;
+                color.paint_fn(
+                    Box::new(|f| {
+                        format_repeat_char(indent_char, indent_length + 1, f)?;
+                        format_repeat_char(mark, range.1 - indent_length, f)
+                    }),
+                    f,
+                )?;
 
                 if !is_annotation_empty(&annotation) {
                     f.write_char(' ')?;
-                    color.paint(
-                        &DisplayFn::new(|f| {
+                    color.paint_fn(
+                        Box::new(|f| {
                             self.format_annotation(
                                 annotation,
                                 annotation_part == &DisplayAnnotationPart::LabelContinuation,
                                 true,
                                 f,
                             )
-                        })
-                        .to_string(),
+                        }),
                         f,
                     )?;
                 }
@@ -254,14 +256,6 @@ impl DisplayList {
         }
     }
 
-    #[inline]
-    fn format_lineno(&self, lineno: Option<usize>, lineno_width: usize) -> String {
-        match lineno {
-            Some(n) => format!("{:>width$}", n, width = lineno_width),
-            None => repeat_char(' ', lineno_width),
-        }
-    }
-
     #[inline]
     fn format_raw_line(
         &self,
@@ -337,10 +331,22 @@ impl DisplayList {
             } => {
                 let lineno_color = self.stylesheet.get_style(StyleClass::LineNo);
                 if self.anonymized_line_numbers && lineno.is_some() {
-                    lineno_color.paint(&format!("{} |", Self::ANONYMIZED_LINE_NUM), f)?;
+                    lineno_color.paint_fn(
+                        Box::new(|f| {
+                            f.write_str(Self::ANONYMIZED_LINE_NUM)?;
+                            f.write_str(" |")
+                        }),
+                        f,
+                    )?;
                 } else {
-                    lineno_color.paint(
-                        &format!("{} |", self.format_lineno(*lineno, lineno_width)),
+                    lineno_color.paint_fn(
+                        Box::new(|f| {
+                            match lineno {
+                                Some(n) => write!(f, "{:>width$}", n, width = lineno_width),
+                                None => format_repeat_char(' ', lineno_width, f),
+                            }?;
+                            f.write_str(" |")
+                        }),
                         f,
                     )?;
                 }
@@ -376,11 +382,13 @@ impl DisplayList {
     ) -> fmt::Result {
         format_repeat_char(' ', inline_marks_width - inline_marks.len(), f)?;
         for mark in inline_marks {
-            self.get_annotation_style(&mark.annotation_type).paint(
-                match mark.mark_type {
-                    DisplayMarkType::AnnotationThrough => "|",
-                    DisplayMarkType::AnnotationStart => "/",
-                },
+            self.get_annotation_style(&mark.annotation_type).paint_fn(
+                Box::new(|f| {
+                    f.write_char(match mark.mark_type {
+                        DisplayMarkType::AnnotationThrough => '|',
+                        DisplayMarkType::AnnotationStart => '/',
+                    })
+                }),
                 f,
             )?;
         }
diff --git a/src/formatter/style.rs b/src/formatter/style.rs
index c011ce3..f76e6b0 100644
--- a/src/formatter/style.rs
+++ b/src/formatter/style.rs
@@ -1,5 +1,4 @@
-//! Set of structures required to implement a stylesheet for
-//! [DisplayListFormatter](super::DisplayListFormatter).
+//! Set of structures required to implement a stylesheet
 //!
 //! In order to provide additional styling information for the
 //! formatter, a structs can implement `Stylesheet` and `Style`
@@ -8,7 +7,6 @@
 use std::fmt;
 
 /// StyleClass is a collection of named variants of style classes
-/// that DisplayListFormatter uses.
 pub enum StyleClass {
     /// Message indicating an error.
     Error,
@@ -33,8 +31,14 @@ pub enum StyleClass {
 
 /// This trait implements a return value for the `Stylesheet::get_style`.
 pub trait Style {
-    /// The method used by the DisplayListFormatter to style the message.
+    /// The method used to write text with formatter
     fn paint(&self, text: &str, f: &mut fmt::Formatter<'_>) -> fmt::Result;
+    /// The method used to write display function with formatter
+    fn paint_fn<'a>(
+        &self,
+        c: Box<dyn FnOnce(&mut fmt::Formatter<'_>) -> fmt::Result + 'a>,
+        f: &mut fmt::Formatter<'_>,
+    ) -> fmt::Result;
     /// The method used by the DisplayListFormatter to display the message
     /// in bold font.
     fn bold(&self) -> Box<dyn Style>;
diff --git a/src/stylesheets/color.rs b/src/stylesheets/color.rs
index c22f07e..024dd06 100644
--- a/src/stylesheets/color.rs
+++ b/src/stylesheets/color.rs
@@ -1,6 +1,6 @@
-use std::fmt;
+use std::fmt::{self, Display};
 
-use ansi_term::{Color::Fixed, Style as AnsiTermStyle};
+use yansi_term::{Color::Fixed, Style as AnsiTermStyle};
 
 use crate::formatter::style::{Style, StyleClass, Stylesheet};
 
@@ -10,7 +10,15 @@ struct AnsiTermStyleWrapper {
 
 impl Style for AnsiTermStyleWrapper {
     fn paint(&self, text: &str, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        fmt::Display::fmt(&self.style.paint(text), f)
+        self.style.paint(text).fmt(f)
+    }
+
+    fn paint_fn<'a>(
+        &self,
+        c: Box<dyn FnOnce(&mut fmt::Formatter<'_>) -> fmt::Result + 'a>,
+        f: &mut fmt::Formatter<'_>,
+    ) -> fmt::Result {
+        self.style.paint_fn(c).fmt(f)
     }
 
     fn bold(&self) -> Box<dyn Style> {
diff --git a/src/stylesheets/no_color.rs b/src/stylesheets/no_color.rs
index 1a5be0a..21cb269 100644
--- a/src/stylesheets/no_color.rs
+++ b/src/stylesheets/no_color.rs
@@ -9,6 +9,14 @@ impl Style for NoOpStyle {
         f.write_str(text)
     }
 
+    fn paint_fn<'a>(
+        &self,
+        c: Box<dyn FnOnce(&mut fmt::Formatter<'_>) -> fmt::Result + 'a>,
+        f: &mut fmt::Formatter<'_>,
+    ) -> fmt::Result {
+        c(f)
+    }
+
     fn bold(&self) -> Box<dyn Style> {
         Box::new(NoOpStyle {})
     }


Test patch:
diff --git a/tests/diff/mod.rs b/tests/diff/mod.rs
index ae847c1..576c6c4 100644
--- a/tests/diff/mod.rs
+++ b/tests/diff/mod.rs
@@ -1,5 +1,5 @@
-use ansi_term::Color::{Black, Green, Red};
 use difference::{Changeset, Difference};
+use yansi_term::Color::{Black, Green, Red};
 
 pub fn get_diff(left: &str, right: &str) -> String {
     let mut output = String::new();