diff --git a/tooling/nargo_fmt/src/chunks.rs b/tooling/nargo_fmt/src/chunks.rs index 0e8429e1971..84673a23e6b 100644 --- a/tooling/nargo_fmt/src/chunks.rs +++ b/tooling/nargo_fmt/src/chunks.rs @@ -6,7 +6,7 @@ //! //! //! However, some changes were introduces to handle comments and other particularities of Noir. -use std::ops::Deref; +use std::{fmt::Display, ops::Deref}; use noirc_frontend::token::Token; @@ -113,6 +113,26 @@ impl Chunk { } } +impl Display for Chunk { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Chunk::Text(text_chunk, _) + | Chunk::TrailingComment(text_chunk) + | Chunk::LeadingComment(text_chunk) => { + write!(f, "{}", text_chunk.string) + } + Chunk::TrailingComma => write!(f, ","), + Chunk::Group(chunk_group) => chunk_group.fmt(f), + Chunk::SpaceOrLine => write!(f, " "), + Chunk::Line { .. } + | Chunk::IncreaseIndentation + | Chunk::DecreaseIndentation + | Chunk::PushIndentation + | Chunk::PopIndentation => Ok(()), + } + } +} + #[derive(Debug)] pub(crate) struct ChunkGroup { pub(crate) chunks: Vec, @@ -433,6 +453,44 @@ impl ChunkGroup { false } + + /// Assuming this is a MethodCall group, if the ExpressionList nested in it + /// has a LambdaAsLastExpressionInList, returns its `first_line_width`. + fn method_call_lambda_first_line_width(&self) -> Option { + for chunk in &self.chunks { + let Chunk::Group(group) = chunk else { + continue; + }; + + let GroupKind::ExpressionList { expressions_count: 1, .. } = group.kind else { + continue; + }; + + for chunk in &group.chunks { + let Chunk::Group(group) = chunk else { + continue; + }; + + let GroupKind::LambdaAsLastExpressionInList { first_line_width, .. } = group.kind + else { + continue; + }; + + return Some(first_line_width); + } + } + + None + } +} + +impl Display for ChunkGroup { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + for chunk in &self.chunks { + chunk.fmt(f)?; + } + Ok(()) + } } #[derive(Debug, Copy, Clone, PartialEq, Eq)] @@ -610,10 +668,24 @@ impl<'a> Formatter<'a> { lhs: false, } = group.kind { - let total_width = self.current_line_width() + width_until_left_paren_inclusive; + let mut total_width = self.current_line_width() + width_until_left_paren_inclusive; + + if total_width <= self.max_width { + // Check if this method call has a single lambda argument, like this: + // + // foo.bar.baz(|x| { x + 1 }) + // + // In this case we can't just consider the width up to `(`, we need to include `|x| {` + // because we'll want to format it in the same line as the `(` (we could put the lambda + // on a separate line but it would be less readable). + if let Some(first_line_width) = group.method_call_lambda_first_line_width() { + total_width += first_line_width; + } + } + if total_width <= self.max_width { // Check if this method call has another call or method call nested in it. - // If not, it means tis is the last nested call and after it we'll need to start + // If not, it means this is the last nested call and after it we'll need to start // writing at least one closing parentheses. So the argument list will actually // have one less character available for writing, and that's why we (temporarily) decrease // max width. @@ -681,7 +753,6 @@ impl<'a> Formatter<'a> { return; } - // if chunks.has_newlines() { // When formatting an expression list we have to check if the last argument is a lambda, // because we format that in a special way: // 1. to compute the group width we'll consider only the `|...| {` part of the lambda diff --git a/tooling/nargo_fmt/src/formatter/expression.rs b/tooling/nargo_fmt/src/formatter/expression.rs index 49697c6ce00..5c95925324b 100644 --- a/tooling/nargo_fmt/src/formatter/expression.rs +++ b/tooling/nargo_fmt/src/formatter/expression.rs @@ -271,15 +271,10 @@ impl ChunkFormatter<'_, '_> { group.group(body_group); - let first_line_width = params_and_return_type_chunk_width - + (if block_statement_count.is_some() { - // 1 because we already have `|param1, param2, ..., paramN| ` (including the space) - // so all that's left is a `{`. - 1 - } else { - // The body is not a block so we can write it right away - 0 - }); + // We assume that if we'll split the lambda into multiple lines, we'll always + // format the lambda expression as a block, so we add 1 to account the added `{`, + // so `first_line_width` is the width of `|...| {`. + let first_line_width = params_and_return_type_chunk_width + 1; FormattedLambda { group, first_line_width } } @@ -1882,7 +1877,7 @@ global y = 1; } #[test] - fn format_method_call_chain() { + fn format_method_call_chain_1() { let src = "global x = bar . baz ( 1, 2 ) . qux ( 1 , 2, 3) . one ( 5, 6) ;"; let expected = "global x = bar .baz(1, 2) @@ -2545,4 +2540,49 @@ global y = 1; "#; assert_format_with_max_width(src, expected, 30); } + + #[test] + fn lambda_in_method_call_that_exceeds_max_width_because_of_curly_1() { + let src = r#"fn foo() { + let _ = foo.bar().x(|y| { + y.qux() + }); +} +"#; + let expected = r#"fn foo() { + let _ = foo + .bar() + .x(|y| y.qux()); +} +"#; + assert_format_with_max_width(src, expected, 28); + } + + #[test] + fn lambda_in_method_call_that_exceeds_max_width_because_of_curly_2() { + let src = r#"fn foo() { + let _ = foo.bar().x( + |y| y.qux(), + ); +} +"#; + let expected = r#"fn foo() { + let _ = foo + .bar() + .x(|y| y.qux()); +} +"#; + assert_format_with_max_width(src, expected, 28); + } + + #[test] + fn lambda_in_method_call_that_exceeds_max_width_because_of_curly_3() { + let src = r#"fn foo() { + let _ = foo + .bar() + .x(|y| y.qux()); +} +"#; + assert_format_with_max_width(src, src, 28); + } }