Skip to content

Commit

Permalink
feat: let the formatter remove lambda block braces for single-stateme…
Browse files Browse the repository at this point in the history
…nt blocks (#6335)

…nt blocks

# Description

## Problem

Not a problem, but rustfmt seems to do the same and it can help reduce
some of the code noise.

## Summary

## Additional Context

Maybe this feature is controversial? I didn't find a way to turn it off
in rustfmt.

## Documentation

Check one:
- [x] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[For Experimental Features]** Documentation to be submitted in a
separate PR.

# PR Checklist

- [x] I have tested the changes locally.
- [x] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.
  • Loading branch information
asterite authored Oct 24, 2024
1 parent fb5d16b commit 52f7c0b
Show file tree
Hide file tree
Showing 8 changed files with 112 additions and 20 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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; };
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
55 changes: 51 additions & 4 deletions tooling/nargo_fmt/src/chunks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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 {
Expand All @@ -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);
}

Expand Down Expand Up @@ -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() {
Expand Down
57 changes: 51 additions & 6 deletions tooling/nargo_fmt/src/formatter/expression.rs
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,8 @@ impl<'a, 'b> ChunkFormatter<'a, 'b> {
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() {
Expand All @@ -218,7 +220,7 @@ impl<'a, 'b> ChunkFormatter<'a, 'b> {
}
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);
Expand All @@ -230,16 +232,27 @@ impl<'a, 'b> ChunkFormatter<'a, 'b> {

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
Expand Down Expand Up @@ -1980,12 +1993,44 @@ global y = 1;
}

#[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 } ;";
Expand Down
2 changes: 1 addition & 1 deletion tooling/nargo_fmt/tests/expected/contract.nr
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit 52f7c0b

Please sign in to comment.