Add an optional max width argument to rewrite_string()

1 files changed, 11 insertions, 3 deletions

diff --git a/src/string.rs b/src/string.rs
index c3f1c5faf1e..43b1ccbcb93 100644
--- a/src/string.rs
+++ b/src/string.rs
@@ -44,7 +44,11 @@ impl<'a> StringFormat<'a> {
 }
 
 // FIXME: simplify this!
-pub fn rewrite_string<'a>(orig: &str, fmt: &StringFormat<'a>) -> Option<String> {
+pub fn rewrite_string<'a>(
+    orig: &str,
+    fmt: &StringFormat<'a>,
+    max_width: Option<usize>,
+) -> Option<String> {
     // Strip line breaks.
     let re = Regex::new(r"([^\\](\\\\)*)\\[\n\r][[:space:]]*").unwrap();
     let stripped_str = re.replace_all(orig, "$1");
@@ -67,7 +71,7 @@ pub fn rewrite_string<'a>(orig: &str, fmt: &StringFormat<'a>) -> Option<String>
     let ender_length = fmt.line_end.len();
     // If we cannot put at least a single character per line, the rewrite won't
     // succeed.
-    let max_chars = shape
+    let mut max_chars = shape
         .width
         .checked_sub(fmt.opener.len() + ender_length + 1)? + 1;
 
@@ -135,6 +139,10 @@ pub fn rewrite_string<'a>(orig: &str, fmt: &StringFormat<'a>) -> Option<String>
 
         // The next line starts where the current line ends.
         cur_start = cur_end;
+
+        if let Some(new_max_chars) = max_width {
+            max_chars = new_max_chars.checked_sub(fmt.opener.len() + ender_length + 1)? + 1;
+        }
     }
 
     result.push_str(fmt.closer);
@@ -150,6 +158,6 @@ mod test {
     fn issue343() {
         let config = Default::default();
         let fmt = StringFormat::new(Shape::legacy(2, Indent::empty()), &config);
-        rewrite_string("eq_", &fmt);
+        rewrite_string("eq_", &fmt, None);
     }
 }