chore(fmt): fix emit comments (#3406)

Co-authored-by: kevaundray <[email protected]>

tooling/nargo_fmt/src/utils.rs

@@ -189,6 +189,10 @@ fn comment_len(comment: &str) -> usize {
     }
 }
 
+pub(crate) fn count_newlines(slice: &str) -> usize {
+    bytecount::count(slice.as_bytes(), b'\n')
+}
+
 pub(crate) trait Item {
     fn span(&self) -> Span;
 

tooling/nargo_fmt/src/visitor.rs

@@ -72,13 +72,6 @@ impl<'me> FmtVisitor<'me> {
         self.buffer
     }
 
-    fn with_indent<T>(&mut self, f: impl FnOnce(&mut Self) -> T) -> T {
-        self.indent.block_indent(self.config);
-        let ret = f(self);
-        self.indent.block_unindent(self.config);
-        ret
-    }
-
     fn at_start(&self) -> bool {
         self.buffer.is_empty()
     }
@@ -102,16 +95,12 @@ impl<'me> FmtVisitor<'me> {
         self.push_str(&rewrite);
     }
 
-    fn format_missing(&mut self, end: u32) {
-        self.format_missing_inner(end, |this, slice, _| this.push_str(slice));
-    }
-
     #[track_caller]
     fn format_missing_indent(&mut self, end: u32, should_indent: bool) {
         self.format_missing_inner(end, |this, last_slice, slice| {
             this.push_str(last_slice.trim_end());
 
-            if last_slice == slice && !this.at_start() {
+            if (last_slice == slice && !this.at_start()) || this.buffer.ends_with("*/") {
                 this.push_str("\n");
             }
 
@@ -144,49 +133,85 @@ impl<'me> FmtVisitor<'me> {
             self.push_vertical_spaces(slice);
             process_last_slice(self, "", slice);
         } else {
-            if !self.at_start() {
-                if self.buffer.ends_with('{') {
-                    self.push_str("\n");
-                } else {
-                    self.push_vertical_spaces(slice);
-                }
-            }
-
-            let (result, last_end) = self.format_comment_in_block(slice);
-
-            if result.is_empty() {
+            let (result, last_end) = self.format_comment_in_block(slice, start);
+            if result.trim().is_empty() {
                 process_last_slice(self, slice, slice);
             } else {
-                self.push_str(result.trim_end());
-                let subslice = &slice[last_end as usize..];
-                process_last_slice(self, subslice, subslice);
+                let last_snippet = &slice[last_end as usize..];
+                self.push_str(&result);
+                process_last_slice(self, last_snippet, &result);
             }
         }
     }
 
-    fn format_comment_in_block(&mut self, slice: &str) -> (String, u32) {
+    fn format_comment_in_block(&mut self, slice: &str, start: u32) -> (String, u32) {
         let mut result = String::new();
         let mut last_end = 0;
 
-        for spanned in Lexer::new(slice).skip_comments(false).flatten() {
+        let mut comments = Lexer::new(slice).skip_comments(false).flatten();
+
+        if let Some(comment) = comments.next() {
+            let span = comment.to_span();
+            let diff = start;
+            let big_snippet = &self.source[..(span.start() + diff) as usize];
+            let last_char = big_snippet.chars().rev().find(|rev_c| ![' ', '\t'].contains(rev_c));
+            let fix_indent = last_char.map_or(true, |rev_c| ['{', '\n'].contains(&rev_c));
+
+            if let Token::LineComment(_, _) | Token::BlockComment(_, _) = comment.into_token() {
+                let starts_with_newline = slice.starts_with('\n');
+                let comment = &slice[span.start() as usize..span.end() as usize];
+
+                if fix_indent {
+                    if let Some('{') = last_char {
+                        result.push('\n');
+                    }
+
+                    if let Some('\n') = last_char {
+                        result.push('\n');
+                    }
+
+                    let indent_str = self.indent.to_string();
+                    result.push_str(&indent_str);
+                } else {
+                    match (starts_with_newline, self.at_start()) {
+                        (false, false) => {
+                            result.push(' ');
+                        }
+                        (true, _) => {
+                            result.push_str(&self.indent.to_string_with_newline());
+                        }
+                        _ => {}
+                    };
+                }
+
+                result.push_str(comment);
+            }
+        }
+
+        for spanned in comments {
             let span = spanned.to_span();
             last_end = span.end();
 
             if let Token::LineComment(_, _) | Token::BlockComment(_, _) = spanned.token() {
-                result.push_str(&self.indent.to_string());
-                result.push_str(&slice[span.start() as usize..span.end() as usize]);
-                result.push('\n');
+                let comment = &slice[span.start() as usize..span.end() as usize];
+
+                result.push_str(&self.indent.to_string_with_newline());
+                result.push_str(comment);
             }
         }
 
+        if slice.trim_end_matches([' ', '\t']).ends_with(['\n', '\r']) {
+            result.push('\n');
+        }
+
         (result, last_end)
     }
 
     fn push_vertical_spaces(&mut self, slice: &str) {
         let newline_upper_bound = 2;
         let newline_lower_bound = 1;
 
-        let mut newline_count = bytecount::count(slice.as_bytes(), b'\n');
+        let mut newline_count = utils::count_newlines(slice);
         let offset = self.buffer.chars().rev().take_while(|c| *c == '\n').count();
 
         if newline_count + offset > newline_upper_bound {

tooling/nargo_fmt/src/visitor/expr.rs

@@ -1,5 +1,5 @@
 use noirc_frontend::{
-    hir::resolution::errors::Span, token::Token, ArrayLiteral, BlockExpression,
+    hir::resolution::errors::Span, lexer::Lexer, token::Token, ArrayLiteral, BlockExpression,
     ConstructorExpression, Expression, ExpressionKind, IfExpression, Literal, Statement,
     StatementKind, UnaryOp,
 };
@@ -290,17 +290,13 @@ impl FmtVisitor<'_> {
 
         self.trim_spaces_after_opening_brace(&block.0);
 
-        self.with_indent(|this| {
-            this.visit_stmts(block.0);
-        });
+        self.indent.block_indent(self.config);
 
-        let slice = self.slice(self.last_position..block_span.end() - 1).trim_end();
-        self.push_str(slice);
+        self.visit_stmts(block.0);
 
+        let span = (self.last_position..block_span.end() - 1).into();
+        self.close_block(span);
         self.last_position = block_span.end();
-
-        self.push_str(&self.indent.to_string_with_newline());
-        self.push_str("}");
     }
 
     fn trim_spaces_after_opening_brace(&mut self, block: &[Statement]) {
@@ -315,18 +311,48 @@ impl FmtVisitor<'_> {
     pub(crate) fn visit_empty_block(&mut self, block_span: Span) {
         let slice = self.slice(block_span);
         let comment_str = slice[1..slice.len() - 1].trim();
-        let block_str = if comment_str.is_empty() {
-            "{}".to_string()
+
+        if comment_str.is_empty() {
+            self.push_str("{}");
         } else {
+            self.push_str("{");
             self.indent.block_indent(self.config);
-            let (comment_str, _) = self.format_comment_in_block(comment_str);
-            let comment_str = comment_str.trim_matches('\n');
-            self.indent.block_unindent(self.config);
-            let close_indent = self.indent.to_string();
-            format!("{{\n{comment_str}\n{close_indent}}}")
+            self.close_block(block_span);
         };
+
         self.last_position = block_span.end();
-        self.push_str(&block_str);
+    }
+
+    pub(crate) fn close_block(&mut self, span: Span) {
+        let slice = self.slice(span);
+
+        for spanned in Lexer::new(slice).skip_comments(false).flatten() {
+            match spanned.token() {
+                Token::LineComment(_, _) | Token::BlockComment(_, _) => {
+                    let token_span = spanned.to_span();
+
+                    let offset = token_span.start();
+                    let sub_slice = &slice[token_span.start() as usize..token_span.end() as usize];
+
+                    let span_in_between = span.start()..span.start() + offset;
+                    let slice_in_between = self.slice(span_in_between);
+                    let comment_on_same_line = !slice_in_between.contains('\n');
+
+                    if comment_on_same_line {
+                        self.push_str(" ");
+                        self.push_str(sub_slice);
+                    } else {
+                        self.push_str(&self.indent.to_string_with_newline());
+                        self.push_str(sub_slice);
+                    }
+                }
+                _ => {}
+            }
+        }
+
+        self.indent.block_unindent(self.config);
+        self.push_str(&self.indent.to_string_with_newline());
+        self.push_str("}");
     }
 }
 

tooling/nargo_fmt/src/visitor/item.rs

@@ -33,36 +33,25 @@ impl super::FmtVisitor<'_> {
                 ItemKind::Submodules(module) => {
                     let name = module.name;
 
-                    self.format_missing(span.start());
+                    self.format_missing_indent(span.start(), true);
 
                     let after_brace = self.span_after(span, Token::LeftBrace).start();
                     self.last_position = after_brace;
 
                     let keyword = if module.is_contract { "contract" } else { "mod" };
 
-                    let indent = if self.at_start()
-                        || self.buffer.ends_with(|ch: char| ch.is_whitespace())
-                    {
-                        self.indent.to_string()
-                    } else {
-                        self.indent.to_string_with_newline()
-                    };
-                    self.push_str(&format!("{indent}{keyword} {name} "));
+                    self.push_str(&format!("{keyword} {name} "));
 
                     if module.contents.items.is_empty() {
                         self.visit_empty_block((after_brace - 1..span.end()).into());
+                        continue;
                     } else {
                         self.push_str("{");
-                        let indent = self.with_indent(|this| {
-                            this.visit_module(module.contents);
-
-                            let mut indent = this.indent;
-                            indent.block_unindent(self.config);
-                            indent.to_string_with_newline()
-                        });
-                        self.push_str(&format!("{indent}}}"));
+                        self.indent.block_indent(self.config);
+                        self.visit_module(module.contents);
                     }
 
+                    self.close_block((self.last_position..span.end() - 1).into());
                     self.last_position = span.end();
                 }
                 ItemKind::Import(_)

tooling/nargo_fmt/tests/expected/comment.nr

@@ -1,10 +1,8 @@
 fn comment1() {
     //
 }
-
 // random comment
-fn comment2() {
-    // Test
+fn comment2() { // Test
 }
 
 fn comment3() // some comment

tooling/nargo_fmt/tests/expected/contract.nr

@@ -34,13 +34,11 @@ contract Benchmarking {
 
     #[aztec(private)]
     fn constructor() {}
-
     // Nec tincidunt praesent semper feugiat nibh sed pulvinar. Nibh nisl condimentum id venenatis a.
     #[aztec(private)]
     fn create_note(owner: Field, value: Field) {
         increment(storage.notes.at(owner), value, owner);
     }
-
     // Diam quam nulla porttitor massa id. Elit ullamcorper dignissim cras tincidunt lobortis feugiat.
     #[aztec(private)]
     fn recreate_note(owner: Field, index: u32) {
@@ -51,15 +49,13 @@ contract Benchmarking {
         owner_notes.remove(note);
         increment(owner_notes, note.value, owner);
     }
-
     // Ultrices in iaculis nunc sed augue lacus.
     #[aztec(public)]
     fn increment_balance(owner: Field, value: Field) {
         let current = storage.balances.at(owner).read();
         storage.balances.at(owner).write(current + value);
         let _callStackItem1 = context.call_public_function(context.this_address(), compute_selector("broadcast(Field)"), [owner]);
     }
-
     // Est ultricies integer quis auctor elit sed. In nibh mauris cursus mattis molestie a iaculis.
     #[aztec(public)]
     fn broadcast(owner: Field) {
@@ -71,3 +67,5 @@ contract Benchmarking {
         note_utils::compute_note_hash_and_nullifier(ValueNoteMethods, note_header, preimage)
     }
 }
+// Uses the token bridge contract, which tells which input token we need to talk to and handles the exit funds to L1 
+contract Uniswap {}

tooling/nargo_fmt/tests/expected/expr.nr

@@ -2,8 +2,7 @@
 fn qux() {
     {}
 
-    {
-        /* a block with a comment */
+    { /* a block with a comment */
     }
     {}
     {
@@ -72,6 +71,22 @@ fn test() {
     34
 }
 
+fn line() {
+    42; // 42
+}
+
+fn line() {
+    42; // 42
+    42;
+    // hello
+}
+
+fn line() {
+    42; // 42
+    // 42
+    // hello
+}
+
 fn parenthesized() {
     value + (x as Field)
 }

tooling/nargo_fmt/tests/expected/single_fn.nr

@@ -0,0 +1,3 @@
+fn main() {
+    // hello
+}

tooling/nargo_fmt/tests/expected/single_mod.nr

@@ -0,0 +1 @@
+mod hello {}

tooling/nargo_fmt/tests/expected/struct.nr

@@ -55,14 +55,12 @@ fn main(x: Field, y: Field) {
     assert(p.bar() == x);
     assert(p.second == y);
     assert(p.first.array[0] != p.first.array[1]);
-
     // Nested structs
     let (struct_from_tuple, a_bool) = test_struct_in_tuple(true, x, y);
     assert(struct_from_tuple.my_bool == true);
     assert(a_bool == true);
     assert(struct_from_tuple.my_int == 5);
     assert(struct_from_tuple.my_nest.a == 0);
-
     // Regression test for issue #670
     let Animal { legs, eyes } = get_dog();
     let six = legs + eyes as Field;

tooling/nargo_fmt/tests/expected/tuple.nr

@@ -22,7 +22,6 @@ fn main() {
         2);
 
     (/*1*/ 1, /*2*/ 2);
-
     // FIXME:
     (((//2
                 1,),),);

tooling/nargo_fmt/tests/expected/vec.nr

@@ -1,7 +1,6 @@
 struct Vec<T> { 
     slice: [T]
 }
-
 // A mutable vector type implemented as a wrapper around immutable slices.
 // A separate type is technically not needed but helps differentiate which operations are mutable.
 impl<T> Vec<T> {