Skip to content

Commit

Permalink
refactor rewrite_closure
Browse files Browse the repository at this point in the history
  • Loading branch information
ding-young authored and ytmimi committed Jul 15, 2024
1 parent bb56224 commit 071ca7d
Show file tree
Hide file tree
Showing 3 changed files with 51 additions and 36 deletions.
82 changes: 48 additions & 34 deletions src/closures.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use crate::expr::{block_contains_comment, is_simple_block, is_unsafe_block, rewr
use crate::items::{span_hi_for_param, span_lo_for_param};
use crate::lists::{definitive_tactic, itemize_list, write_list, ListFormatting, Separator};
use crate::overflow::OverflowableItem;
use crate::rewrite::{Rewrite, RewriteContext};
use crate::rewrite::{Rewrite, RewriteContext, RewriteError, RewriteErrorExt, RewriteResult};
use crate::shape::Shape;
use crate::source_map::SpanUtils;
use crate::types::rewrite_bound_params;
Expand All @@ -36,7 +36,7 @@ pub(crate) fn rewrite_closure(
span: Span,
context: &RewriteContext<'_>,
shape: Shape,
) -> Option<String> {
) -> RewriteResult {
debug!("rewrite_closure {:?}", body);

let (prefix, extra_offset) = rewrite_closure_fn_decl(
Expand All @@ -52,29 +52,31 @@ pub(crate) fn rewrite_closure(
shape,
)?;
// 1 = space between `|...|` and body.
let body_shape = shape.offset_left(extra_offset)?;
let body_shape = shape
.offset_left(extra_offset)
.max_width_error(shape.width, span)?;

if let ast::ExprKind::Block(ref block, _) = body.kind {
// The body of the closure is an empty block.
if block.stmts.is_empty() && !block_contains_comment(context, block) {
return body
.rewrite(context, shape)
.rewrite_result(context, shape)
.map(|s| format!("{} {}", prefix, s));
}

let result = match fn_decl.output {
ast::FnRetTy::Default(_) if !context.inside_macro() => {
try_rewrite_without_block(body, &prefix, context, shape, body_shape)
}
_ => None,
_ => Err(RewriteError::Unknown),
};

result.or_else(|| {
result.or_else(|_| {
// Either we require a block, or tried without and failed.
rewrite_closure_block(block, &prefix, context, body_shape)
})
} else {
rewrite_closure_expr(body, &prefix, context, body_shape).or_else(|| {
rewrite_closure_expr(body, &prefix, context, body_shape).or_else(|_| {
// The closure originally had a non-block expression, but we can't fit on
// one line, so we'll insert a block.
rewrite_closure_with_block(body, &prefix, context, body_shape)
Expand All @@ -88,7 +90,7 @@ fn try_rewrite_without_block(
context: &RewriteContext<'_>,
shape: Shape,
body_shape: Shape,
) -> Option<String> {
) -> RewriteResult {
let expr = get_inner_expr(expr, prefix, context);

if is_block_closure_forced(context, expr) {
Expand Down Expand Up @@ -152,11 +154,11 @@ fn rewrite_closure_with_block(
prefix: &str,
context: &RewriteContext<'_>,
shape: Shape,
) -> Option<String> {
) -> RewriteResult {
let left_most = left_most_sub_expr(body);
let veto_block = veto_block(body) && !expr_requires_semi_to_be_stmt(left_most);
if veto_block {
return None;
return Err(RewriteError::Unknown);
}

let block = ast::Block {
Expand All @@ -183,9 +185,8 @@ fn rewrite_closure_with_block(
None,
shape,
false,
)
.ok()?;
Some(format!("{prefix} {block}"))
)?;
Ok(format!("{prefix} {block}"))
}

// Rewrite closure with a single expression without wrapping its body with block.
Expand All @@ -194,7 +195,7 @@ fn rewrite_closure_expr(
prefix: &str,
context: &RewriteContext<'_>,
shape: Shape,
) -> Option<String> {
) -> RewriteResult {
fn allow_multi_line(expr: &ast::Expr) -> bool {
match expr.kind {
ast::ExprKind::Match(..)
Expand All @@ -217,12 +218,12 @@ fn rewrite_closure_expr(
// unless it is a block-like expression or we are inside macro call.
let veto_multiline = (!allow_multi_line(expr) && !context.inside_macro())
|| context.config.force_multiline_blocks();
expr.rewrite(context, shape)
expr.rewrite_result(context, shape)
.and_then(|rw| {
if veto_multiline && rw.contains('\n') {
None
Err(RewriteError::Unknown)
} else {
Some(rw)
Ok(rw)
}
})
.map(|rw| format!("{} {}", prefix, rw))
Expand All @@ -234,8 +235,12 @@ fn rewrite_closure_block(
prefix: &str,
context: &RewriteContext<'_>,
shape: Shape,
) -> Option<String> {
Some(format!("{} {}", prefix, block.rewrite(context, shape)?))
) -> RewriteResult {
Ok(format!(
"{} {}",
prefix,
block.rewrite_result(context, shape)?
))
}

// Return type is (prefix, extra_offset)
Expand All @@ -250,13 +255,14 @@ fn rewrite_closure_fn_decl(
span: Span,
context: &RewriteContext<'_>,
shape: Shape,
) -> Option<(String, usize)> {
) -> Result<(String, usize), RewriteError> {
let binder = match binder {
ast::ClosureBinder::For { generic_params, .. } if generic_params.is_empty() => {
"for<> ".to_owned()
}
ast::ClosureBinder::For { generic_params, .. } => {
let lifetime_str = rewrite_bound_params(context, shape, generic_params)?;
let lifetime_str =
rewrite_bound_params(context, shape, generic_params).unknown_error()?;
format!("for<{lifetime_str}> ")
}
ast::ClosureBinder::NotPresent => "".to_owned(),
Expand Down Expand Up @@ -287,13 +293,17 @@ fn rewrite_closure_fn_decl(
// 4 = "|| {".len(), which is overconservative when the closure consists of
// a single expression.
let nested_shape = shape
.shrink_left(binder.len() + const_.len() + immovable.len() + coro.len() + mover.len())?
.sub_width(4)?;
.shrink_left(binder.len() + const_.len() + immovable.len() + coro.len() + mover.len())
.and_then(|shape| shape.sub_width(4))
.max_width_error(shape.width, span)?;

// 1 = |
let param_offset = nested_shape.indent + 1;
let param_shape = nested_shape.offset_left(1)?.visual_indent(0);
let ret_str = fn_decl.output.rewrite(context, param_shape)?;
let param_shape = nested_shape
.offset_left(1)
.max_width_error(nested_shape.width, span)?
.visual_indent(0);
let ret_str = fn_decl.output.rewrite_result(context, param_shape)?;

let param_items = itemize_list(
context.snippet_provider,
Expand All @@ -317,14 +327,16 @@ fn rewrite_closure_fn_decl(
horizontal_budget,
);
let param_shape = match tactic {
DefinitiveListTactic::Horizontal => param_shape.sub_width(ret_str.len() + 1)?,
DefinitiveListTactic::Horizontal => param_shape
.sub_width(ret_str.len() + 1)
.max_width_error(param_shape.width, span)?,
_ => param_shape,
};

let fmt = ListFormatting::new(param_shape, context.config)
.tactic(tactic)
.preserve_newline(true);
let list_str = write_list(&item_vec, &fmt)?;
let list_str = write_list(&item_vec, &fmt).unknown_error()?;
let mut prefix = format!("{binder}{const_}{immovable}{coro}{mover}|{list_str}|");

if !ret_str.is_empty() {
Expand All @@ -339,7 +351,7 @@ fn rewrite_closure_fn_decl(
// 1 = space between `|...|` and body.
let extra_offset = last_line_width(&prefix) + 1;

Some((prefix, extra_offset))
Ok((prefix, extra_offset))
}

// Rewriting closure which is placed at the end of the function call's arg.
Expand All @@ -348,7 +360,7 @@ pub(crate) fn rewrite_last_closure(
context: &RewriteContext<'_>,
expr: &ast::Expr,
shape: Shape,
) -> Option<String> {
) -> RewriteResult {
if let ast::ExprKind::Closure(ref closure) = expr.kind {
let ast::Closure {
ref binder,
Expand Down Expand Up @@ -385,10 +397,12 @@ pub(crate) fn rewrite_last_closure(
)?;
// If the closure goes multi line before its body, do not overflow the closure.
if prefix.contains('\n') {
return None;
return Err(RewriteError::Unknown);
}

let body_shape = shape.offset_left(extra_offset)?;
let body_shape = shape
.offset_left(extra_offset)
.max_width_error(shape.width, expr.span)?;

// We force to use block for the body of the closure for certain kinds of expressions.
if is_block_closure_forced(context, body) {
Expand All @@ -400,7 +414,7 @@ pub(crate) fn rewrite_last_closure(
// closure. However, if the closure has a return type, then we must
// keep the blocks.
match rewrite_closure_expr(body, &prefix, context, shape) {
Some(single_line_body_str)
Ok(single_line_body_str)
if !single_line_body_str.contains('\n') =>
{
single_line_body_str
Expand All @@ -424,9 +438,9 @@ pub(crate) fn rewrite_last_closure(
}

// Seems fine, just format the closure in usual manner.
return expr.rewrite(context, shape);
return expr.rewrite_result(context, shape);
}
None
Err(RewriteError::Unknown)
}

/// Returns `true` if the given vector of arguments has more than one `ast::ExprKind::Closure`.
Expand Down
3 changes: 2 additions & 1 deletion src/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -230,7 +230,8 @@ pub(crate) fn format_expr(
expr.span,
context,
shape,
),
)
.ok(),
ast::ExprKind::Try(..)
| ast::ExprKind::Field(..)
| ast::ExprKind::MethodCall(..)
Expand Down
2 changes: 1 addition & 1 deletion src/overflow.rs
Original file line number Diff line number Diff line change
Expand Up @@ -427,7 +427,7 @@ impl<'a> Context<'a> {
if closures::args_have_many_closure(&self.items) {
None
} else {
closures::rewrite_last_closure(self.context, expr, shape)
closures::rewrite_last_closure(self.context, expr, shape).ok()
}
}

Expand Down

0 comments on commit 071ca7d

Please sign in to comment.