impl rewrite_result for ControlFlow, Stmt, update rewrite_index

6 files changed, 101 insertions, 60 deletions

diff --git a/src/comment.rs b/src/comment.rs
index 5bf3c1a725c..c8cadf364da 100644
--- a/src/comment.rs
+++ b/src/comment.rs
@@ -1703,12 +1703,11 @@ impl<'a> Iterator for CommentCodeSlices<'a> {
 }
 
 /// Checks is `new` didn't miss any comment from `span`, if it removed any, return previous text
-/// (if it fits in the width/offset, else return `None`), else return `new`
 pub(crate) fn recover_comment_removed(
     new: String,
     span: Span,
     context: &RewriteContext<'_>,
-) -> Option<String> {
+) -> String {
     let snippet = context.snippet(span);
     if snippet != new && changed_comment_content(snippet, &new) {
         // We missed some comments. Warn and keep the original text.
@@ -1722,9 +1721,9 @@ pub(crate) fn recover_comment_removed(
                 )],
             );
         }
-        Some(snippet.to_owned())
+        snippet.to_owned()
     } else {
-        Some(new)
+        new
     }
 }
 
diff --git a/src/expr.rs b/src/expr.rs
index 35b7bada92d..8381e3ab083 100644
--- a/src/expr.rs
+++ b/src/expr.rs
@@ -278,7 +278,7 @@ pub(crate) fn format_expr(
         )
         .ok(),
         ast::ExprKind::Index(ref expr, ref index, _) => {
-            rewrite_index(&**expr, &**index, context, shape)
+            rewrite_index(&**expr, &**index, context, shape).ok()
         }
         ast::ExprKind::Repeat(ref expr, ref repeats) => rewrite_pair(
             &**expr,
@@ -435,7 +435,7 @@ pub(crate) fn format_expr(
     };
 
     expr_rw
-        .and_then(|expr_str| recover_comment_removed(expr_str, expr.span, context))
+        .map(|expr_str| recover_comment_removed(expr_str, expr.span, context))
         .and_then(|expr_str| {
             let attrs = outer_attributes(&expr.attrs);
             let attrs_str = attrs.rewrite(context, shape)?;
@@ -672,6 +672,7 @@ pub(crate) fn rewrite_cond(
                 String::from("\n") + &shape.indent.block_only().to_string(context.config);
             control_flow
                 .rewrite_cond(context, shape, &alt_block_sep)
+                .ok()
                 .map(|rw| rw.0)
         }),
     }
@@ -896,10 +897,12 @@ impl<'a> ControlFlow<'a> {
         expr: &ast::Expr,
         shape: Shape,
         offset: usize,
-    ) -> Option<String> {
+    ) -> RewriteResult {
         debug!("rewrite_pat_expr {:?} {:?} {:?}", shape, self.pat, expr);
 
-        let cond_shape = shape.offset_left(offset)?;
+        let cond_shape = shape
+            .offset_left(offset)
+            .max_width_error(shape.width, expr.span)?;
         if let Some(pat) = self.pat {
             let matcher = if self.matcher.is_empty() {
                 self.matcher.to_owned()
@@ -907,9 +910,10 @@ impl<'a> ControlFlow<'a> {
                 format!("{} ", self.matcher)
             };
             let pat_shape = cond_shape
-                .offset_left(matcher.len())?
-                .sub_width(self.connector.len())?;
-            let pat_string = pat.rewrite(context, pat_shape)?;
+                .offset_left(matcher.len())
+                .and_then(|s| s.sub_width(self.connector.len()))
+                .max_width_error(cond_shape.width, pat.span)?;
+            let pat_string = pat.rewrite_result(context, pat_shape)?;
             let comments_lo = context
                 .snippet_provider
                 .span_after(self.span.with_lo(pat.span.hi()), self.connector.trim());
@@ -923,14 +927,13 @@ impl<'a> ControlFlow<'a> {
                 RhsTactics::Default,
                 comments_span,
                 true,
-            )
-            .ok();
+            );
         }
 
-        let expr_rw = expr.rewrite(context, cond_shape);
+        let expr_rw = expr.rewrite_result(context, cond_shape);
         // The expression may (partially) fit on the current line.
         // We do not allow splitting between `if` and condition.
-        if self.keyword == "if" || expr_rw.is_some() {
+        if self.keyword == "if" || expr_rw.is_ok() {
             return expr_rw;
         }
 
@@ -939,7 +942,7 @@ impl<'a> ControlFlow<'a> {
             .block_indent(context.config.tab_spaces())
             .with_max_width(context.config);
         let nested_indent_str = nested_shape.indent.to_string_with_newline(context.config);
-        expr.rewrite(context, nested_shape)
+        expr.rewrite_result(context, nested_shape)
             .map(|expr_rw| format!("{}{}", nested_indent_str, expr_rw))
     }
 
@@ -948,7 +951,7 @@ impl<'a> ControlFlow<'a> {
         context: &RewriteContext<'_>,
         shape: Shape,
         alt_block_sep: &str,
-    ) -> Option<(String, usize)> {
+    ) -> Result<(String, usize), RewriteError> {
         // Do not take the rhs overhead from the upper expressions into account
         // when rewriting pattern.
         let new_width = context.budget(shape.used_width());
@@ -959,7 +962,9 @@ impl<'a> ControlFlow<'a> {
         let constr_shape = if self.nested_if {
             // We are part of an if-elseif-else chain. Our constraints are tightened.
             // 7 = "} else " .len()
-            fresh_shape.offset_left(7)?
+            fresh_shape
+                .offset_left(7)
+                .max_width_error(fresh_shape.width, self.span)?
         } else {
             fresh_shape
         };
@@ -995,7 +1000,7 @@ impl<'a> ControlFlow<'a> {
 
             if let Some(cond_str) = trial {
                 if cond_str.len() <= context.config.single_line_if_else_max_width() {
-                    return Some((cond_str, 0));
+                    return Ok((cond_str, 0));
                 }
             }
         }
@@ -1048,7 +1053,7 @@ impl<'a> ControlFlow<'a> {
             label_string.len() + self.keyword.len() + pat_expr_string.len() + 2
         };
 
-        Some((
+        Ok((
             format!(
                 "{}{}{}{}{}",
                 label_string,
@@ -1114,13 +1119,17 @@ pub(crate) fn rewrite_else_kw_with_comments(
 
 impl<'a> Rewrite for ControlFlow<'a> {
     fn rewrite(&self, context: &RewriteContext<'_>, shape: Shape) -> Option<String> {
+        self.rewrite_result(context, shape).ok()
+    }
+
+    fn rewrite_result(&self, context: &RewriteContext<'_>, shape: Shape) -> RewriteResult {
         debug!("ControlFlow::rewrite {:?} {:?}", self, shape);
 
         let alt_block_sep = &shape.indent.to_string_with_newline(context.config);
         let (cond_str, used_width) = self.rewrite_cond(context, shape, alt_block_sep)?;
         // If `used_width` is 0, it indicates that whole control flow is written in a single line.
         if used_width == 0 {
-            return Some(cond_str);
+            return Ok(cond_str);
         }
 
         let block_width = shape.width.saturating_sub(used_width);
@@ -1138,8 +1147,7 @@ impl<'a> Rewrite for ControlFlow<'a> {
         let block_str = {
             let old_val = context.is_if_else_block.replace(self.else_block.is_some());
             let result =
-                rewrite_block_with_visitor(context, "", self.block, None, None, block_shape, true)
-                    .ok();
+                rewrite_block_with_visitor(context, "", self.block, None, None, block_shape, true);
             context.is_if_else_block.replace(old_val);
             result?
         };
@@ -1165,7 +1173,7 @@ impl<'a> Rewrite for ControlFlow<'a> {
                         true,
                         mk_sp(else_block.span.lo(), self.span.hi()),
                     )
-                    .rewrite(context, shape)
+                    .rewrite_result(context, shape)
                 }
                 _ => {
                     last_in_chain = true;
@@ -1176,6 +1184,7 @@ impl<'a> Rewrite for ControlFlow<'a> {
                         ..shape
                     };
                     format_expr(else_block, ExprType::Statement, context, else_shape)
+                        .unknown_error()
                 }
             };
 
@@ -1190,7 +1199,7 @@ impl<'a> Rewrite for ControlFlow<'a> {
             result.push_str(&rewrite?);
         }
 
-        Some(result)
+        Ok(result)
     }
 }
 
@@ -1567,8 +1576,8 @@ fn rewrite_index(
     index: &ast::Expr,
     context: &RewriteContext<'_>,
     shape: Shape,
-) -> Option<String> {
-    let expr_str = expr.rewrite(context, shape)?;
+) -> RewriteResult {
+    let expr_str = expr.rewrite_result(context, shape)?;
 
     let offset = last_line_width(&expr_str) + 1;
     let rhs_overhead = shape.rhs_overhead(context.config);
@@ -1583,37 +1592,45 @@ fn rewrite_index(
                 .and_then(|shape| shape.sub_width(1)),
             IndentStyle::Visual => shape.visual_indent(offset).sub_width(offset + 1),
         }
-    };
-    let orig_index_rw = index_shape.and_then(|s| index.rewrite(context, s));
+    }
+    .max_width_error(shape.width, index.span());
+    let orig_index_rw = index_shape.and_then(|s| index.rewrite_result(context, s));
 
     // Return if index fits in a single line.
     match orig_index_rw {
-        Some(ref index_str) if !index_str.contains('\n') => {
-            return Some(format!("{expr_str}[{index_str}]"));
+        Ok(ref index_str) if !index_str.contains('\n') => {
+            return Ok(format!("{expr_str}[{index_str}]"));
         }
         _ => (),
     }
 
     // Try putting index on the next line and see if it fits in a single line.
     let indent = shape.indent.block_indent(context.config);
-    let index_shape = Shape::indented(indent, context.config).offset_left(1)?;
-    let index_shape = index_shape.sub_width(1 + rhs_overhead)?;
-    let new_index_rw = index.rewrite(context, index_shape);
+    let index_shape = Shape::indented(indent, context.config)
+        .offset_left(1)
+        .max_width_error(shape.width, index.span())?;
+    let index_shape = index_shape
+        .sub_width(1 + rhs_overhead)
+        .max_width_error(index_shape.width, index.span())?;
+    let new_index_rw = index.rewrite_result(context, index_shape);
     match (orig_index_rw, new_index_rw) {
-        (_, Some(ref new_index_str)) if !new_index_str.contains('\n') => Some(format!(
+        (_, Ok(ref new_index_str)) if !new_index_str.contains('\n') => Ok(format!(
             "{}{}[{}]",
             expr_str,
             indent.to_string_with_newline(context.config),
             new_index_str,
         )),
-        (None, Some(ref new_index_str)) => Some(format!(
+        (Err(_), Ok(ref new_index_str)) => Ok(format!(
             "{}{}[{}]",
             expr_str,
             indent.to_string_with_newline(context.config),
             new_index_str,
         )),
-        (Some(ref index_str), _) => Some(format!("{expr_str}[{index_str}]")),
-        _ => None,
+        (Ok(ref index_str), _) => Ok(format!("{expr_str}[{index_str}]")),
+        // When both orig_index_rw and new_index_rw result in errors, we currently propagate the
+        // error from the second attempt since it is more generous with width constraints.
+        // This decision is somewhat arbitrary and is open to change.
+        (Err(_), Err(new_index_rw_err)) => Err(new_index_rw_err),
     }
 }
 
diff --git a/src/items.rs b/src/items.rs
index 4ea9751e816..89fb9bda9df 100644
--- a/src/items.rs
+++ b/src/items.rs
@@ -2077,7 +2077,7 @@ fn rewrite_static(
             true,
         )
         .ok()
-        .and_then(|res| recover_comment_removed(res, static_parts.span, context))
+        .map(|res| recover_comment_removed(res, static_parts.span, context))
         .map(|s| if s.ends_with(';') { s } else { s + ";" })
     } else {
         Some(format!("{prefix}{ty_str};"))
diff --git a/src/overflow.rs b/src/overflow.rs
index 510ad3c642e..dc4716a81af 100644
--- a/src/overflow.rs
+++ b/src/overflow.rs
@@ -19,7 +19,7 @@ use crate::lists::{
 };
 use crate::macros::MacroArg;
 use crate::patterns::{can_be_overflowed_pat, TuplePatField};
-use crate::rewrite::{Rewrite, RewriteContext, RewriteErrorExt, RewriteResult};
+use crate::rewrite::{Rewrite, RewriteContext, RewriteError, RewriteErrorExt, RewriteResult};
 use crate::shape::Shape;
 use crate::source_map::SpanUtils;
 use crate::spanned::Spanned;
@@ -90,6 +90,10 @@ impl<'a> Rewrite for OverflowableItem<'a> {
     fn rewrite(&self, context: &RewriteContext<'_>, shape: Shape) -> Option<String> {
         self.map(|item| item.rewrite(context, shape))
     }
+
+    fn rewrite_result(&self, context: &RewriteContext<'_>, shape: Shape) -> RewriteResult {
+        self.map(|item| item.rewrite_result(context, shape))
+    }
 }
 
 impl<'a> Spanned for OverflowableItem<'a> {
@@ -617,7 +621,7 @@ impl<'a> Context<'a> {
         tactic
     }
 
-    fn rewrite_items(&self) -> Option<(bool, String)> {
+    fn rewrite_items(&self) -> Result<(bool, String), RewriteError> {
         let span = self.items_span();
         debug!("items: {:?}", self.items);
 
@@ -661,7 +665,6 @@ impl<'a> Context<'a> {
             .ends_with_newline(ends_with_newline);
 
         write_list(&list_items, &fmt)
-            .ok()
             .map(|items_str| (tactic == DefinitiveListTactic::Horizontal, items_str))
     }
 
@@ -718,7 +721,7 @@ impl<'a> Context<'a> {
     }
 
     fn rewrite(&self, shape: Shape) -> RewriteResult {
-        let (extendable, items_str) = self.rewrite_items().unknown_error()?;
+        let (extendable, items_str) = self.rewrite_items()?;
 
         // If we are using visual indent style and failed to format, retry with block indent.
         if !self.context.use_block_indent()
diff --git a/src/stmt.rs b/src/stmt.rs
index 38433433c6b..4f31705295e 100644
--- a/src/stmt.rs
+++ b/src/stmt.rs
@@ -4,7 +4,7 @@ use rustc_span::Span;
 use crate::comment::recover_comment_removed;
 use crate::config::StyleEdition;
 use crate::expr::{format_expr, is_simple_block, ExprType};
-use crate::rewrite::{Rewrite, RewriteContext};
+use crate::rewrite::{Rewrite, RewriteContext, RewriteError, RewriteErrorExt, RewriteResult};
 use crate::shape::Shape;
 use crate::source_map::LineRangeUtils;
 use crate::spanned::Spanned;
@@ -90,6 +90,14 @@ impl<'a> Stmt<'a> {
 
 impl<'a> Rewrite for Stmt<'a> {
     fn rewrite(&self, context: &RewriteContext<'_>, shape: Shape) -> Option<String> {
+        self.rewrite_result(context, shape).ok()
+    }
+
+    fn rewrite_result(
+        &self,
+        context: &RewriteContext<'_>,
+        shape: Shape,
+    ) -> crate::rewrite::RewriteResult {
         let expr_type =
             if context.config.style_edition() >= StyleEdition::Edition2024 && self.is_last_expr() {
                 ExprType::SubExpression
@@ -112,11 +120,11 @@ fn format_stmt(
     stmt: &ast::Stmt,
     expr_type: ExprType,
     is_last_expr: bool,
-) -> Option<String> {
-    skip_out_of_file_lines_range!(context, stmt.span());
+) -> RewriteResult {
+    skip_out_of_file_lines_range_err!(context, stmt.span());
 
     let result = match stmt.kind {
-        ast::StmtKind::Let(ref local) => local.rewrite(context, shape),
+        ast::StmtKind::Let(ref local) => local.rewrite_result(context, shape),
         ast::StmtKind::Expr(ref ex) | ast::StmtKind::Semi(ref ex) => {
             let suffix = if semicolon_for_stmt(context, stmt, is_last_expr) {
                 ";"
@@ -124,10 +132,16 @@ fn format_stmt(
                 ""
             };
 
-            let shape = shape.sub_width(suffix.len())?;
-            format_expr(ex, expr_type, context, shape).map(|s| s + suffix)
+            let shape = shape
+                .sub_width(suffix.len())
+                .max_width_error(shape.width, ex.span())?;
+            format_expr(ex, expr_type, context, shape)
+                .map(|s| s + suffix)
+                .unknown_error()
+        }
+        ast::StmtKind::MacCall(..) | ast::StmtKind::Item(..) | ast::StmtKind::Empty => {
+            Err(RewriteError::Unknown)
         }
-        ast::StmtKind::MacCall(..) | ast::StmtKind::Item(..) | ast::StmtKind::Empty => None,
     };
-    result.and_then(|res| recover_comment_removed(res, stmt.span(), context))
+    result.map(|res| recover_comment_removed(res, stmt.span(), context))
 }
diff --git a/src/types.rs b/src/types.rs
index d1cd291d51f..af57a5843db 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -480,8 +480,8 @@ impl Rewrite for ast::WherePredicate {
             ast::WherePredicate::RegionPredicate(ast::WhereRegionPredicate {
                 ref lifetime,
                 ref bounds,
-                ..
-            }) => rewrite_bounded_lifetime(lifetime, bounds, context, shape)?,
+                span,
+            }) => rewrite_bounded_lifetime(lifetime, bounds, span, context, shape).ok()?,
             ast::WherePredicate::EqPredicate(ast::WhereEqPredicate {
                 ref lhs_ty,
                 ref rhs_ty,
@@ -553,23 +553,27 @@ fn rewrite_generic_args(
 fn rewrite_bounded_lifetime(
     lt: &ast::Lifetime,
     bounds: &[ast::GenericBound],
+    span: Span,
     context: &RewriteContext<'_>,
     shape: Shape,
-) -> Option<String> {
-    let result = lt.rewrite(context, shape)?;
+) -> RewriteResult {
+    let result = lt.rewrite_result(context, shape)?;
 
     if bounds.is_empty() {
-        Some(result)
+        Ok(result)
     } else {
         let colon = type_bound_colon(context);
         let overhead = last_line_width(&result) + colon.len();
+        let shape = shape
+            .sub_width(overhead)
+            .max_width_error(shape.width, span)?;
         let result = format!(
             "{}{}{}",
             result,
             colon,
-            join_bounds(context, shape.sub_width(overhead)?, bounds, true).ok()?
+            join_bounds(context, shape, bounds, true)?
         );
-        Some(result)
+        Ok(result)
     }
 }
 
@@ -580,8 +584,12 @@ impl Rewrite for ast::AnonConst {
 }
 
 impl Rewrite for ast::Lifetime {
-    fn rewrite(&self, context: &RewriteContext<'_>, _: Shape) -> Option<String> {
-        Some(rewrite_ident(context, self.ident).to_owned())
+    fn rewrite(&self, context: &RewriteContext<'_>, shape: Shape) -> Option<String> {
+        self.rewrite_result(context, shape).ok()
+    }
+
+    fn rewrite_result(&self, context: &RewriteContext<'_>, _: Shape) -> RewriteResult {
+        Ok(rewrite_ident(context, self.ident).to_owned())
     }
 }