From 3205e17cc379684baa7ec098add6929096d28644 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?St=C3=A9phane=20Campinas?= <stephane.campinas@gmail.com>
Date: Thu, 4 Oct 2018 13:56:57 +0200
Subject: [PATCH] rewrite_string: handle newlines in the last line that fits in
 the shape for cases where line_start is not a whitespace

---
 src/string.rs | 64 ++++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 51 insertions(+), 13 deletions(-)

diff --git a/src/string.rs b/src/string.rs
index 7589dca1bf95..39a0596b9a9d 100644
--- a/src/string.rs
+++ b/src/string.rs
@@ -96,15 +96,35 @@ pub fn rewrite_string<'a>(orig: &str, fmt: &StringFormat<'a>) -> Option<String>
     // Snip a line at a time from `stripped_str` until it is used up. Push the snippet
     // onto result.
     let mut cur_max_chars = max_chars_with_indent;
-    let is_overflow_allowed = is_whitespace(fmt.line_start);
+    let is_bareline_ok = fmt.line_start.is_empty() || is_whitespace(fmt.line_start);
     loop {
         // All the input starting at cur_start fits on the current line
         if graphemes.len() - cur_start <= cur_max_chars {
-            let last_line = graphemes[cur_start..].join("");
-            if fmt.trim_end {
-                result.push_str(&last_line.trim_right());
+            // trim trailing whitespaces
+            let graphemes_minus_ws = if !fmt.trim_end {
+                &graphemes[cur_start..]
             } else {
-                result.push_str(&last_line);
+                match graphemes[cur_start..]
+                    .iter()
+                    .rposition(|grapheme| !is_whitespace(grapheme))
+                {
+                    Some(index) => &graphemes[cur_start..=cur_start + index],
+                    None => &graphemes[cur_start..],
+                }
+            };
+            if is_bareline_ok {
+                // new lines don't need to start with line_start
+                result.push_str(&graphemes_minus_ws.join(""));
+            } else {
+                // new lines need to be indented and prefixed with line_start
+                for grapheme in graphemes_minus_ws {
+                    if is_line_feed(grapheme) {
+                        result.push_str(&indent_with_newline);
+                        result.push_str(fmt.line_start);
+                    } else {
+                        result.push_str(grapheme);
+                    }
+                }
             }
             break;
         }
@@ -121,7 +141,7 @@ pub fn rewrite_string<'a>(orig: &str, fmt: &StringFormat<'a>) -> Option<String>
             }
             SnippetState::EndWithLineFeed(line, len) => {
                 result.push_str(&line);
-                if is_overflow_allowed {
+                if is_bareline_ok {
                     // the next line can benefit from the full width
                     cur_max_chars = max_chars_without_indent;
                 } else {
@@ -171,13 +191,10 @@ fn break_string(max_chars: usize, trim_end: bool, input: &[&str]) -> SnippetStat
     let break_at = |index /* grapheme at index is included */| {
         // Take in any whitespaces to the left/right of `input[index]` and
         // check if there is a line feed, in which case whitespaces needs to be kept.
-        let mut index_minus_ws = index;
-        for (i, grapheme) in input[0..=index].iter().enumerate().rev() {
-            if !is_whitespace(grapheme) {
-                index_minus_ws = i;
-                break;
-            }
-        }
+        let index_minus_ws = input[0..=index]
+            .iter()
+            .rposition(|grapheme| !is_whitespace(grapheme))
+            .unwrap_or(index);
         // Take into account newlines occuring in input[0..=index], i.e., the possible next new
         // line. If there is one, then text after it could be rewritten in a way that the available
         // space is fully used.
@@ -398,6 +415,27 @@ mod test {
         assert_eq!(rewritten_string, Some("\"Vivamus id mi.  \"".to_string()));
     }
 
+    #[test]
+    fn last_line_fit_with_newline() {
+        let string = "Vivamus id mi.\nVivamus id mi.";
+        let config: Config = Default::default();
+        let fmt = StringFormat {
+            opener: "",
+            closer: "",
+            line_start: "// ",
+            line_end: "",
+            shape: Shape::legacy(100, Indent::from_width(&config, 4)),
+            trim_end: true,
+            config: &config,
+        };
+
+        let rewritten_string = rewrite_string(string, &fmt);
+        assert_eq!(
+            rewritten_string,
+            Some("Vivamus id mi.\n    // Vivamus id mi.".to_string())
+        );
+    }
+
     #[test]
     fn overflow_in_non_string_content() {
         let comment = "Aenean metus.\nVestibulum ac lacus. Vivamus porttitor";