chore: format submodule/contract (#3286)

tooling/nargo_fmt/src/lib.rs

@@ -30,6 +30,6 @@ pub use config::Config;
 
 pub fn format(source: &str, parsed_module: ParsedModule, config: &Config) -> String {
     let mut fmt = FmtVisitor::new(source, config);
-    fmt.visit_module(parsed_module);
+    fmt.visit_file(parsed_module);
     fmt.finish()
 }

tooling/nargo_fmt/src/utils.rs

@@ -119,15 +119,13 @@ impl<'me, T> Exprs<'me, T> {
 }
 
 pub(crate) trait FindToken {
-    fn find_token(&self, token: Token) -> Option<u32>;
+    fn find_token(&self, token: Token) -> Option<Span>;
     fn find_token_with(&self, f: impl Fn(&Token) -> bool) -> 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().start()))
+    fn find_token(&self, token: Token) -> Option<Span> {
+        Lexer::new(self).flatten().find_map(|it| (it.token() == &token).then(|| it.to_span()))
     }
 
     fn find_token_with(&self, f: impl Fn(&Token) -> bool) -> Option<u32> {
@@ -163,7 +161,9 @@ pub(crate) fn find_comment_end(slice: &str, is_last: bool) -> usize {
     }
 
     let newline_index = slice.find('\n');
-    if let Some(separator_index) = slice.find_token(Token::Comma).map(|index| index as usize) {
+    if let Some(separator_index) =
+        slice.find_token(Token::Comma).map(|index| index.start() as usize)
+    {
         match (block_open_index, newline_index) {
             (Some(block), None) if block > separator_index => separator_index + 1,
             (Some(block), None) => {

tooling/nargo_fmt/src/visitor.rs

@@ -2,7 +2,7 @@ mod expr;
 mod item;
 mod stmt;
 
-use noirc_frontend::{hir::resolution::errors::Span, token::Token};
+use noirc_frontend::{hir::resolution::errors::Span, lexer::Lexer, token::Token};
 
 use crate::{
     config::Config,
@@ -33,11 +33,20 @@ impl<'me> FmtVisitor<'me> {
         &self.source[span.start() as usize..span.end() as usize]
     }
 
+    fn span_after(&self, span: impl Into<Span>, token: Token) -> Span {
+        let span = span.into();
+
+        let slice = self.slice(span);
+        let offset = slice.find_token(token).unwrap().end();
+
+        (span.start() + offset..span.end()).into()
+    }
+
     fn span_before(&self, span: impl Into<Span>, token: Token) -> Span {
         let span = span.into();
 
         let slice = self.slice(span);
-        let offset = slice.find_token(token).unwrap();
+        let offset = slice.find_token(token).unwrap().start();
 
         (span.start() + offset..span.end()).into()
     }
@@ -127,8 +136,42 @@ impl<'me> FmtVisitor<'me> {
             self.push_vertical_spaces(slice);
             process_last_slice(self, "", slice);
         } else {
-            process_last_slice(self, slice, slice);
+            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() {
+                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);
+            }
+        }
+    }
+
+    fn format_comment_in_block(&mut self, slice: &str) -> (String, u32) {
+        let mut result = String::new();
+        let mut last_end = 0;
+
+        for spanned in Lexer::new(slice).skip_comments(false).flatten() {
+            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');
+            }
         }
+
+        (result, last_end)
     }
 
     fn push_vertical_spaces(&mut self, slice: &str) {

tooling/nargo_fmt/src/visitor/expr.rs

@@ -32,7 +32,7 @@ impl FmtVisitor<'_> {
         match kind {
             ExpressionKind::Block(block) => {
                 let mut visitor = self.fork();
-                visitor.visit_block(block, span, true);
+                visitor.visit_block(block, span);
                 visitor.buffer
             }
             ExpressionKind::Prefix(prefix) => {
@@ -192,7 +192,7 @@ impl FmtVisitor<'_> {
 
                 self.format_if(*if_expr)
             }
-            ExpressionKind::Variable(_) | ExpressionKind::Lambda(_) => self.slice(span).to_string(),
+            ExpressionKind::Lambda(_) | ExpressionKind::Variable(_) => self.slice(span).to_string(),
             ExpressionKind::Error => unreachable!(),
         }
     }
@@ -281,14 +281,9 @@ impl FmtVisitor<'_> {
         format!("{type_name} {{{fields}}}")
     }
 
-    pub(crate) fn visit_block(
-        &mut self,
-        block: BlockExpression,
-        block_span: Span,
-        should_indent: bool,
-    ) {
+    pub(crate) fn visit_block(&mut self, block: BlockExpression, block_span: Span) {
         if block.is_empty() {
-            self.visit_empty_block(block_span, should_indent);
+            self.visit_empty_block(block_span);
             return;
         }
 
@@ -306,10 +301,7 @@ impl FmtVisitor<'_> {
 
         self.last_position = block_span.end();
 
-        self.push_str("\n");
-        if should_indent {
-            self.push_str(&self.indent.to_string());
-        }
+        self.push_str(&self.indent.to_string_with_newline());
         self.push_str("}");
     }
 
@@ -322,19 +314,18 @@ impl FmtVisitor<'_> {
         }
     }
 
-    fn visit_empty_block(&mut self, block_span: Span, should_indent: bool) {
+    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()
         } else {
             self.indent.block_indent(self.config);
-            let open_indent = self.indent.to_string();
+            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 = if should_indent { self.indent.to_string() } else { String::new() };
-
-            let ret = format!("{{\n{open_indent}{comment_str}\n{close_indent}}}");
-            ret
+            let close_indent = self.indent.to_string();
+            format!("{{\n{comment_str}\n{close_indent}}}")
         };
         self.last_position = block_span.end();
         self.push_str(&block_str);

tooling/nargo_fmt/src/visitor/item.rs

@@ -1,5 +1,6 @@
 use noirc_frontend::{
     parser::{Item, ItemKind},
+    token::Token,
     NoirFunction, ParsedModule,
 };
 
@@ -10,24 +11,72 @@ impl super::FmtVisitor<'_> {
         (slice.trim_end().to_string(), force_brace_newline)
     }
 
-    pub(crate) fn visit_module(&mut self, module: ParsedModule) {
+    pub(crate) fn visit_file(&mut self, module: ParsedModule) {
+        self.visit_module(module);
+        self.format_missing_indent(self.source.len() as u32, false);
+    }
+
+    fn visit_module(&mut self, module: ParsedModule) {
         for Item { kind, span } in module.items {
             match kind {
                 ItemKind::Function(func) => {
                     let (fn_before_block, force_brace_newline) =
                         self.format_fn_before_block(func.clone(), span.start());
 
-                    self.format_missing_indent(span.start(), false);
+                    self.format_missing_indent(span.start(), true);
 
                     self.push_str(&fn_before_block);
                     self.push_str(if force_brace_newline { "\n" } else { " " });
 
-                    self.visit_block(func.def.body, func.def.span, false);
+                    self.visit_block(func.def.body, func.def.span);
+                }
+                ItemKind::Submodules(module) => {
+                    let name = module.name;
+
+                    self.format_missing(span.start());
+
+                    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} "));
+
+                    if module.contents.items.is_empty() {
+                        self.visit_empty_block((after_brace - 1..span.end()).into());
+                    } 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.last_position = span.end();
+                }
+                ItemKind::Import(_)
+                | ItemKind::Struct(_)
+                | ItemKind::Trait(_)
+                | ItemKind::TraitImpl(_)
+                | ItemKind::Impl(_)
+                | ItemKind::TypeAlias(_)
+                | ItemKind::Global(_)
+                | ItemKind::ModuleDecl(_) => {
+                    self.push_rewrite(self.slice(span).to_string(), span);
+                    self.last_position = span.end();
                 }
-                _ => self.format_missing(span.end()),
             }
         }
-
-        self.format_missing_indent(self.source.len() as u32, false);
     }
 }

tooling/nargo_fmt/src/visitor/stmt.rs

@@ -55,7 +55,9 @@ impl super::FmtVisitor<'_> {
 
                     self.push_rewrite(constrain, span);
                 }
-                StatementKind::Assign(_) | StatementKind::For(_) => self.format_missing(span.end()),
+                StatementKind::For(_) | StatementKind::Assign(_) => {
+                    self.push_rewrite(self.slice(span).to_string(), span);
+                }
                 StatementKind::Error => unreachable!(),
             }
 

tooling/nargo_fmt/tests/expected/contract.nr

@@ -0,0 +1,73 @@
+// Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. 
+// Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. 
+// Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. 
+// Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
+contract Benchmarking {
+    use dep::value_note::{
+        utils::{increment, decrement},
+        value_note::{VALUE_NOTE_LEN, ValueNote, ValueNoteMethods},
+    };
+
+    use dep::aztec::{
+        context::{Context},
+        note::{utils as note_utils, note_getter_options::NoteGetterOptions, note_header::NoteHeader},
+        selector::compute_selector,
+        log::emit_unencrypted_log,
+        state_vars::{map::Map, public_state::PublicState, set::Set},
+        types::type_serialization::field_serialization::{FieldSerializationMethods, FIELD_SERIALIZED_LEN},
+        types::address::{AztecAddress},
+    };
+
+    struct Storage {
+        notes: Map<Set<ValueNote, VALUE_NOTE_LEN>>,
+        balances: Map<PublicState<Field, FIELD_SERIALIZED_LEN>>,
+    }
+
+    impl Storage {
+        fn init(context: Context) -> pub Self {
+            Storage {
+                notes: Map::new(context, 1, |context, slot| { Set::new(context, slot, ValueNoteMethods) }),
+                balances: Map::new(context, 2, |context, slot| { PublicState::new(context, slot, FieldSerializationMethods) }),
+            }
+        }
+    }
+
+    #[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) {
+        let owner_notes = storage.notes.at(owner);
+        let getter_options = NoteGetterOptions::new().set_limit(1).set_offset(index);
+        let notes = owner_notes.get_notes(getter_options);
+        let note = notes[0].unwrap_unchecked();
+        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) {
+        emit_unencrypted_log(&mut context, storage.balances.at(owner).read());
+    }
+
+    unconstrained fn compute_note_hash_and_nullifier(contract_address: Field, nonce: Field, storage_slot: Field, preimage: [Field; VALUE_NOTE_LEN]) -> [Field; 4] {
+        let note_header = NoteHeader::new(contract_address, nonce, storage_slot);
+        note_utils::compute_note_hash_and_nullifier(ValueNoteMethods, note_header, preimage)
+    }
+}

tooling/nargo_fmt/tests/expected/expr.nr

@@ -50,7 +50,7 @@ fn only_comment() {
 
 fn only_comments() {
     // Keep this here
-// Keep this here
+    // Keep this here
 }
 
 fn only_comments() {

tooling/nargo_fmt/tests/expected/submodule.nr

@@ -0,0 +1,30 @@
+mod a {
+    // hello
+}
+
+mod a {
+    // hello
+    mod b {
+        // hello
+    }
+}
+
+mod a {
+    mod b {
+        mod c {}
+    }
+}
+
+mod a {
+    // 1
+    // 2
+    // 3
+    /*test*/
+}
+
+mod a {
+    /**/
+    mod b {
+        mod c {}
+    }
+}