feat(fmt): comments in item abi (FuelLabs#4005)

## Description

Part of FuelLabs#3938

This PR handles **comment formatting** for the abi,
as well as comments before modules (this is to pass the test suite)

This was a pretty tricky PR to tackle, since I cycled through several
thought processes while trying to finish this.

1) I tried to solve the problem at the parse/lex level - cycled back to
the thought of re-parsing the file to create a special commented AST -
worked a bit on that before I scrapped the whole idea because it just
didn't seem necessary.

2) I then started trying to solve it in the originally discussed way,
attempting to format the comments within each `Format` implementation.
This started off well - it was trivial to navigate within a span, but it
was particularly tricky navigating _between_ spans, and for formatting
the ABI, it turns out that trailing comments can appear after an
attribute, or even after the fn signature. I initially tried to
implement a way to 'look forward', but soon realized it was impossible
unless we want to hold a reference to the unformatted code in each
`Format` implementation (which I guess is what I ended up doing anyway)

3) So, I ended up looking for a way to 'look back', but still used 'look
ahead' for comments:

Comments will always be 'anchored' to the next span - meaning to say,
when a trailing comment exists, it will 'belong' to the next span
instead of the span(s) on the same line. This means that when we know
the locations of 2 spans, we can easily call `write_comments` to insert
the comments in between.

While we mentioned that looking forward is impossible, one place we
_can_ look forward though, is between comments in the unformatted code,
if there are multiple comments to be written at once. This is basically
what `collect_newlines_after_comment_span()` is doing.

### Other notable changes

- Renamed `write_comments_from_comment_map` to `write_comments` to be
more concise.
- Introduce `CommentsContext` to hold both `CommentMap` and
`unformatted_code`, instead of exposing `unformatted_code` at the
Formatter level.
- Changed generic `T` to `ItemKind` within attribute.rs. This is so we
can have access to `span()`.
- Lexing the commented token tree now gives us `CommentKind` info.

## 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.

---------

Co-authored-by: Kaya Gökalp <[email protected]>

examples/wallet_abi/src/main.sw

@@ -7,7 +7,7 @@ abi Wallet {
     #[storage(read, write)]
     fn receive_funds();
     // ANCHOR_END: receive_funds
-    
+
     // ANCHOR: send_funds
     #[storage(read, write)]
     fn send_funds(amount_to_send: u64, recipient_address: Address);

sway-ast/src/item/item_configurable.rs

@@ -20,3 +20,15 @@ pub struct ConfigurableField {
     pub eq_token: EqToken,
     pub initializer: Expr,
 }
+
+impl Spanned for ConfigurableField {
+    fn span(&self) -> Span {
+        Span::join_all([
+            self.name.span(),
+            self.colon_token.span(),
+            self.ty.span(),
+            self.eq_token.span(),
+            self.initializer.span(),
+        ])
+    }
+}

sway-ast/src/token.rs

@@ -106,9 +106,35 @@ impl Delimiter {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
+pub enum CommentKind {
+    /// A newlined comment is a comment with a preceding newline before another token.
+    ///
+    /// # Examples
+    ///
+    /// ```sway
+    /// pub fn main() -> bool {
+    ///
+    ///     // Newlined comment
+    ///     true
+    /// }
+    /// ```
+    Newlined,
+
+    /// A trailing comment is a comment without a preceding newline before another token.
+    ///
+    /// # Examples
+    ///
+    /// ```sway
+    /// var foo = 1; // Trailing comment
+    /// ```
+    Trailing,
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
 pub struct Comment {
     pub span: Span,
+    pub comment_kind: CommentKind,
 }
 
 impl Spanned for Comment {

sway-parse/src/token.rs

@@ -5,13 +5,13 @@ use std::path::PathBuf;
 use std::sync::Arc;
 use sway_ast::literal::{LitChar, LitInt, LitIntType, LitString, Literal};
 use sway_ast::token::{
-    Comment, CommentedGroup, CommentedTokenStream, CommentedTokenTree, Delimiter, DocComment,
-    DocStyle, GenericTokenTree, Punct, PunctKind, Spacing, TokenStream,
+    Comment, CommentKind, CommentedGroup, CommentedTokenStream, CommentedTokenTree, Delimiter,
+    DocComment, DocStyle, GenericTokenTree, Punct, PunctKind, Spacing, TokenStream,
 };
 use sway_error::error::CompileError;
 use sway_error::handler::{ErrorEmitted, Handler};
 use sway_error::lex_error::{LexError, LexErrorKind};
-use sway_types::{Ident, Span};
+use sway_types::{Ident, Span, Spanned};
 use unicode_xid::UnicodeXID;
 
 #[extension_trait]
@@ -138,8 +138,41 @@ pub fn lex_commented(
         if character == '/' {
             match l.stream.peek() {
                 Some((_, '/')) => {
-                    let ctt =
-                        lex_line_comment(&mut l, end, index, file_start_offset, gather_module_docs);
+                    // search_end is the index at which we stop looking backwards for
+                    // a newline
+                    let search_end = token_trees
+                        .last()
+                        .map(|tt| {
+                            if let CommentedTokenTree::Tree(t) = tt {
+                                t.span().end()
+                            } else {
+                                0
+                            }
+                        })
+                        .unwrap_or_default();
+
+                    let has_newline = src[search_end..index]
+                        .chars()
+                        .rev()
+                        .take_while(|c| c.is_whitespace())
+                        .filter(|&c| c == '\n')
+                        .count()
+                        > 0;
+
+                    let comment_kind = if has_newline {
+                        CommentKind::Newlined
+                    } else {
+                        CommentKind::Trailing
+                    };
+
+                    let ctt = lex_line_comment(
+                        &mut l,
+                        end,
+                        index,
+                        comment_kind,
+                        file_start_offset,
+                        gather_module_docs,
+                    );
                     if let CommentedTokenTree::Tree(GenericTokenTree::DocComment(DocComment {
                         doc_style: DocStyle::Inner,
                         ..
@@ -310,6 +343,7 @@ fn lex_line_comment(
     l: &mut Lexer<'_>,
     end: usize,
     index: usize,
+    comment_kind: CommentKind,
     offset: usize,
     gather_module_docs: bool,
 ) -> CommentedTokenTree {
@@ -348,7 +382,11 @@ fn lex_line_comment(
         };
         CommentedTokenTree::Tree(doc_comment.into())
     } else {
-        Comment { span: sp }.into()
+        Comment {
+            span: sp,
+            comment_kind,
+        }
+        .into()
     }
 }
 
@@ -378,7 +416,13 @@ fn lex_multiline_comment(l: &mut Lexer<'_>, index: usize) -> Option<CommentedTok
                         // We could represent them as several ones, but that's unnecessary.
                         let end = slash_ix + '/'.len_utf8();
                         let span = span(l, start, end);
-                        return Some(Comment { span }.into());
+                        return Some(
+                            Comment {
+                                span,
+                                comment_kind: CommentKind::Newlined,
+                            }
+                            .into(),
+                        );
                     }
                 }
                 Some(_) => {}
@@ -742,7 +786,10 @@ mod tests {
     use std::sync::Arc;
     use sway_ast::{
         literal::{LitChar, Literal},
-        token::{Comment, CommentedTokenTree, CommentedTree, DocComment, DocStyle, TokenTree},
+        token::{
+            Comment, CommentKind, CommentedTokenTree, CommentedTree, DocComment, DocStyle,
+            TokenTree,
+        },
     };
     use sway_error::handler::Handler;
 
@@ -755,7 +802,7 @@ mod tests {
             /* multi-
              * line-
              * comment */
-            bar: i32,
+            bar: i32, // trailing comment
         }
         "#;
         let start = 0;
@@ -786,11 +833,82 @@ mod tests {
             assert_eq!(tts.next().unwrap().span().as_str(), ":");
             assert_eq!(tts.next().unwrap().span().as_str(), "i32");
             assert_eq!(tts.next().unwrap().span().as_str(), ",");
+            assert_matches!(
+                tts.next(),
+                Some(CommentedTokenTree::Comment(Comment {
+                    span,
+                    comment_kind: CommentKind::Trailing,
+                })) if span.as_str() ==  "// trailing comment"
+            );
             assert!(tts.next().is_none());
         }
         assert!(tts.next().is_none());
     }
 
+    #[test]
+    fn lex_comments_check_comment_kind() {
+        let input = r#"
+        // CommentKind::Newlined
+        abi Foo {
+            // CommentKind::Newlined
+            fn bar(); // CommentKind::Trailing
+            // CommentKind::Newlined
+        }
+        "#;
+        let start = 0;
+        let end = input.len();
+        let path = None;
+        let handler = Handler::default();
+        let stream = lex_commented(&handler, &Arc::from(input), start, end, &path).unwrap();
+        assert!(handler.consume().0.is_empty());
+        let mut tts = stream.token_trees().iter();
+
+        assert_matches!(
+            tts.next(),
+            Some(CommentedTokenTree::Comment(Comment {
+                span,
+                comment_kind: CommentKind::Newlined,
+            })) if span.as_str() ==  "// CommentKind::Newlined"
+        );
+        assert_eq!(tts.next().unwrap().span().as_str(), "abi");
+        assert_eq!(tts.next().unwrap().span().as_str(), "Foo");
+
+        {
+            let group = match tts.next() {
+                Some(CommentedTokenTree::Tree(CommentedTree::Group(group))) => group,
+                _ => panic!("expected group"),
+            };
+            let mut tts = group.token_stream.token_trees().iter();
+
+            assert_matches!(
+                tts.next(),
+                Some(CommentedTokenTree::Comment(Comment {
+                    span,
+                    comment_kind: CommentKind::Newlined,
+                })) if span.as_str() ==  "// CommentKind::Newlined"
+            );
+            assert_eq!(tts.next().unwrap().span().as_str(), "fn");
+            assert_eq!(tts.next().unwrap().span().as_str(), "bar");
+            assert_eq!(tts.next().unwrap().span().as_str(), "()");
+            assert_eq!(tts.next().unwrap().span().as_str(), ";");
+            assert_matches!(
+                tts.next(),
+                Some(CommentedTokenTree::Comment(Comment {
+                    span,
+                    comment_kind: CommentKind::Trailing,
+                })) if span.as_str() ==  "// CommentKind::Trailing"
+            );
+            assert_matches!(
+                tts.next(),
+                Some(CommentedTokenTree::Comment(Comment {
+                    span,
+                    comment_kind: CommentKind::Newlined,
+                })) if span.as_str() ==  "// CommentKind::Newlined"
+            );
+            assert!(tts.next().is_none());
+        }
+    }
+
     #[test]
     fn lex_doc_comments() {
         let input = r#"
@@ -811,25 +929,29 @@ mod tests {
         assert_matches!(
             tts.next(),
             Some(CommentedTokenTree::Comment(Comment {
-                span
+                span,
+                comment_kind: CommentKind::Newlined,
             })) if span.as_str() ==  "//none"
         );
         assert_matches!(
             tts.next(),
             Some(CommentedTokenTree::Comment(Comment {
-                span
+                span,
+                comment_kind: CommentKind::Newlined,
             })) if span.as_str() ==  "////none"
         );
         assert_matches!(
             tts.next(),
             Some(CommentedTokenTree::Comment(Comment {
                 span,
+                comment_kind: CommentKind::Newlined,
             })) if span.as_str() ==  "//!inner"
         );
         assert_matches!(
             tts.next(),
             Some(CommentedTokenTree::Comment(Comment {
                 span,
+                comment_kind: CommentKind::Newlined,
             })) if span.as_str() ==  "//! inner"
         );
         assert_matches!(