Add storage formatting to sway-fmt-v2 (#2124)

sway-fmt-v2/src/config/heuristics.rs

@@ -1,8 +1,8 @@
 //! Configuration options related to heuristics.
 use crate::constants::{
     DEFAULT_ARRAY_WIDTH, DEFAULT_ATTR_FN_LIKE_WIDTH, DEFAULT_CHAIN_WIDTH, DEFAULT_FN_CALL_WIDTH,
-    DEFAULT_MAX_LINE_WIDTH, DEFAULT_SINGLE_LINE_IF_ELSE_WIDTH, DEFAULT_STRUCT_LIT_WIDTH,
-    DEFAULT_STRUCT_VAR_WIDTH,
+    DEFAULT_MAX_LINE_WIDTH, DEFAULT_SINGLE_LINE_IF_ELSE_WIDTH, DEFAULT_STRUCTURE_LIT_WIDTH,
+    DEFAULT_STRUCTURE_VAR_WIDTH,
 };
 use serde::{Deserialize, Serialize};
 
@@ -131,9 +131,9 @@ impl WidthHeuristics {
             fn_call_width: (DEFAULT_FN_CALL_WIDTH as f32 * max_width_ratio).round() as usize,
             attr_fn_like_width: (DEFAULT_ATTR_FN_LIKE_WIDTH as f32 * max_width_ratio).round()
                 as usize,
-            structure_lit_width: (DEFAULT_STRUCT_LIT_WIDTH as f32 * max_width_ratio).round()
+            structure_lit_width: (DEFAULT_STRUCTURE_LIT_WIDTH as f32 * max_width_ratio).round()
                 as usize,
-            structure_field_width: (DEFAULT_STRUCT_VAR_WIDTH as f32 * max_width_ratio).round()
+            structure_field_width: (DEFAULT_STRUCTURE_VAR_WIDTH as f32 * max_width_ratio).round()
                 as usize,
             array_width: (DEFAULT_ARRAY_WIDTH as f32 * max_width_ratio).round() as usize,
             chain_width: (DEFAULT_CHAIN_WIDTH as f32 * max_width_ratio).round() as usize,

sway-fmt-v2/src/constants.rs

@@ -14,10 +14,10 @@ pub const DEFAULT_TAB_SPACES: usize = 4;
 pub const DEFAULT_FN_CALL_WIDTH: usize = 60;
 /// Default max width of the args of a function-like attributes before falling back to vertical formatting.
 pub const DEFAULT_ATTR_FN_LIKE_WIDTH: usize = 70;
-/// Default max width in the body of a struct literal before falling back to vertical formatting.
-pub const DEFAULT_STRUCT_LIT_WIDTH: usize = 18;
-/// Default max width in the body of a struct variant before falling back to vertical formatting.
-pub const DEFAULT_STRUCT_VAR_WIDTH: usize = 35;
+/// Default max width in the body of a user-defined structure literal before falling back to vertical formatting.
+pub const DEFAULT_STRUCTURE_LIT_WIDTH: usize = 18;
+/// Default max width in the body of a user-defined structure field before falling back to vertical formatting.
+pub const DEFAULT_STRUCTURE_VAR_WIDTH: usize = 35;
 /// Default Maximum width of an array literal before falling back to vertical formatting.
 pub const DEFAULT_ARRAY_WIDTH: usize = 60;
 /// Defalt width threshold for an array element to be considered short.

sway-fmt-v2/src/fmt.rs

@@ -271,4 +271,69 @@ pub const TEST1: u16 = 10;"#;
             Formatter::format(&mut formatter, Arc::from(sway_code_to_format), None).unwrap();
         assert_eq!(correct_sway_code, formatted_sway_code)
     }
+    #[test]
+    fn test_storage_without_alignment() {
+        let sway_code_to_format = r#"contract;
+
+storage{foo:Test,bar
+: 
+    Test
+, baz: u64 } 
+"#;
+        let correct_sway_code = r#"contract;
+
+storage {
+    foo: Test,
+    bar: Test,
+    baz: u64,
+}"#;
+
+        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_with_alignment() {
+        let sway_code_to_format = r#"contract;
+
+storage {
+ long_var_name: Type1,
+      var2: Type2,
+}
+"#;
+        let correct_sway_code = r#"contract;
+
+storage {
+    long_var_name : Type1,
+    var2          : Type2,
+}"#;
+
+        let mut config = Config::default();
+        config.structures.field_alignment = FieldAlignment::AlignFields(50);
+        let mut formatter = get_formatter(config, Shape::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_single_line() {
+        let sway_code_to_format = r#"contract;
+
+storage {
+ long_var_name: Type1,
+      var2: Type2,
+}
+"#;
+        let correct_sway_code = r#"contract;
+
+storage { long_var_name: Type1, var2: Type2 }"#;
+        let mut config = Config::default();
+        config.structures.small_structures_single_line = true;
+        config.whitespace.max_width = 300;
+        let mut formatter = get_formatter(config, Shape::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_enum.rs

@@ -30,7 +30,7 @@ impl Format for ItemEnum {
             .to_width_heuristics(&formatter.config.whitespace);
         let enum_lit_width = width_heuristics.structure_lit_width;
 
-        let multiline = !enum_lit_single_line || self.get_formatted_len() > enum_lit_width;
+        let multiline = !enum_lit_single_line || self.get_formatted_len()? > enum_lit_width;
 
         format_enum(self, formatted_code, formatter, multiline)?;
         Ok(())
@@ -181,10 +181,10 @@ fn format_enum(
 }
 
 impl ItemLen for ItemEnum {
-    fn get_formatted_len(&self) -> usize {
+    fn get_formatted_len(&self) -> Result<usize, FormatterError> {
         // TODO while determininig the length we may want to format to some degree and take length.
         let str_item = &self.span().as_str().len();
-        *str_item as usize
+        Ok(*str_item as usize)
     }
 }
 

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

@@ -1,15 +1,199 @@
 use crate::{
+    config::{items::ItemBraceStyle, user_def::FieldAlignment},
     fmt::{Format, FormattedCode, Formatter},
+    utils::{bracket::CurlyBrace, item_len::ItemLen},
     FormatterError,
 };
-use sway_parse::ItemStorage;
+use std::fmt::Write;
+use sway_parse::{
+    token::{Delimiter, PunctKind},
+    ItemStorage,
+};
+use sway_types::Spanned;
 
 impl Format for ItemStorage {
     fn format(
         &self,
-        _formatted_code: &mut FormattedCode,
-        _formatter: &mut Formatter,
+        formatted_code: &mut FormattedCode,
+        formatter: &mut Formatter,
+    ) -> Result<(), FormatterError> {
+        // Should we format small storage into single line
+        let storage_single_line = formatter.config.structures.small_structures_single_line;
+
+        // Get the width limit of a storage to be formatted into single line if storage_single_line is true
+        let config_whitespace = formatter.config.whitespace;
+        let width_heuristics = formatter
+            .config
+            .heuristics
+            .heuristics_pref
+            .to_width_heuristics(&config_whitespace);
+        let storage_width = width_heuristics.structure_lit_width;
+
+        let multiline = !storage_single_line || self.get_formatted_len()? > storage_width;
+        format_storage(self, formatter, formatted_code, multiline)?;
+        Ok(())
+    }
+}
+
+impl ItemLen for ItemStorage {
+    fn get_formatted_len(&self) -> Result<usize, FormatterError> {
+        // Format to single line and return the length
+        let mut str_item = String::new();
+        let mut formatter = Formatter::default();
+        format_storage(self, &mut formatter, &mut str_item, false)?;
+        Ok(str_item.len() as usize)
+    }
+}
+
+fn format_storage(
+    item_storage: &ItemStorage,
+    formatter: &mut Formatter,
+    formatted_code: &mut String,
+    multiline: bool,
+) -> Result<(), FormatterError> {
+    // Add storage token
+    formatted_code.push_str(item_storage.storage_token.span().as_str());
+    let fields = item_storage.fields.clone().into_inner();
+
+    // Handle openning brace
+    ItemStorage::open_curly_brace(formatted_code, formatter)?;
+    if multiline {
+        formatted_code.push('\n');
+        // Determine alignment tactic
+        match formatter.config.structures.field_alignment {
+            FieldAlignment::AlignFields(storage_field_align_threshold) => {
+                let value_pairs = fields.value_separator_pairs;
+                // In first iteration we are going to be collecting the lengths of the struct fields.
+                let field_length: Vec<usize> = value_pairs
+                    .iter()
+                    .map(|field| field.0.name.as_str().len())
+                    .collect();
+
+                // Find the maximum length in the `field_length` vector that is still smaller than `storage_field_align_threshold`.
+                // `max_valid_field_length`: the length of the field that we are taking as a reference to align.
+                let mut max_valid_field_length = 0;
+                field_length.iter().for_each(|length| {
+                    if *length > max_valid_field_length && *length < storage_field_align_threshold {
+                        max_valid_field_length = *length;
+                    }
+                });
+
+                let mut value_pairs_iter = value_pairs.iter().enumerate().peekable();
+                for (field_index, field) in value_pairs_iter.clone() {
+                    formatted_code.push_str(&formatter.shape.indent.to_string(formatter));
+
+                    let type_field = &field.0;
+                    // Add name
+                    formatted_code.push_str(type_field.name.as_str());
+
+                    // `current_field_length`: the length of the current field that we are trying to format.
+                    let current_field_length = field_length[field_index];
+                    if current_field_length < max_valid_field_length {
+                        // We need to add alignment between `:` and `ty`
+                        let mut required_alignment = max_valid_field_length - current_field_length;
+                        while required_alignment != 0 {
+                            formatted_code.push(' ');
+                            required_alignment -= 1;
+                        }
+                    }
+                    // Add `:`, `ty` & `CommaToken`
+                    //
+                    // TODO(#2101): We are currently converting ty to string directly but
+                    // we will probably need to format `ty` before adding.
+                    write!(
+                        formatted_code,
+                        " {} {}",
+                        type_field.colon_token.ident().as_str(),
+                        type_field.ty.span().as_str(),
+                    )?;
+                    if value_pairs_iter.peek().is_some() {
+                        writeln!(formatted_code, "{}", field.1.span().as_str())?;
+                    } else if let Some(final_value) = &fields.final_value_opt {
+                        final_value.format(formatted_code, formatter)?;
+                    }
+                }
+            }
+            FieldAlignment::Off => {
+                let mut value_pairs_iter = fields.value_separator_pairs.iter().peekable();
+                for field in value_pairs_iter.clone() {
+                    formatted_code.push_str(&formatter.shape.indent.to_string(formatter));
+                    let item_field = &field.0;
+                    item_field.format(formatted_code, formatter)?;
+
+                    if value_pairs_iter.peek().is_some() {
+                        writeln!(formatted_code, "{}", field.1.span().as_str())?;
+                    }
+                }
+                if let Some(final_value) = &fields.final_value_opt {
+                    formatted_code.push_str(&formatter.shape.indent.to_string(formatter));
+                    final_value.format(formatted_code, formatter)?;
+                    writeln!(formatted_code, "{}", PunctKind::Comma.as_char())?;
+                }
+            }
+        }
+    } else {
+        // non-multiline formatting
+        formatted_code.push(' ');
+        let mut value_pairs_iter = fields.value_separator_pairs.iter().peekable();
+        for field in value_pairs_iter.clone() {
+            let type_field = &field.0;
+            type_field.format(formatted_code, formatter)?;
+
+            if value_pairs_iter.peek().is_some() {
+                write!(formatted_code, "{} ", field.1.span().as_str())?;
+            }
+        }
+        if let Some(final_value) = &fields.final_value_opt {
+            final_value.format(formatted_code, formatter)?;
+            formatted_code.push(' ');
+        } else {
+            formatted_code.pop();
+            formatted_code.pop();
+            formatted_code.push(' ');
+        }
+    }
+
+    // Handle closing brace
+    ItemStorage::close_curly_brace(formatted_code, formatter)?;
+    Ok(())
+}
+
+impl CurlyBrace for ItemStorage {
+    fn open_curly_brace(
+        line: &mut String,
+        formatter: &mut Formatter,
+    ) -> Result<(), FormatterError> {
+        let brace_style = formatter.config.items.item_brace_style;
+        let extra_width = formatter.config.whitespace.tab_spaces;
+        let mut shape = formatter.shape;
+        let open_brace = Delimiter::Brace.as_open_char();
+        match brace_style {
+            ItemBraceStyle::AlwaysNextLine => {
+                // Add opening brace to the next line.
+                write!(line, "\n{}", open_brace)?;
+                shape = shape.block_indent(extra_width);
+            }
+            _ => {
+                // Add opening brace to the same line
+                write!(line, " {}", open_brace)?;
+                shape = shape.block_indent(extra_width);
+            }
+        }
+
+        formatter.shape = shape;
+        Ok(())
+    }
+    fn close_curly_brace(
+        line: &mut String,
+        formatter: &mut Formatter,
     ) -> Result<(), FormatterError> {
-        todo!()
+        line.push(Delimiter::Brace.as_close_char());
+        // shrink_left would return error if the current indentation level is becoming < 0, in that
+        // case we should use the Shape::default() which has 0 indentation level.
+        formatter.shape = formatter
+            .shape
+            .shrink_left(formatter.config.whitespace.tab_spaces)
+            .unwrap_or_default();
+        Ok(())
     }
 }

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

@@ -30,7 +30,7 @@ impl Format for ItemStruct {
             .to_width_heuristics(&formatter.config.whitespace);
         let struct_lit_width = width_heuristics.structure_lit_width;
 
-        let multiline = !struct_lit_single_line || self.get_formatted_len() > struct_lit_width;
+        let multiline = !struct_lit_single_line || self.get_formatted_len()? > struct_lit_width;
 
         format_struct(self, formatted_code, formatter, multiline)?;
         Ok(())
@@ -186,10 +186,10 @@ fn format_struct(
 }
 
 impl ItemLen for ItemStruct {
-    fn get_formatted_len(&self) -> usize {
+    fn get_formatted_len(&self) -> Result<usize, FormatterError> {
         // TODO while determininig the length we may want to format to some degree and take length.
         let str_item = &self.span().as_str().len();
-        *str_item as usize
+        Ok(*str_item as usize)
     }
 }
 

sway-fmt-v2/src/utils/item_len.rs

@@ -1,3 +1,4 @@
+use crate::fmt::FormatterError;
 pub trait ItemLen {
-    fn get_formatted_len(&self) -> usize;
+    fn get_formatted_len(&self) -> Result<usize, FormatterError>;
 }

sway-fmt-v2/src/utils/punctuated.rs

@@ -3,7 +3,7 @@ use crate::{
     FormatterError,
 };
 use std::fmt::Write;
-use sway_parse::{punctuated::Punctuated, TypeField};
+use sway_parse::{punctuated::Punctuated, StorageField, TypeField};
 use sway_types::Spanned;
 
 impl<T, P> Format for Punctuated<T, P>
@@ -55,3 +55,20 @@ impl Format for TypeField {
         Ok(())
     }
 }
+
+impl Format for StorageField {
+    fn format(
+        &self,
+        formatted_code: &mut FormattedCode,
+        _formatter: &mut Formatter,
+    ) -> Result<(), FormatterError> {
+        write!(
+            formatted_code,
+            "{}{} {}",
+            self.name.span().as_str(),
+            self.colon_token.span().as_str(),
+            self.ty.span().as_str()
+        )?;
+        Ok(())
+    }
+}