feat(fmt): format comments supertraits (FuelLabs#4129)

## Description

Followup to FuelLabs#4005

makes `abi_supertraits` test pass comment formatting without
`handle_comments`.


In fixing this, this also inadvertently led to me fixing other
formatting issues:

- There may be comments before the module declaration - the first one
was tagged as a `Trailing` comment, even though it should not be. Added
a 0 check within lexing for this.
- Massively simplified `comments.rs` - before, we used to arbitrarily
compare to the previous token, but now that we have the concept of
`Trailing` and `Newlined` comments, we do not need to do this. We also
used to write an indent after writing a trailing comment, but I think
this should be moved up-one-level to whereever the trailing comment was
written. There is also an update to a test that retains the trailing
nature of a trailing comment in an existing test.


For the rest of changes, they seem largely unrelated to what this PR
claims to solve, so I'll go change by change:
- within `item_fn/mod.rs`: This fixes the supertraits test example
(which had a weird spacing for expected comment output)
- within `module/mod.rs`: We format comments between items (this fixes
the the tests like `abi_comments` and `abi_supertrait` that have
comments written between items.


## Checklist

- [x] I have linked to any relevant issues.
- [x] I have commented my code, particularly in hard-to-understand
areas.
- [ ] I have updated the documentation where relevant (API docs, the
reference, and the Sway book).
- [x] I have added tests that prove my fix is effective or that my
feature works.
- [ ] I have added (or requested a maintainer to add) the necessary
`Breaking*` or `New Feature` labels where relevant.
- [ ] I have done my best to ensure that my PR adheres to [the Fuel Labs
Code Review
Standards](https://github.com/FuelLabs/rfcs/blob/master/text/code-standards/external-contributors.md).
- [x] I have requested a review from the relevant team or maintainers.

examples/abi_supertraits/src/main.sw

@@ -8,7 +8,7 @@ abi MyAbi : ABIsupertrait {
     fn bar();
 } {
     fn baz() {
-        Self::foo()     // supertrait method usage
+        Self::foo() // supertrait method usage
     }
 }
 
@@ -19,6 +19,6 @@ impl ABIsupertrait for Contract {
 // The implementation of MyAbi for Contract must also implement ABIsupertrait
 impl MyAbi for Contract {
     fn bar() {
-        Self::foo()     // supertrait method usage
+        Self::foo() // supertrait method usage
     }
 }

sway-parse/src/token.rs

@@ -159,7 +159,10 @@ pub fn lex_commented(
                         .count()
                         > 0;
 
-                    let comment_kind = if has_newline {
+                    // We found a comment at the start of file, which should be accounted for as a Newlined comment.
+                    let start_of_file_found = search_end == 0 && index == 0;
+
+                    let comment_kind = if has_newline || start_of_file_found {
                         CommentKind::Newlined
                     } else {
                         CommentKind::Trailing

swayfmt/src/comments.rs

@@ -86,38 +86,22 @@ pub fn write_comments(
             writeln!(formatted_code)?;
         }
 
-        while let Some(comment) = comments_iter.next() {
+        for comment in comments_iter {
             let newlines = collect_newlines_after_comment(&formatter.comments_context, comment);
 
-            if formatted_code.trim_end().ends_with(&[']', ';', ')']) {
-                match comment.comment_kind {
-                    CommentKind::Newlined => {
-                        write!(
-                            formatted_code,
-                            "{}{}{}",
-                            formatter.shape.indent.to_string(&formatter.config)?,
-                            comment.span().as_str(),
-                            newlines
-                        )?;
-                    }
-                    CommentKind::Trailing => {
-                        write_trailing_comment(formatted_code, comment)?;
-                        if comments_iter.peek().is_none() {
-                            write!(
-                                formatted_code,
-                                "{}",
-                                formatter.shape.indent.to_string(&formatter.config)?,
-                            )?;
-                        }
-                    }
+            match comment.comment_kind {
+                CommentKind::Newlined => {
+                    write!(
+                        formatted_code,
+                        "{}{}{}",
+                        formatter.shape.indent.to_string(&formatter.config)?,
+                        comment.span().as_str(),
+                        newlines
+                    )?;
+                }
+                CommentKind::Trailing => {
+                    write_trailing_comment(formatted_code, comment)?;
                 }
-            } else {
-                writeln!(
-                    formatted_code,
-                    "{}{}",
-                    formatter.shape.indent.to_string(&formatter.config)?,
-                    comment.span().as_str(),
-                )?;
             }
         }
     }

swayfmt/src/items/item_fn/mod.rs

@@ -38,6 +38,15 @@ impl Format for ItemFn {
                     Self::open_curly_brace(formatted_code, formatter)?;
                     formatter.shape.block_indent(&formatter.config);
                     body.format(formatted_code, formatter)?;
+
+                    if let Some(final_expr_opt) = body.final_expr_opt.as_ref() {
+                        write_comments(
+                            formatted_code,
+                            final_expr_opt.span().end()..self.span().end(),
+                            formatter,
+                        )?;
+                    }
+
                     Self::close_curly_brace(formatted_code, formatter)?;
                 } else {
                     Self::open_curly_brace(formatted_code, formatter)?;
@@ -46,6 +55,9 @@ impl Format for ItemFn {
                     if has_comments(comments) {
                         formatter.shape.block_indent(&formatter.config);
                         write_comments(formatted_code, range, formatter)?;
+                        if !formatted_code.ends_with('\n') {
+                            writeln!(formatted_code)?;
+                        }
                     }
                     Self::close_curly_brace(formatted_code, formatter)?;
                 }

swayfmt/src/module/mod.rs

@@ -4,7 +4,7 @@ use crate::{
     utils::map::byte_span::{self, ByteSpan, LeafSpans},
 };
 use std::fmt::Write;
-use sway_ast::{ItemKind, Module, ModuleKind};
+use sway_ast::{Item, ItemKind, Module, ModuleKind};
 use sway_types::Spanned;
 
 pub(crate) mod dependency;
@@ -30,13 +30,23 @@ impl Format for Module {
         }
 
         let iter = self.items.iter();
+        let mut prev_item: Option<&Item> = None;
         for item in iter.clone() {
+            if let Some(prev_item) = prev_item {
+                write_comments(
+                    formatted_code,
+                    prev_item.span().end()..item.span().start(),
+                    formatter,
+                )?;
+            }
             item.format(formatted_code, formatter)?;
             if let ItemKind::Dependency { .. } = item.value {
                 // Do not print a newline after a dependency
             } else {
                 writeln!(formatted_code)?;
             }
+
+            prev_item = Some(item);
         }
 
         Ok(())

swayfmt/src/utils/language/attribute.rs

@@ -24,7 +24,19 @@ impl<T: Format + Spanned> Format for Annotated<T> {
         for attr in &self.attribute_list {
             // Write trailing comments after the end of the previous attribute
             if let Some(end) = attr_end {
-                write_comments(formatted_code, end..attr.span().start(), formatter)?;
+                write_comments(formatted_code, end..attr.span().start(), formatter).and_then(
+                    |w| {
+                        if w {
+                            write!(
+                                formatted_code,
+                                "{}",
+                                &formatter.shape.indent.to_string(&formatter.config)?,
+                            )?;
+                        };
+
+                        Ok(())
+                    },
+                )?;
             };
             attr.format(formatted_code, formatter)?;
             write!(
@@ -37,7 +49,19 @@ impl<T: Format + Spanned> Format for Annotated<T> {
 
         // Write trailing comments after the end of the last attribute
         if let Some(end) = attr_end {
-            write_comments(formatted_code, end..self.value.span().start(), formatter)?;
+            write_comments(formatted_code, end..self.value.span().start(), formatter).and_then(
+                |w| {
+                    if w {
+                        write!(
+                            formatted_code,
+                            "{}",
+                            &formatter.shape.indent.to_string(&formatter.config)?,
+                        )?;
+                    };
+
+                    Ok(())
+                },
+            )?;
         };
         // format `ItemKind`
         self.value.format(formatted_code, formatter)?;

swayfmt/tests/mod.rs

@@ -1319,7 +1319,7 @@ impl MyContract for Contract {
             // Overindented comment, underindented else
     else if self == PrimaryColor::Red {
             true
-        } // Trailing multiline comment should be newlined
+        } // Trailing comment
     // Underindented comment
             // Overindented else
                 else {
@@ -1337,8 +1337,7 @@ impl MyContract for Contract {
         // Overindented comment, underindented else
         else if self == PrimaryColor::Red {
             true
-        }
-        // Trailing multiline comment should be newlined
+        } // Trailing comment
         // Underindented comment
         // Overindented else
         else {
@@ -1468,7 +1467,7 @@ abi MyAbi : ABIsupertrait {
     fn bar();
 } {
     fn baz() {
-        Self::foo()     // supertrait method usage
+        Self::foo() // supertrait method usage
     }
 }
 
@@ -1479,7 +1478,7 @@ impl ABIsupertrait for Contract {
 // The implementation of MyAbi for Contract must also implement ABIsupertrait
 impl MyAbi for Contract {
     fn bar() {
-        Self::foo()     // supertrait method usage
+        Self::foo() // supertrait method usage
     }
 }
 "#,