chore(fmt): don't lose line breaks between comments (#3505)

Co-authored-by: jfecher <jfecher11@gmail.com>

compiler/noirc_frontend/src/lexer/lexer.rs

@@ -16,6 +16,7 @@ pub struct Lexer<'a> {
     position: Position,
     done: bool,
     skip_comments: bool,
+    skip_whitespaces: bool,
 }
 
 pub type SpannedTokenResult = Result<SpannedToken, LexerErrorKind>;
@@ -37,14 +38,25 @@ impl<'a> Lexer<'a> {
     }
 
     pub fn new(source: &'a str) -> Self {
-        Lexer { chars: source.char_indices(), position: 0, done: false, skip_comments: true }
+        Lexer {
+            chars: source.char_indices(),
+            position: 0,
+            done: false,
+            skip_comments: true,
+            skip_whitespaces: true,
+        }
     }
 
     pub fn skip_comments(mut self, flag: bool) -> Self {
         self.skip_comments = flag;
         self
     }
 
+    pub fn skip_whitespaces(mut self, flag: bool) -> Self {
+        self.skip_whitespaces = flag;
+        self
+    }
+
     /// Iterates the cursor and returns the char at the new cursor position
     fn next_char(&mut self) -> Option<char> {
         let (position, ch) = self.chars.next()?;
@@ -82,9 +94,13 @@ impl<'a> Lexer<'a> {
 
     fn next_token(&mut self) -> SpannedTokenResult {
         match self.next_char() {
-            Some(x) if { x.is_whitespace() } => {
-                self.eat_whitespace();
-                self.next_token()
+            Some(x) if x.is_whitespace() => {
+                let spanned = self.eat_whitespace(x);
+                if self.skip_whitespaces {
+                    self.next_token()
+                } else {
+                    Ok(spanned)
+                }
             }
             Some('<') => self.glue(Token::Less),
             Some('>') => self.glue(Token::Greater),
@@ -454,8 +470,10 @@ impl<'a> Lexer<'a> {
     }
 
     /// Skips white space. They are not significant in the source language
-    fn eat_whitespace(&mut self) {
-        self.eat_while(None, |ch| ch.is_whitespace());
+    fn eat_whitespace(&mut self, initial_char: char) -> SpannedToken {
+        let start = self.position;
+        let whitespace = self.eat_while(initial_char.into(), |ch| ch.is_whitespace());
+        SpannedToken::new(Token::Whitespace(whitespace), Span::inclusive(start, self.position))
     }
 }
 

compiler/noirc_frontend/src/lexer/token.rs

@@ -88,6 +88,8 @@ pub enum Token {
     #[allow(clippy::upper_case_acronyms)]
     EOF,
 
+    Whitespace(String),
+
     /// An invalid character is one that is not in noir's language or grammar.
     ///
     /// We don't report invalid tokens in the source as errors until parsing to
@@ -194,6 +196,7 @@ impl fmt::Display for Token {
             Token::Bang => write!(f, "!"),
             Token::EOF => write!(f, "end of input"),
             Token::Invalid(c) => write!(f, "{c}"),
+            Token::Whitespace(ref s) => write!(f, "{s}"),
         }
     }
 }

tooling/nargo_fmt/src/rewrite/array.rs

@@ -36,14 +36,14 @@ pub(crate) fn rewrite(mut visitor: FmtVisitor, array: Vec<Expression>, array_spa
         let trailing = trailing[..offset].trim_end_matches(',').trim_matches(pattern);
         last_position = item_span.end() + offset as u32;
 
-        let (leading, _) = visitor.format_comment_in_block(leading, 0, false);
-        let (trailing, _) = visitor.format_comment_in_block(trailing, 0, false);
+        let (leading, _) = visitor.format_comment_in_block(leading);
+        let (trailing, _) = visitor.format_comment_in_block(trailing);
 
         result.push(Expr { leading, value: item, trailing, different_line: false });
     }
 
     let slice = visitor.slice(last_position..end_position);
-    let (comment, _) = visitor.format_comment_in_block(slice, 0, false);
+    let (comment, _) = visitor.format_comment_in_block(slice);
     result.push(Expr {
         leading: "".into(),
         value: "".into(),

tooling/nargo_fmt/src/visitor.rs

@@ -163,7 +163,7 @@ impl<'me> FmtVisitor<'me> {
             self.push_vertical_spaces(slice);
             process_last_slice(self, "", slice);
         } else {
-            let (result, last_end) = self.format_comment_in_block(slice, start, true);
+            let (result, last_end) = self.format_comment_in_block(slice);
             if result.trim().is_empty() {
                 process_last_slice(self, slice, slice);
             } else {
@@ -174,75 +174,36 @@ impl<'me> FmtVisitor<'me> {
         }
     }
 
-    pub(crate) fn format_comment_in_block(
-        &mut self,
-        slice: &str,
-        start: u32,
-        fix_indent: bool,
-    ) -> (String, u32) {
+    pub(crate) fn format_comment_in_block(&mut self, slice: &str) -> (String, u32) {
         let mut result = String::new();
-        let mut last_end = 0;
-
-        let mut comments = Lexer::new(slice).skip_comments(false).flatten();
+        let comments = Lexer::new(slice).skip_comments(false).skip_whitespaces(false).flatten();
 
-        if let Some(comment) = comments.next() {
+        let indent = self.indent.to_string();
+        for comment in comments {
             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 =
-                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');
+            match comment.token() {
+                Token::LineComment(_, _) | Token::BlockComment(_, _) => {
+                    let comment = &slice[span.start() as usize..span.end() as usize];
+                    if result.ends_with('\n') {
+                        result.push_str(&indent);
+                    } else if !self.at_start() {
+                        result.push(' ');
                     }
-
-                    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());
-                        }
-                        (false, _) => {
-                            result.push(' ');
-                        }
-                    };
+                    result.push_str(comment);
                 }
-
-                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() {
-                let comment = &slice[span.start() as usize..span.end() as usize];
-
-                result.push_str(&self.indent.to_string_with_newline());
-                result.push_str(comment);
+                Token::Whitespace(whitespaces) => {
+                    let mut visitor = self.fork();
+                    if whitespaces.contains('\n') {
+                        visitor.push_vertical_spaces(whitespaces.trim_matches(' '));
+                        result.push_str(&visitor.finish());
+                    }
+                }
+                _ => {}
             }
         }
 
-        if slice.trim_end_matches([' ', '\t']).ends_with(['\n', '\r']) {
-            result.push('\n');
-        }
-
-        (result, last_end)
+        (result, slice.len() as u32)
     }
 
     fn push_vertical_spaces(&mut self, slice: &str) {

tooling/nargo_fmt/tests/expected/comment.nr

@@ -1,7 +1,9 @@
 fn comment1() {
     //
 }
+
 // random comment
+
 fn comment2() { // Test
 }
 

tooling/nargo_fmt/tests/expected/contract.nr

@@ -34,11 +34,13 @@ 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) {
@@ -49,6 +51,7 @@ 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) {
@@ -58,6 +61,7 @@ contract Benchmarking {
             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) {
@@ -69,5 +73,6 @@ 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/global.nr

@@ -0,0 +1,9 @@
+//! Super module :]
+
+// super global variable
+global answer = 42;
+
+// Super module :]
+
+// super global variable
+global answer = 42;

tooling/nargo_fmt/tests/expected/struct.nr

@@ -55,12 +55,14 @@ 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,6 +22,7 @@ fn main() {
         2);
 
     (/*1*/ 1, /*2*/ 2);
+
     // FIXME:
     (((//2
                 1,),),);

tooling/nargo_fmt/tests/expected/vec.nr

@@ -1,6 +1,7 @@
 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> {

tooling/nargo_fmt/tests/input/global.nr

@@ -0,0 +1,17 @@
+//! Super module :]
+
+// super global variable
+global answer = 42;
+
+// Super module :]
+
+
+
+
+
+
+
+
+
+// super global variable
+global answer = 42;