fix(formatter): account lambda beginning when formatting method call

tooling/nargo_fmt/src/chunks.rs

@@ -6,7 +6,7 @@
 //! <https://yorickpeterse.com/articles/how-to-write-a-code-formatter/>
 //!
 //! 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<Chunk>,
@@ -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<usize> {
+        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

tooling/nargo_fmt/src/formatter/expression.rs

@@ -98,9 +98,9 @@
                     false, // force multiple lines
                 ));
             }
             ExpressionKind::Unsafe(unsafe_xpression) => {
                 group.group(self.format_unsafe_expression(
                     unsafe_xpression.block,
                     false, // force multiple lines
                 ));
             }
@@ -271,15 +271,10 @@
 
         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 @@
     }
 
     #[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)
@@ -1907,9 +1902,9 @@
 
     #[test]
     fn format_method_call_chain_3() {
         let src = "fn foo() {     assert(p4_affine.eq(Gaffine::new(6890855772600357754907169075114257697580319025794532037257385534741338397365, 4338620300185947561074059802482547481416142213883829469920100239455078257889)));  }";
         let expected = "fn foo() {
     assert(p4_affine.eq(Gaffine::new(
         6890855772600357754907169075114257697580319025794532037257385534741338397365,
         4338620300185947561074059802482547481416142213883829469920100239455078257889,
     )));
@@ -1945,7 +1940,7 @@
     fn format_nested_method_call_with_maximum_width_2() {
         let src = "fn foo() {
     assert(
         p4_affine.eq(Gaffine::new(
             6890855772600357754907169075114257697580319025794532037257385534741338397365,
             4338620300185947561074059802482547481416142213883829469920100239455078257889,
         )),
@@ -1953,7 +1948,7 @@
 }
 ";
         let expected = "fn foo() {
     assert(p4_affine.eq(Gaffine::new(
         6890855772600357754907169075114257697580319025794532037257385534741338397365,
         4338620300185947561074059802482547481416142213883829469920100239455078257889,
     )));
@@ -2545,4 +2540,49 @@
 "#;
         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);
+    }
 }