feat: let the formatter remove lambda block braces for single-statement blocks

test_programs/compile_success_empty/arithmetic_generics/src/main.nr

@@ -77,25 +77,25 @@ fn equiv_trans<T, TU, U, UT, UV, V, VU>(
     x: Equiv<T, TU, U, UT>,
     y: Equiv<U, UV, V, VU>,
     ) -> Equiv<T, (Equiv<U, UV, V, VU>, Equiv<T, TU, U, UT>), V, (Equiv<T, TU, U, UT>, Equiv<U, UV, V, VU>)> {
-    Equiv { to_: |z| { y.to(x.to(z)) }, fro_: |z| { x.fro(y.fro(z)) } }
+    Equiv { to_: |z| y.to(x.to(z)), fro_: |z| x.fro(y.fro(z)) }
 }
 
 fn mul_one_r<let N: u32>() -> Equiv<W<N>, (), W<N>, ()> {
-    Equiv { to_: |_x| { W {} }, fro_: |_x| { W {} } }
+    Equiv { to_: |_x| W {}, fro_: |_x| W {} }
 }
 
 fn add_equiv_r<let C: u32, let N: u32, let M: u32, EN, EM>(
     _: Equiv<W<N>, EN, W<M>, EM>,
 ) -> Equiv<W<C + N>, (), W<C + M>, ()> {
-    Equiv { to_: |_x| { W {} }, fro_: |_x| { W {} } }
+    Equiv { to_: |_x| W {}, fro_: |_x| W {} }
 }
 
 fn mul_comm<let N: u32, let M: u32>() -> Equiv<W<N * M>, (), W<M * N>, ()> {
-    Equiv { to_: |_x| { W {} }, fro_: |_x| { W {} } }
+    Equiv { to_: |_x| W {}, fro_: |_x| W {} }
 }
 
 fn mul_add<let N: u32, let M: u32, let P: u32>() -> Equiv<W<N * (M + P)>, (), W<N * M + N * P>, ()> {
-    Equiv { to_: |_x| { W {} }, fro_: |_x| { W {} } }
+    Equiv { to_: |_x| W {}, fro_: |_x| W {} }
 }
 
 // (N + 1) * N == N * N + N

test_programs/compile_success_empty/comptime_closures/src/main.nr

@@ -6,7 +6,7 @@ fn main() {
 
 fn closure_test(mut x: Field) {
     let one = 1;
-    let add1 = |z| { (|| { *z += one; })() };
+    let add1 = |z| (|| { *z += one; })();
 
     let two = 2;
     let add2 = |z| { *z = *z + two; };

test_programs/compile_success_empty/inner_outer_cl/src/main.nr

@@ -2,7 +2,7 @@ fn main() {
     let z1 = 0;
     let z2 = 1;
     let cl_outer = |x| {
-        let cl_inner = |y| { x + y + z2 };
+        let cl_inner = |y| x + y + z2;
         cl_inner(1) + z1
     };
     let result = cl_outer(1);

test_programs/compile_success_empty/use_callers_scope/src/main.nr

@@ -19,7 +19,7 @@ mod bar {
 
         // Ensure closures can still access Bar even
         // though `map` separates them from `fn_attr`.
-        let _ = &[1, 2, 3].map(|_| { quote { Bar }.as_type() });
+        let _ = &[1, 2, 3].map(|_| quote { Bar }.as_type());
     }
 
     // use_callers_scope should also work nested

test_programs/execution_success/higher_order_functions/src/main.nr

@@ -5,7 +5,7 @@ fn main(w: Field) -> pub Field {
     assert(twice(|x| x * 2, 5) == 20);
     assert((|x, y| x + y + 1)(2, 3) == 6);
     // nested lambdas
-    assert((|a, b| { a + (|c| c + 2)(b) })(0, 1) == 3);
+    assert((|a, b| a + (|c| c + 2)(b))(0, 1) == 3);
     // Closures:
     let a = 42;
     let g = || a;

tooling/nargo_fmt/src/chunks.rs

@@ -466,7 +466,11 @@ pub(crate) enum GroupKind {
     /// fit in the current line and it's not a block, instead of splitting that expression
     /// somewhere that's probably undesired, we'll "turn it" into a block
     /// (write the "{" and "}" delimiters) and write the lambda body in the next line.
-    LambdaBody { is_block: bool },
+    LambdaBody {
+        block_statement_count: Option<usize>,
+        has_comments: bool,
+        lambda_has_return_type: bool,
+    },
     /// A method call.
     /// We track all this information to see, if we end up needing to format this call
     /// in multiple lines, if we can write everything up to the left parentheses (inclusive)
@@ -762,14 +766,14 @@ impl<'a> Formatter<'a> {
 
             // If a lambda body doesn't fit in the current line and it's not a block,
             // we can turn it into a block and write it in the next line, so its contents fit.
-            if let GroupKind::LambdaBody { is_block: false } = group.kind {
+            if let GroupKind::LambdaBody { block_statement_count: None, .. } = group.kind {
                 // Try to format it again in the next line, but we don't want to recurse
                 // infinitely so we change the group kind.
                 group.kind = GroupKind::Regular;
                 self.write("{");
                 self.trim_spaces();
                 self.increase_indentation();
-                self.write_line();
+                self.write_line_without_skipping_whitespace_and_comments();
                 self.write_indentation();
                 self.format_chunk_group_impl(group);
 
@@ -803,7 +807,7 @@ impl<'a> Formatter<'a> {
                 // )
                 let comma_trimmed = self.trim_comma();
                 self.decrease_indentation();
-                self.write_line();
+                self.write_line_without_skipping_whitespace_and_comments();
                 self.write_indentation();
                 self.write("}");
                 if comma_trimmed {
@@ -816,6 +820,24 @@ impl<'a> Formatter<'a> {
             return;
         }
 
+        // At this point we determined we are going to write this group in a single line.
+        // If the current group is a lambda body that is a block with a single statement, like this:
+        //
+        // |x| { 1 + 2 }
+        //
+        // given that we determined the block fits the current line, if we remove the surrounding
+        // `{ .. }` it will still fit the current line, and reduce some noise from the code
+        // (this is what rustfmt seems to do too).
+        if let GroupKind::LambdaBody {
+            block_statement_count: Some(1),
+            has_comments: false,
+            lambda_has_return_type: false,
+        } = group.kind
+        {
+            self.format_lambda_body_removing_braces(group);
+            return;
+        }
+
         self.format_chunk_group_in_one_line(group);
     }
 
@@ -939,6 +961,31 @@ impl<'a> Formatter<'a> {
         }
     }
 
+    fn format_lambda_body_removing_braces(&mut self, group: ChunkGroup) {
+        // Write to an intermediate string so we can remove the braces if needed.
+        let text_chunk = self.chunk_formatter().chunk(|formatter| {
+            formatter.format_chunk_group_in_one_line(group);
+        });
+        let string = text_chunk.string;
+
+        // Don't remove the braces if the lambda's body is a Semi expression.
+        if string.ends_with("; }") || string.ends_with("; },") {
+            self.write(&string);
+            return;
+        }
+
+        let string = string.strip_prefix("{ ").unwrap();
+
+        // The lambda might have a trailing comma if it's inside an arguments list
+        if let Some(string) = string.strip_suffix(" },") {
+            self.write(string);
+            self.write(",");
+        } else {
+            let string = string.strip_suffix(" }").unwrap();
+            self.write(string);
+        }
+    }
+
     /// Appends the string to the current buffer line by line, with some pre-checks.
     fn write_chunk_lines(&mut self, string: &str) {
         for (index, line) in string.lines().enumerate() {

tooling/nargo_fmt/src/formatter/expression.rs

@@ -198,6 +198,8 @@
     fn format_lambda(&mut self, lambda: Lambda) -> FormattedLambda {
         let mut group = ChunkGroup::new();
 
+        let lambda_has_return_type = lambda.return_type.typ != UnresolvedTypeData::Unspecified;
+
         let params_and_return_type_chunk = self.chunk(|formatter| {
             formatter.write_token(Token::Pipe);
             for (index, (pattern, typ)) in lambda.parameters.into_iter().enumerate() {
@@ -218,7 +220,7 @@
             }
             formatter.write_token(Token::Pipe);
             formatter.write_space();
-            if lambda.return_type.typ != UnresolvedTypeData::Unspecified {
+            if lambda_has_return_type {
                 formatter.write_token(Token::Arrow);
                 formatter.write_space();
                 formatter.format_type(lambda.return_type);
@@ -230,16 +232,27 @@
 
         group.text(params_and_return_type_chunk);
 
-        let body_is_block = matches!(lambda.body.kind, ExpressionKind::Block(..));
+        let block_statement_count = if let ExpressionKind::Block(block) = &lambda.body.kind {
+            Some(block.statements.len())
+        } else {
+            None
+        };
 
         let mut body_group = ChunkGroup::new();
-        body_group.kind = GroupKind::LambdaBody { is_block: body_is_block };
 
+        let comments_count_before_body = self.written_comments_count;
         self.format_expression(lambda.body, &mut body_group);
+
+        body_group.kind = GroupKind::LambdaBody {
+            block_statement_count,
+            has_comments: self.written_comments_count > comments_count_before_body,
+            lambda_has_return_type,
+        };
+
         group.group(body_group);
 
         let first_line_width = params_and_return_type_chunk_width
-            + (if body_is_block {
+            + (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
@@ -1671,9 +1684,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,
     )));
@@ -1709,7 +1722,7 @@
     fn format_nested_method_call_with_maximum_width_2() {
         let src = "fn foo() {
     assert(
         p4_affine.eq(Gaffine::new(
             6890855772600357754907169075114257697580319025794532037257385534741338397365,
             4338620300185947561074059802482547481416142213883829469920100239455078257889,
         )),
@@ -1717,7 +1730,7 @@
 }
 ";
         let expected = "fn foo() {
     assert(p4_affine.eq(Gaffine::new(
         6890855772600357754907169075114257697580319025794532037257385534741338397365,
         4338620300185947561074059802482547481416142213883829469920100239455078257889,
     )));
@@ -1980,12 +1993,44 @@
     }
 
     #[test]
-    fn format_lambda_with_block() {
+    fn format_lambda_with_block_simplifies() {
         let src = "global x = | |  {  1  } ;";
-        let expected = "global x = || { 1 };\n";
+        let expected = "global x = || 1;\n";
+        assert_format(src, expected);
+    }
+
+    #[test]
+    fn format_lambda_with_block_does_not_simplify_if_it_ends_with_semicolon() {
+        let src = "global x = | |  {  1;  } ;";
+        let expected = "global x = || { 1; };\n";
+        assert_format(src, expected);
+    }
+
+    #[test]
+    fn format_lambda_with_block_does_not_simplify_if_it_has_return_type() {
+        let src = "global x = | | -> i32  {  1  } ;";
+        let expected = "global x = || -> i32 { 1 };\n";
+        assert_format(src, expected);
+    }
+
+    #[test]
+    fn format_lambda_with_simplifies_block_with_quote() {
+        let src = "global x = | | {  quote { 1 }   } ;";
+        let expected = "global x = || quote { 1 };\n";
         assert_format(src, expected);
     }
 
+    #[test]
+    fn format_lambda_with_block_simplifies_inside_arguments_list() {
+        let src = "global x = some_call(this_is_a_long_argument, | |  {  1  });";
+        let expected = "global x = some_call(
+    this_is_a_long_argument,
+    || 1,
+);
+";
+        assert_format_with_max_width(src, expected, 20);
+    }
+
     #[test]
     fn format_lambda_with_block_multiple_statements() {
         let src = "global x = | a, b |  {  1; 2  } ;";

tooling/nargo_fmt/tests/expected/contract.nr

@@ -34,7 +34,7 @@ contract Benchmarking {
                 notes: Map::new(
                     context,
                     1,
-                    |context, slot| { PrivateSet::new(context, slot, ValueNoteMethods) },
+                    |context, slot| PrivateSet::new(context, slot, ValueNoteMethods),
                 ),
                 balances: Map::new(
                     context,