From e6a10887f3b5fcec0c85a157cb51414b4c2fee75 Mon Sep 17 00:00:00 2001
From: Ary Borenszweig <asterite@gmail.com>
Date: Mon, 21 Oct 2024 14:21:41 -0300
Subject: [PATCH] Fix an issue introduced in the new lambda formatting

---
 tooling/nargo_fmt/src/formatter/buffer.rs     | 10 ++++++
 tooling/nargo_fmt/src/formatter/chunks.rs     | 33 +++++++++++++++++++
 .../src/formatter/comments_and_whitespace.rs  |  5 +++
 tooling/nargo_fmt/src/formatter/expression.rs | 31 +++++++++++++++++
 4 files changed, 79 insertions(+)

diff --git a/tooling/nargo_fmt/src/formatter/buffer.rs b/tooling/nargo_fmt/src/formatter/buffer.rs
index ec3643c82f..e4740311bf 100644
--- a/tooling/nargo_fmt/src/formatter/buffer.rs
+++ b/tooling/nargo_fmt/src/formatter/buffer.rs
@@ -44,6 +44,16 @@ impl Buffer {
         }
     }
 
+    /// Trim commas from the end of the buffer. Returns true if a comma was trimmed.
+    pub(super) fn trim_comma(&mut self) -> bool {
+        if self.buffer.ends_with(',') {
+            self.buffer.truncate(self.buffer.len() - 1);
+            true
+        } else {
+            false
+        }
+    }
+
     pub(crate) fn contents(self) -> String {
         self.buffer
     }
diff --git a/tooling/nargo_fmt/src/formatter/chunks.rs b/tooling/nargo_fmt/src/formatter/chunks.rs
index d674aec5ad..77bc32476b 100644
--- a/tooling/nargo_fmt/src/formatter/chunks.rs
+++ b/tooling/nargo_fmt/src/formatter/chunks.rs
@@ -729,10 +729,43 @@ impl<'a> Formatter<'a> {
                 self.write_line();
                 self.write_indentation();
                 self.format_chunk_group_impl(group);
+
+                // If this lambda was in an expression list and it was formatted in multiple
+                // lines, it might be that the trailing comma happened after the lambda body:
+                //
+                // foo(
+                //     1,
+                //     |lambda| body,
+                // )
+                //
+                // Because we attach commas to the last text to avoid splitting it, the body
+                // in this case is "body,", so if we end up writing it as a block it will
+                // look like this:
+                //
+                // foo(
+                //     1,
+                //     |lambda| {
+                //         body,
+                //     }
+                // )
+                //
+                // So, if after writing the body we find a comma (there will be at most one)
+                // we remove it, but place it after the right brace, so it looks like this:
+                //
+                // foo(
+                //     1,
+                //     |lambda| {
+                //         body
+                //     },
+                // )
+                let comma_trimmed = self.trim_comma();
                 self.decrease_indentation();
                 self.write_line();
                 self.write_indentation();
                 self.write("}");
+                if comma_trimmed {
+                    self.write(",");
+                }
                 return;
             }
 
diff --git a/tooling/nargo_fmt/src/formatter/comments_and_whitespace.rs b/tooling/nargo_fmt/src/formatter/comments_and_whitespace.rs
index 817df30886..6d5fbf8719 100644
--- a/tooling/nargo_fmt/src/formatter/comments_and_whitespace.rs
+++ b/tooling/nargo_fmt/src/formatter/comments_and_whitespace.rs
@@ -237,6 +237,11 @@ impl<'a> Formatter<'a> {
     pub(super) fn trim_spaces(&mut self) {
         self.buffer.trim_spaces();
     }
+
+    /// Trim commas from the end of the buffer. Returns true if a comma was trimmed.
+    pub(super) fn trim_comma(&mut self) -> bool {
+        self.buffer.trim_comma()
+    }
 }
 
 #[cfg(test)]
diff --git a/tooling/nargo_fmt/src/formatter/expression.rs b/tooling/nargo_fmt/src/formatter/expression.rs
index 436a42677d..c04ba8eb70 100644
--- a/tooling/nargo_fmt/src/formatter/expression.rs
+++ b/tooling/nargo_fmt/src/formatter/expression.rs
@@ -2015,6 +2015,37 @@ global y = 1;
         assert_format(src, expected);
     }
 
+    #[test]
+    fn format_lambda_as_last_method_call_argument_3() {
+        let src = "mod moo {
+    fn foo() {
+        let mut sorted_write_tuples = unsafe {
+            get_sorted_tuple(
+                final_public_data_writes.storage,
+                |(_, leaf_a): (u32, PublicDataTreeLeaf), (_, leaf_b): (u32, PublicDataTreeLeaf)| full_field_less_than(
+                    1, 2,
+                ),
+            )
+        };
+    }
+}
+";
+        let expected = "mod moo {
+    fn foo() {
+        let mut sorted_write_tuples = unsafe {
+            get_sorted_tuple(
+                final_public_data_writes.storage,
+                |(_, leaf_a): (u32, PublicDataTreeLeaf), (_, leaf_b): (u32, PublicDataTreeLeaf)| {
+                    full_field_less_than(1, 2)
+                },
+            )
+        };
+    }
+}
+";
+        assert_format(src, expected);
+    }
+
     #[test]
     fn format_lambda_as_last_method_call_chain_argument() {
         let src = "global x = foo.bar(1).baz(2, |x| { 1; 2 });";