From 2d1cd591510690b19f8640335cfcb914581d2711 Mon Sep 17 00:00:00 2001
From: f01dab1e <andriy.n@obox.systems>
Date: Thu, 12 Oct 2023 22:38:18 +0000
Subject: [PATCH] chore: add tuple formatter

---
 tooling/nargo_fmt/src/visitor.rs          |  2 +-
 tooling/nargo_fmt/src/visitor/expr.rs     | 87 +++++++++++++++++++----
 tooling/nargo_fmt/tests/expected/tuple.nr | 14 ++++
 tooling/nargo_fmt/tests/input/tuple.nr    | 22 ++++++
 4 files changed, 110 insertions(+), 15 deletions(-)
 create mode 100644 tooling/nargo_fmt/tests/expected/tuple.nr
 create mode 100644 tooling/nargo_fmt/tests/input/tuple.nr

diff --git a/tooling/nargo_fmt/src/visitor.rs b/tooling/nargo_fmt/src/visitor.rs
index cf2b4d9adfd..bfc06b2f61e 100644
--- a/tooling/nargo_fmt/src/visitor.rs
+++ b/tooling/nargo_fmt/src/visitor.rs
@@ -1,7 +1,7 @@
 /// A macro to create a slice from a given data source, helping to avoid borrow checker errors.
 #[macro_export]
 macro_rules! slice {
-    ($this:ident, $start:expr, $end:expr) => {
+    ($this:expr, $start:expr, $end:expr) => {
         &$this.source[$start as usize..$end as usize]
     };
 }
diff --git a/tooling/nargo_fmt/src/visitor/expr.rs b/tooling/nargo_fmt/src/visitor/expr.rs
index d8bea6c7530..5e07ada3fb6 100644
--- a/tooling/nargo_fmt/src/visitor/expr.rs
+++ b/tooling/nargo_fmt/src/visitor/expr.rs
@@ -79,6 +79,17 @@ impl FmtVisitor<'_> {
                 let formatted_index = self.format_expr(index_expr.index);
                 format!("{}[{}]", formatted_collection, formatted_index)
             }
+            ExpressionKind::Tuple(elements) => {
+                let mut elements = TupleElements::new(self, span, elements).collect::<Vec<_>>();
+
+                let elements = if elements.len() == 1 {
+                    format!("{},", elements.pop().unwrap())
+                } else {
+                    elements.join(", ")
+                };
+
+                format!("({elements})")
+            }
             ExpressionKind::Literal(literal) => match literal {
                 Literal::Integer(_) => slice!(self, span.start(), span.end()).to_string(),
                 Literal::Array(ArrayLiteral::Repeated { repeated_element, length }) => {
@@ -171,20 +182,68 @@ fn recover_comment_removed(original: &str, new: String) -> String {
 }
 
 fn changed_comment_content(original: &str, new: &str) -> bool {
-    comments(original) != comments(new)
+    comments(original).ne(comments(new))
 }
 
-fn comments(source: &str) -> Vec<String> {
-    Lexer::new(source)
-        .skip_comments(false)
-        .flatten()
-        .filter_map(|spanned| {
-            if let Token::LineComment(content) | Token::BlockComment(content) = spanned.into_token()
-            {
-                Some(content)
-            } else {
-                None
-            }
-        })
-        .collect()
+fn comments(source: &str) -> impl Iterator<Item = String> + '_ {
+    Lexer::new(source).skip_comments(false).flatten().filter_map(|spanned| {
+        if let Token::LineComment(content) | Token::BlockComment(content) = spanned.into_token() {
+            Some(content)
+        } else {
+            None
+        }
+    })
+}
+
+struct TupleElements<'me> {
+    visitor: &'me FmtVisitor<'me>,
+    elements: std::iter::Peekable<std::vec::IntoIter<Expression>>,
+    last_position: u32,
+    end_position: u32,
+}
+
+impl<'me> TupleElements<'me> {
+    fn new(visitor: &'me FmtVisitor, span: Span, elements: Vec<Expression>) -> Self {
+        Self {
+            visitor,
+            last_position: span.start() + 1, /*(*/
+            end_position: span.end() - 1,    /*)*/
+            elements: elements.into_iter().peekable(),
+        }
+    }
+}
+
+impl Iterator for TupleElements<'_> {
+    type Item = String;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let element = self.elements.next()?;
+        let element_span = element.span;
+
+        let start = self.last_position;
+        let end = element_span.start();
+
+        let next_start = self.elements.peek().map_or(self.end_position, |expr| expr.span.start());
+
+        let leading = slice!(self.visitor, start, end).trim();
+        let trailing = slice!(self.visitor, element_span.end(), next_start);
+        let element_str = self.visitor.format_expr(element);
+
+        let end = trailing.find_token(Token::Comma).unwrap_or(trailing.len() as u32);
+
+        let trailing = trailing[..end as usize].trim_matches(',').trim();
+        self.last_position = element_span.end() + end;
+
+        format!("{leading}{element_str}{trailing}").into()
+    }
+}
+
+trait FindToken {
+    fn find_token(&self, token: Token) -> Option<u32>;
+}
+
+impl FindToken for str {
+    fn find_token(&self, token: Token) -> Option<u32> {
+        Lexer::new(self).flatten().find_map(|it| (it.token() == &token).then(|| it.to_span().end()))
+    }
 }
diff --git a/tooling/nargo_fmt/tests/expected/tuple.nr b/tooling/nargo_fmt/tests/expected/tuple.nr
new file mode 100644
index 00000000000..db74c23b4db
--- /dev/null
+++ b/tooling/nargo_fmt/tests/expected/tuple.nr
@@ -0,0 +1,14 @@
+fn main() {
+    (1,);
+    (/*test*/1,);
+    (/*a*/1/*b*/,);
+    (/*a*/1/*b*/, /*c*/2/*d*/, /*c*/2/*d*/);
+    (/*a*/1/*b*/, /*c*/2/*d*/, /*c*/2/*d*/, /*e*/3/*f*/);
+
+    (1/*1*/, 2/* 2*/);
+
+    (/*a*/1/*b*/, /*c*/2/*d*/, /*c*/2/*d*/, /*e*/3/*f*/);
+
+    // FIXME:
+    ( 1, /*test*/ );
+}
diff --git a/tooling/nargo_fmt/tests/input/tuple.nr b/tooling/nargo_fmt/tests/input/tuple.nr
new file mode 100644
index 00000000000..e0e2cdf963c
--- /dev/null
+++ b/tooling/nargo_fmt/tests/input/tuple.nr
@@ -0,0 +1,22 @@
+fn main() {
+    ( 1, );
+    ( /*test*/1, );
+    ( /*a*/1/*b*/, );
+    ( /*a*/1/*b*/, /*c*/2/*d*/, /*c*/2/*d*/ );
+    ( /*a*/1/*b*/, /*c*/2/*d*/, /*c*/2/*d*/, /*e*/3/*f*/ );
+
+    (   1  /*1*/  ,  2 /* 2*/  );
+
+    ( 
+        /*a*/
+        1
+        /*b*/,
+    /*c*/
+    2/*d*/, 
+    /*c*/2/*d*/, 
+    /*e*/3/*f*/ 
+    );
+
+    // FIXME:
+    ( 1, /*test*/ );
+}