Add logic for inserting comments for sway-fmt-v2 (#2311)

sway-fmt-v2/Cargo.toml

@@ -12,6 +12,7 @@ description = "Sway sway-fmt-v2."
 [dependencies]
 anyhow = "1"
 forc-util = { version = "0.18.1", path = "../forc-util" }
+ropey = "1.5"
 serde = { version = "1.0", features = ["derive"] }
 serde_ignored = "0.1"
 sway-core = { version = "0.18.1", path = "../sway-core" }

sway-fmt-v2/src/error.rs

@@ -9,6 +9,8 @@ pub enum FormatterError {
     FormatError(#[from] std::fmt::Error),
     #[error("Error while lexing file: {0}")]
     LexError(#[from] sway_parse::LexError),
+    #[error("Error while adding comments")]
+    CommentError,
 }
 
 #[derive(Debug, Error)]

sway-fmt-v2/src/fmt.rs

@@ -1,5 +1,6 @@
 use crate::utils::{
-    indent_style::Shape, newline_style::apply_newline_style, program_type::insert_program_type,
+    comments::handle_comments, indent_style::Shape, newline_style::apply_newline_style,
+    program_type::insert_program_type,
 };
 use std::{path::Path, sync::Arc};
 use sway_core::BuildConfig;
@@ -42,11 +43,11 @@ impl Formatter {
     ) -> Result<FormattedCode, FormatterError> {
         let path = build_config.map(|build_config| build_config.canonical_root_module());
         let src_len = src.len();
-        let module = sway_parse::parse_file(src, path)?;
+        let module = sway_parse::parse_file(src.clone(), path.clone())?;
         // Get parsed items
-        let items = module.items;
+        let items = &module.items;
         // Get the program type (script, predicate, contract or library)
-        let program_type = module.kind;
+        let program_type = &module.kind;
 
         // Formatted code will be pushed here with raw newline stlye.
         // Which means newlines are not converted into system-specific versions until `apply_newline_style()`.
@@ -68,6 +69,14 @@ impl Formatter {
         }
 
         let mut formatted_code = String::from(&raw_formatted_code);
+        // Add comments
+        handle_comments(
+            src,
+            &module,
+            Arc::from(formatted_code.clone()),
+            path,
+            &mut formatted_code,
+        )?;
         // Replace newlines with specified `NewlineStyle`
         apply_newline_style(
             self.config.whitespace.newline_style,
@@ -337,8 +346,10 @@ storage {
         let correct_sway_code = r#"contract;
 
 storage {
-    long_var_name : Type1,
-    var2          : Type2,
+    long_var_name : Type1 = Type1 {
+    },
+    var2          : Type2 = Type2 {
+    },
 }"#;
 
         let mut formatter = Formatter::default();
@@ -358,10 +369,12 @@ storage {
 "#;
         let correct_sway_code = r#"contract;
 
-storage { long_var_name: Type1, var2: Type2 }"#;
+storage { long_var_name: Type1 = Type1 {
+    }, var2: Type2 = Type2 {
+    } }"#;
         let mut formatter = Formatter::default();
         formatter.config.structures.small_structures_single_line = true;
-        formatter.config.whitespace.max_width = 300;
+        formatter.config.whitespace.max_width = 700;
         let formatted_sway_code =
             Formatter::format(&mut formatter, Arc::from(sway_code_to_format), None).unwrap();
         assert_eq!(correct_sway_code, formatted_sway_code)
@@ -391,12 +404,10 @@ struct Type1 {
     x: u64,
     y: u64,
 }
-
 struct Type2 {
     w: b256,
     z: bool,
 }
-
 storage {
     var1: Type1 = Type1 {
         x: 0,
@@ -485,4 +496,188 @@ trait CompSciStudent: Programmer + Student {
             Formatter::format(&mut formatter, Arc::from(sway_code_to_format), None).unwrap();
         assert_eq!(correct_sway_code, formatted_sway_code)
     }
+
+    #[test]
+    fn test_struct_comments() {
+        let sway_code_to_format = r#"contract;
+// This is a comment, for this one to be placed correctly we need to have Module visitor implemented
+pub struct Foo { // Here is a comment
+
+
+
+    // Trying some ASCII art
+    baz:u64,
+
+
+
+
+    bazzz:u64//  ________ ___  ___  _______   ___               ___       ________  ________  ________
+             // |\  _____\\  \|\  \|\  ___ \ |\  \             |\  \     |\   __  \|\   __  \|\   ____\
+             // \ \  \__/\ \  \\\  \ \   __/|\ \  \            \ \  \    \ \  \|\  \ \  \|\ /\ \  \___|_
+             //  \ \   __\\ \  \\\  \ \  \_|/_\ \  \            \ \  \    \ \   __  \ \   __  \ \_____  \
+             //   \ \  \_| \ \  \\\  \ \  \_|\ \ \  \____        \ \  \____\ \  \ \  \ \  \|\  \|____|\  \
+             //    \ \__\   \ \_______\ \_______\ \_______\       \ \_______\ \__\ \__\ \_______\____\_\  \
+             //     \|__|    \|_______|\|_______|\|_______|        \|_______|\|__|\|__|\|_______|\_________\
+             //                                                                                  \|_________|
+}
+// This is a comment
+"#;
+        let correct_sway_code = r#"contract;
+
+// This is a comment, for this one to be placed correctly we need to have Module visitor implemented
+pub struct Foo { // Here is a comment
+
+
+
+    // Trying some ASCII art
+    baz: u64,
+    bazzz: u64,//  ________ ___  ___  _______   ___               ___       ________  ________  ________
+             // |\  _____\\  \|\  \|\  ___ \ |\  \             |\  \     |\   __  \|\   __  \|\   ____\
+             // \ \  \__/\ \  \\\  \ \   __/|\ \  \            \ \  \    \ \  \|\  \ \  \|\ /\ \  \___|_
+             //  \ \   __\\ \  \\\  \ \  \_|/_\ \  \            \ \  \    \ \   __  \ \   __  \ \_____  \
+             //   \ \  \_| \ \  \\\  \ \  \_|\ \ \  \____        \ \  \____\ \  \ \  \ \  \|\  \|____|\  \
+             //    \ \__\   \ \_______\ \_______\ \_______\       \ \_______\ \__\ \__\ \_______\____\_\  \
+             //     \|__|    \|_______|\|_______|\|_______|        \|_______|\|__|\|__|\|_______|\_________\
+             //                                                                                  \|_________|
+}
+// This is a comment"#;
+        let mut formatter = Formatter::default();
+        let formatted_sway_code =
+            Formatter::format(&mut formatter, Arc::from(sway_code_to_format), None).unwrap();
+        assert_eq!(correct_sway_code, formatted_sway_code)
+    }
+
+    #[test]
+    fn test_enum_comments() {
+        let sway_code_to_format = r#"contract;
+pub enum Bazz { // Here is a comment
+    // Trying some ASCII art
+    baz: (),
+
+
+
+
+
+    bazzz: (),//-----
+              //--D--
+              //-----
+}
+"#;
+        let correct_sway_code = r#"contract;
+
+pub enum Bazz { // Here is a comment
+    // Trying some ASCII art
+    baz: (),
+    bazzz: (),//-----
+              //--D--
+              //-----
+}"#;
+        let mut formatter = Formatter::default();
+        let formatted_sway_code =
+            Formatter::format(&mut formatter, Arc::from(sway_code_to_format), None).unwrap();
+        assert_eq!(correct_sway_code, formatted_sway_code);
+    }
+
+    #[test]
+    fn test_fn_comments() {
+        let sway_code_to_format = r#"contract;
+// This is a comment before a fn
+// This is another comment before a fn
+fn hello_world( baz: /* this is a comment */ u64) { // This is a comment inside the block
+}
+"#;
+        let correct_sway_code = r#"contract;
+
+// This is a comment before a fn
+// This is another comment before a fn
+fn hello_world(baz: /* this is a comment */ u64) { // This is a comment inside the block
+}"#;
+
+        let mut formatter = Formatter::default();
+        let formatted_sway_code =
+            Formatter::format(&mut formatter, Arc::from(sway_code_to_format), None).unwrap();
+        assert_eq!(correct_sway_code, formatted_sway_code);
+    }
+
+    #[test]
+    fn test_abi_comments() {
+        let sway_code_to_format = r#"contract;
+// This is an abi
+abi StorageMapExample {
+    // insert_into_map is blah blah
+    #[storage(write)] // this is some other comment
+    fn insert_into_map(key: u64, value: u64); // this is the last comment inside the StorageMapExample
+}"#;
+        let correct_sway_code = r#"contract;
+
+// This is an abi
+abi StorageMapExample {
+    // insert_into_map is blah blah
+    #[storage(write)] // this is some other comment
+    fn insert_into_map(key: u64, value: u64); // this is the last comment inside the StorageMapExample
+}"#;
+        let mut formatter = Formatter::default();
+        let formatted_sway_code =
+            Formatter::format(&mut formatter, Arc::from(sway_code_to_format), None).unwrap();
+        assert_eq!(correct_sway_code, formatted_sway_code);
+    }
+
+    #[test]
+    fn test_const_comments() {
+        let sway_code_to_format = r#"contract;
+pub const /* TEST: blah blah tests */ TEST: u16 = 10; // This is a comment next to a const"#;
+        let correct_sway_code = r#"contract;
+
+pub const /* TEST: blah blah tests */ TEST: u16 = 10;"#; // Comment next to const is not picked up by the lexer see: #2356
+        let mut formatter = Formatter::default();
+        let formatted_sway_code =
+            Formatter::format(&mut formatter, Arc::from(sway_code_to_format), None).unwrap();
+        assert_eq!(correct_sway_code, formatted_sway_code);
+    }
+    #[test]
+    fn test_storage_comments() {
+        let sway_code_to_format = r#"contract;
+storage {
+    // Testing a comment inside storage
+    long_var_name: Type1=Type1{},
+    // Testing another comment
+    var2: Type2 = Type2{} // This is the last comment
+}"#;
+        let correct_sway_code = r#"contract;
+
+storage {
+    // Testing a comment inside storage
+    long_var_name: Type1 = Type1 {
+    },
+    // Testing another comment
+    var2: Type2 = Type2 {
+    }, // This is the last comment
+}"#;
+        let mut formatter = Formatter::default();
+        let formatted_sway_code =
+            Formatter::format(&mut formatter, Arc::from(sway_code_to_format), None).unwrap();
+        assert_eq!(correct_sway_code, formatted_sway_code);
+    }
+
+    #[test]
+    fn test_trait_comments() {
+        let sway_code_to_format = r#"contract;
+// This is the programmer trait
+trait Programmer {
+    // Returns fav languages of this Programmer.
+    fn fav_language(self) -> String;
+}"#;
+        let correct_sway_code = r#"contract;
+
+// This is the programmer trait
+trait Programmer {
+    // Returns fav languages of this Programmer.
+    fn fav_language(self) -> String;
+}"#;
+
+        let mut formatter = Formatter::default();
+        let formatted_sway_code =
+            Formatter::format(&mut formatter, Arc::from(sway_code_to_format), None).unwrap();
+        assert_eq!(correct_sway_code, formatted_sway_code)
+    }
 }

sway-fmt-v2/src/items/item_abi.rs

@@ -1,7 +1,11 @@
 use crate::{
     config::items::ItemBraceStyle,
     fmt::{Format, FormattedCode, Formatter},
-    utils::{attribute::FormatDecl, bracket::CurlyBrace},
+    utils::{
+        attribute::FormatDecl,
+        bracket::CurlyBrace,
+        comments::{ByteSpan, LeafSpans},
+    },
     FormatterError,
 };
 use std::fmt::Write;
@@ -130,3 +134,15 @@ impl CurlyBrace for ItemAbi {
         Ok(())
     }
 }
+
+impl LeafSpans for ItemAbi {
+    fn leaf_spans(&self) -> Vec<ByteSpan> {
+        let mut collected_spans = vec![ByteSpan::from(self.abi_token.span())];
+        collected_spans.push(ByteSpan::from(self.name.span()));
+        collected_spans.append(&mut self.abi_items.leaf_spans());
+        if let Some(abi_defs) = &self.abi_defs_opt {
+            collected_spans.append(&mut abi_defs.leaf_spans());
+        }
+        collected_spans
+    }
+}

sway-fmt-v2/src/items/item_const.rs

@@ -1,5 +1,6 @@
 use crate::{
     fmt::{Format, FormattedCode, Formatter},
+    utils::comments::{ByteSpan, LeafSpans},
     FormatterError,
 };
 use std::fmt::Write;
@@ -44,3 +45,22 @@ impl Format for ItemConst {
         Ok(())
     }
 }
+
+impl LeafSpans for ItemConst {
+    fn leaf_spans(&self) -> Vec<ByteSpan> {
+        let mut collected_spans = Vec::new();
+        if let Some(visibility) = &self.visibility {
+            collected_spans.push(ByteSpan::from(visibility.span()));
+        }
+        collected_spans.push(ByteSpan::from(self.const_token.span()));
+        collected_spans.push(ByteSpan::from(self.name.span()));
+        if let Some(ty) = &self.ty_opt {
+            collected_spans.append(&mut ty.leaf_spans());
+            // TODO: determine if we allow comments in between `:` and ty
+        }
+        collected_spans.push(ByteSpan::from(self.eq_token.span()));
+        collected_spans.append(&mut self.expr.leaf_spans());
+        collected_spans.push(ByteSpan::from(self.semicolon_token.span()));
+        collected_spans
+    }
+}

sway-fmt-v2/src/items/item_enum.rs

@@ -1,7 +1,11 @@
 use crate::{
     config::{items::ItemBraceStyle, user_def::FieldAlignment},
     fmt::{Format, FormattedCode, Formatter},
-    utils::{bracket::CurlyBrace, item::ItemLenChars},
+    utils::{
+        bracket::CurlyBrace,
+        comments::{ByteSpan, LeafSpans},
+        item::ItemLenChars,
+    },
     FormatterError,
 };
 use std::fmt::Write;
@@ -235,3 +239,18 @@ impl CurlyBrace for ItemEnum {
         Ok(())
     }
 }
+impl LeafSpans for ItemEnum {
+    fn leaf_spans(&self) -> Vec<ByteSpan> {
+        let mut collected_spans = Vec::new();
+        if let Some(visibility) = &self.visibility {
+            collected_spans.push(ByteSpan::from(visibility.span()));
+        }
+        collected_spans.push(ByteSpan::from(self.enum_token.span()));
+        collected_spans.push(ByteSpan::from(self.name.span()));
+        if let Some(generics) = &self.generics {
+            collected_spans.push(ByteSpan::from(generics.parameters.span()))
+        }
+        collected_spans.append(&mut self.fields.leaf_spans());
+        collected_spans
+    }
+}