Fixes the new storage namespaces in swayfmt.

FuelLabs · Jun 6, 2024 · 3941df8 · 3941df8
1 parent bb86372
commit 3941df8
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 12 deletions.
diff --git a/swayfmt/src/items/item_storage/mod.rs b/swayfmt/src/items/item_storage/mod.rs
@@ -80,6 +80,14 @@ impl Format for ItemStorage {
                         value_pairs.iter().for_each(|(storage_entry, _)| {
                             collect_field_lengths(storage_entry, ident_size, 0, &mut field_lengths)
                         });
+                        if let Some(final_value) = &entries.final_value_opt {
+                            collect_field_lengths(
+                                &final_value.value,
+                                ident_size,
+                                0,
+                                &mut field_lengths,
+                            );
+                        }
 
                         // Find the maximum length in the `field_length` vector that is still
                         // smaller than `storage_field_align_threshold`.  `max_valid_field_length`:
@@ -100,20 +108,35 @@ impl Format for ItemStorage {
                             field_lengths: &HashMap<IdentUnique, usize>,
                             max_valid_field_length: usize,
                         ) -> Result<(), FormatterError> {
+                            write!(formatted_code, "{}", formatter.indent_to_str()?)?;
                             if let Some(namespace) = &entry.namespace {
                                 entry.name.format(formatted_code, formatter)?;
                                 ItemStorage::open_curly_brace(formatted_code, formatter)?;
+                                writeln!(formatted_code)?;
 
-                                formatter.shape.code_line.update_expr_new_line(true);
-
-                                for e in namespace.clone().into_inner().into_iter() {
+                                for (e, comma_token) in
+                                    namespace.clone().into_inner().value_separator_pairs
+                                {
                                     format_entry(
                                         formatted_code,
                                         formatter,
                                         &e.value,
                                         field_lengths,
                                         max_valid_field_length,
                                     )?;
+                                    writeln!(formatted_code, "{}", comma_token.ident().as_str())?;
+                                }
+                                if let Some(final_value) =
+                                    &namespace.clone().into_inner().final_value_opt
+                                {
+                                    format_entry(
+                                        formatted_code,
+                                        formatter,
+                                        &final_value.value,
+                                        field_lengths,
+                                        max_valid_field_length,
+                                    )?;
+                                    writeln!(formatted_code)?;
                                 }
 
                                 ItemStorage::close_curly_brace(formatted_code, formatter)?;
@@ -155,7 +178,6 @@ impl Format for ItemStorage {
                             Ok(())
                         }
                         for (storage_entry, comma_token) in value_pairs.iter().clone() {
-                            write!(formatted_code, "{}", formatter.indent_to_str()?)?;
                             format_entry(
                                 formatted_code,
                                 formatter,
@@ -166,7 +188,14 @@ impl Format for ItemStorage {
                             writeln!(formatted_code, "{}", comma_token.ident().as_str())?;
                         }
                         if let Some(final_value) = &entries.final_value_opt {
-                            final_value.format(formatted_code, formatter)?;
+                            format_entry(
+                                formatted_code,
+                                formatter,
+                                &final_value.value,
+                                &field_lengths,
+                                max_valid_field_length,
+                            )?;
+                            writeln!(formatted_code)?;
                         }
                     }
                     FieldAlignment::Off => entries.format(formatted_code, formatter)?,

diff --git a/swayfmt/src/utils/language/punctuated.rs b/swayfmt/src/utils/language/punctuated.rs
@@ -13,6 +13,8 @@ use sway_ast::{
 };
 use sway_types::{ast::PunctKind, Ident, Spanned};
 
+use self::shape::ExprKind;
+
 use super::expr::should_write_multiline;
 
 impl<T, P> Format for Punctuated<T, P>
@@ -287,10 +289,18 @@ impl Format for StorageEntry {
             self.name.format(formatted_code, formatter)?;
             ItemStorage::open_curly_brace(formatted_code, formatter)?;
             formatter.shape.code_line.update_expr_new_line(true);
-            namespace
-                .clone()
-                .into_inner()
-                .format(formatted_code, formatter)?;
+            formatter.with_shape(
+                formatter
+                    .shape
+                    .with_code_line_from(LineStyle::Multiline, ExprKind::Struct),
+                |formatter| -> Result<(), FormatterError> {
+                    namespace
+                        .clone()
+                        .into_inner()
+                        .format(formatted_code, formatter)?;
+                    Ok(())
+                },
+            )?;
             ItemStorage::close_curly_brace(formatted_code, formatter)?;
         }
 

diff --git a/swayfmt/tests/mod.rs b/swayfmt/tests/mod.rs
@@ -296,7 +296,7 @@ fn storage_without_alignment() {
         storage {
          var1: Type1=Type1{ foo: 8 },
               var2: Type2=Type2{ bar: 9 },
-              ns1 { var3: u64 = 1 },
+              ns1 { var3: u64 = 1, ns2 { var4: u64 = 1 } },
         }
         "#,
         r#"contract;
@@ -314,7 +314,10 @@ storage {
     var2: Type2 = Type2 { bar: 9 },
     ns1 {
         var3: u64 = 1,
-    }
+        ns2 {
+            var4: u64 = 1,
+        },
+    },
 }
 "#,
     );
@@ -337,7 +340,7 @@ struct Type2 {
 storage {
  long_var_name: Type1=Type1{ foo: 8 },
       var2: Type2=Type2{ bar: 9 },
-      ns1 { var3: u64 = 1 },
+      ns1 { var3: u64 = 1,  ns2 { var4: u64 = 1, },  }, var5: u64 = 1
 }
 "#,
         r#"contract;
@@ -355,7 +358,11 @@ storage {
     var2          : Type2 = Type2 { bar: 9 },
     ns1 {
         var3      : u64 = 1,
+        ns2 {
+            var4  : u64 = 1,
+        },
     },
+    var5          : u64 = 1
 }
 "#,
         &mut formatter,