chore(fmt): restore removed comment

compiler/noirc_frontend/src/lexer/lexer.rs

@@ -19,6 +19,7 @@ pub struct Lexer<'a> {
     char_iter: Peekable<Zip<Chars<'a>, RangeFrom<u32>>>,
     position: Position,
     done: bool,
+    skip_trivia: bool,
 }
 
 pub type SpannedTokenResult = Result<SpannedToken, LexerErrorKind>;
@@ -39,15 +40,20 @@ impl<'a> Lexer<'a> {
         (Tokens(tokens), errors)
     }
 
-    fn new(source: &'a str) -> Self {
+    pub fn new(source: &'a str) -> Self {
         Lexer {
             // We zip with the character index here to ensure the first char has index 0
             char_iter: source.chars().zip(0..).peekable(),
             position: 0,
             done: false,
+            skip_trivia: true,
         }
     }
 
+    pub fn set_skip_trivia(&mut self, flag: bool) {
+        self.skip_trivia = flag;
+    }
+
     /// Iterates the cursor and returns the char at the new cursor position
     fn next_char(&mut self) -> Option<char> {
         let (c, index) = self.char_iter.next()?;
@@ -176,13 +182,16 @@ impl<'a> Lexer<'a> {
             Token::Minus => self.single_double_peek_token('>', prev_token, Token::Arrow),
             Token::Colon => self.single_double_peek_token(':', prev_token, Token::DoubleColon),
             Token::Slash => {
+                let start = self.position;
+
                 if self.peek_char_is('/') {
                     self.next_char();
-                    return self.parse_comment();
+                    return self.parse_comment(start);
                 } else if self.peek_char_is('*') {
                     self.next_char();
-                    return self.parse_block_comment();
+                    return self.parse_block_comment(start);
                 }
+
                 Ok(spanned_prev_token)
             }
             _ => Err(LexerErrorKind::NotADoubleChar {
@@ -377,15 +386,15 @@ impl<'a> Lexer<'a> {
         }
     }
 
-    fn parse_comment(&mut self) -> SpannedTokenResult {
-        let _ = self.eat_while(None, |ch| ch != '\n');
-        self.next_token()
+    fn parse_comment(&mut self, start: u32) -> SpannedTokenResult {
+        let comment = self.eat_while(None, |ch| ch != '\n');
+        Ok(Token::LineComment(comment).into_span(start, self.position))
     }
 
-    fn parse_block_comment(&mut self) -> SpannedTokenResult {
-        let start = self.position;
+    fn parse_block_comment(&mut self, start: u32) -> SpannedTokenResult {
         let mut depth = 1usize;
 
+        let mut content = String::new();
         while let Some(ch) = self.next_char() {
             match ch {
                 '/' if self.peek_char_is('*') => {
@@ -403,12 +412,12 @@ impl<'a> Lexer<'a> {
                         break;
                     }
                 }
-                _ => {}
+                ch => content.push(ch),
             }
         }
 
         if depth == 0 {
-            self.next_token()
+            Ok(Token::BlockComment(content).into_span(start, self.position))
         } else {
             let span = Span::inclusive(start, self.position);
             Err(LexerErrorKind::UnterminatedBlockComment { span })
@@ -424,10 +433,20 @@ impl<'a> Lexer<'a> {
 impl<'a> Iterator for Lexer<'a> {
     type Item = SpannedTokenResult;
     fn next(&mut self) -> Option<Self::Item> {
-        if self.done {
-            None
-        } else {
-            Some(self.next_token())
+        loop {
+            if self.done {
+                return None;
+            } else {
+                let token = self.next_token();
+
+                if self.skip_trivia
+                    && token.as_ref().map_or(false, |spanned| spanned.token().is_trivia())
+                {
+                    continue;
+                }
+
+                return Some(token);
+            }
         }
     }
 }
@@ -532,7 +551,7 @@ mod tests {
 
         let mut lexer = Lexer::new(input);
         for token in expected.into_iter() {
-            let got = lexer.next_token().unwrap();
+            let got = lexer.next().unwrap().unwrap();
             assert_eq!(got, token);
         }
     }
@@ -638,7 +657,7 @@ mod tests {
 
         let mut lexer = Lexer::new(input);
         for token in expected.into_iter() {
-            let got = lexer.next_token().unwrap();
+            let got = lexer.next().unwrap().unwrap();
             assert_eq!(got, token);
         }
     }
@@ -662,7 +681,7 @@ mod tests {
 
         let mut lexer = Lexer::new(input);
         for token in expected.into_iter() {
-            let got = lexer.next_token().unwrap();
+            let got = lexer.next().unwrap().unwrap();
             assert_eq!(got, token);
         }
     }
@@ -692,7 +711,7 @@ mod tests {
 
         let mut lexer = Lexer::new(input);
         for token in expected.into_iter() {
-            let first_lexer_output = lexer.next_token().unwrap();
+            let first_lexer_output = lexer.next().unwrap().unwrap();
             assert_eq!(first_lexer_output, token);
         }
     }
@@ -714,7 +733,7 @@ mod tests {
 
         let mut lexer = Lexer::new(input);
         for token in expected.into_iter() {
-            let first_lexer_output = lexer.next_token().unwrap();
+            let first_lexer_output = lexer.next().unwrap().unwrap();
             assert_eq!(first_lexer_output, token);
         }
     }
@@ -736,7 +755,7 @@ mod tests {
 
         let mut lexer = Lexer::new(input);
         for token in expected.into_iter() {
-            let first_lexer_output = lexer.next_token().unwrap();
+            let first_lexer_output = lexer.next().unwrap().unwrap();
             assert_eq!(first_lexer_output, token);
         }
     }
@@ -753,7 +772,7 @@ mod tests {
         let mut lexer = Lexer::new(input);
 
         for token in expected.into_iter() {
-            let got = lexer.next_token().unwrap();
+            let got = lexer.next().unwrap().unwrap();
             assert_eq!(got, token);
         }
     }
@@ -766,7 +785,7 @@ mod tests {
         let mut lexer = Lexer::new(input);
 
         for token in expected.into_iter() {
-            let got = lexer.next_token().unwrap();
+            let got = lexer.next().unwrap().unwrap();
             assert_eq!(got, token);
         }
     }
@@ -805,7 +824,7 @@ mod tests {
         let mut lexer = Lexer::new(input);
 
         for spanned_token in expected.into_iter() {
-            let got = lexer.next_token().unwrap();
+            let got = lexer.next().unwrap().unwrap();
             assert_eq!(got.to_span(), spanned_token.to_span());
             assert_eq!(got, spanned_token);
         }
@@ -876,7 +895,7 @@ mod tests {
         let mut lexer = Lexer::new(input);
 
         for token in expected.into_iter() {
-            let got = lexer.next_token().unwrap();
+            let got = lexer.next().unwrap().unwrap();
             assert_eq!(got, token);
         }
     }

compiler/noirc_frontend/src/lexer/token.rs

@@ -19,6 +19,8 @@
     Keyword(Keyword),
     IntType(IntType),
     Attribute(Attribute),
+    LineComment(String),
+    BlockComment(String),
     /// <
     Less,
     /// <=
@@ -95,6 +97,12 @@
     Invalid(char),
 }
 
+impl Token {
+    pub fn is_trivia(&self) -> bool {
+        matches!(self, Token::LineComment(_) | Token::BlockComment(_))
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct SpannedToken(Spanned<Token>);
 
@@ -149,6 +157,8 @@
             Token::FmtStr(ref b) => write!(f, "f{b}"),
             Token::Keyword(k) => write!(f, "{k}"),
             Token::Attribute(ref a) => write!(f, "{a}"),
+            Token::LineComment(ref s) => write!(f, "//{s}"),
+            Token::BlockComment(ref s) => write!(f, "/*{s}*/"),
             Token::IntType(ref i) => write!(f, "{i}"),
             Token::Less => write!(f, "<"),
             Token::LessEqual => write!(f, "<="),
@@ -332,7 +342,7 @@
         match string.trim() {
             "should_fail" => Some(TestScope::ShouldFailWith { reason: None }),
             s if s.starts_with("should_fail_with") => {
                 let parts: Vec<&str> = s.splitn(2, '=').collect();
                 if parts.len() == 2 {
                     let reason = parts[1].trim();
                     let reason = reason.trim_matches('"');
@@ -665,7 +675,7 @@
             Keyword::Field => write!(f, "Field"),
             Keyword::Fn => write!(f, "fn"),
             Keyword::For => write!(f, "for"),
             Keyword::FormatString => write!(f, "fmtstr"),
             Keyword::Global => write!(f, "global"),
             Keyword::If => write!(f, "if"),
             Keyword::Impl => write!(f, "impl"),
@@ -708,7 +718,7 @@
             "Field" => Keyword::Field,
             "fn" => Keyword::Fn,
             "for" => Keyword::For,
             "fmtstr" => Keyword::FormatString,
             "global" => Keyword::Global,
             "if" => Keyword::If,
             "impl" => Keyword::Impl,

tooling/nargo_fmt/src/visitor/expr.rs

@@ -1,6 +1,6 @@
 use noirc_frontend::{
-    hir::resolution::errors::Span, ArrayLiteral, BlockExpression, Expression, ExpressionKind,
-    Literal, Statement,
+    hir::resolution::errors::Span, lexer::Lexer, token::Token, ArrayLiteral, BlockExpression,
+    Expression, ExpressionKind, Literal, Statement,
 };
 
 use super::FmtVisitor;
@@ -10,6 +10,7 @@ impl FmtVisitor<'_> {
         let span = expr.span;
 
         let rewrite = self.format_expr(expr);
+        let rewrite = recover_comment_removed(slice!(self, span.start(), span.end()), rewrite);
         self.push_rewrite(rewrite, span);
 
         self.last_position = span.end();
@@ -159,3 +160,32 @@ impl FmtVisitor<'_> {
         self.push_str(&block_str);
     }
 }
+
+fn recover_comment_removed(original: &str, new: String) -> String {
+    if changed_comment_content(original, &new) {
+        original.to_string()
+    } else {
+        new
+    }
+}
+
+fn changed_comment_content(original: &str, new: &str) -> bool {
+    comments(original) != comments(new)
+}
+
+fn comments(source: &str) -> Vec<String> {
+    let mut lexer = Lexer::new(source);
+    lexer.set_skip_trivia(false);
+
+    lexer
+        .flatten()
+        .filter_map(|spanned| {
+            if let Token::LineComment(content) | Token::BlockComment(content) = spanned.into_token()
+            {
+                Some(content)
+            } else {
+                None
+            }
+        })
+        .collect()
+}

tooling/nargo_fmt/tests/expected/infix.nr

@@ -2,4 +2,7 @@ fn foo() {
     40 + 2;
     !40 + 2;
     40 + 2 == 42;
+
+    40/*test*/ + 2 == 42;
+    40 + 2/*test*/ == 42;
 }

tooling/nargo_fmt/tests/expected/unary_operators.nr

@@ -1,3 +1,5 @@
 fn main() {
-    -1
+    -1;
+    -/*test*/1;
+    -/*test*/1;
 }

tooling/nargo_fmt/tests/input/infix.nr

@@ -2,4 +2,7 @@ fn foo() {
     40 + 2;
     !40+2;
     40 +  2 == 42;
+
+40/*test*/ + 2 == 42;
+                40 + 2/*test*/ == 42;
 }

tooling/nargo_fmt/tests/input/unary_operators.nr

@@ -1,3 +1,5 @@
 fn main() {
-    -1
+    -1;
+-/*test*/1;
+    -/*test*/1;
 }