From 60053682b53d9a7485b2736ab5aec9f78394343d Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Wed, 20 Dec 2023 22:46:20 +0100
Subject: [PATCH] fix(fmt): properly format enums

---
 crates/fmt/src/formatter.rs                   | 40 +++++++++----------
 crates/fmt/testdata/EnumVariants/fmt.sol      | 19 +++++++++
 crates/fmt/testdata/EnumVariants/original.sol | 23 +++++++++++
 crates/fmt/tests/formatter.rs                 |  1 +
 4 files changed, 61 insertions(+), 22 deletions(-)
 create mode 100644 crates/fmt/testdata/EnumVariants/fmt.sol
 create mode 100644 crates/fmt/testdata/EnumVariants/original.sol

diff --git a/crates/fmt/src/formatter.rs b/crates/fmt/src/formatter.rs
index 3312b00b73ca..825cbb0321c0 100644
--- a/crates/fmt/src/formatter.rs
+++ b/crates/fmt/src/formatter.rs
@@ -2038,31 +2038,27 @@ impl<'a, W: Write> Visitor for Formatter<'a, W> {
         let enum_name = enumeration.name.safe_unwrap_mut();
         let mut name =
             self.visit_to_chunk(enum_name.loc.start(), Some(enum_name.loc.end()), enum_name)?;
-        name.content = format!("enum {}", name.content);
-        self.write_chunk(&name)?;
-
+        name.content = format!("enum {} ", name.content);
         if enumeration.values.is_empty() {
+            self.write_chunk(&name)?;
             self.write_empty_brackets()?;
         } else {
-            self.surrounded(
-                SurroundingChunk::new(
-                    "{",
-                    Some(enumeration.values.first_mut().unwrap().safe_unwrap().loc.start()),
-                    None,
-                ),
-                SurroundingChunk::new("}", None, Some(enumeration.loc.end())),
-                |fmt, _multiline| {
-                    let values = fmt.items_to_chunks(
-                        Some(enumeration.loc.end()),
-                        enumeration.values.iter_mut().map(|ident| {
-                            let ident = ident.safe_unwrap_mut();
-                            (ident.loc, ident)
-                        }),
-                    )?;
-                    fmt.write_chunks_separated(&values, ",", true)?;
-                    Ok(())
-                },
-            )?;
+            name.content.push('{');
+            self.write_chunk(&name)?;
+
+            self.indented(1, |fmt| {
+                let values = fmt.items_to_chunks(
+                    Some(enumeration.loc.end()),
+                    enumeration.values.iter_mut().map(|ident| {
+                        let ident = ident.safe_unwrap_mut();
+                        (ident.loc, ident)
+                    }),
+                )?;
+                fmt.write_chunks_separated(&values, ",", true)?;
+                writeln!(fmt.buf())?;
+                Ok(())
+            })?;
+            write_chunk!(self, "}}")?;
         }
 
         Ok(())
diff --git a/crates/fmt/testdata/EnumVariants/fmt.sol b/crates/fmt/testdata/EnumVariants/fmt.sol
new file mode 100644
index 000000000000..b33b8846984d
--- /dev/null
+++ b/crates/fmt/testdata/EnumVariants/fmt.sol
@@ -0,0 +1,19 @@
+interface I {
+    enum Empty {}
+
+    /// A modification applied to either `msg.sender` or `tx.origin`. Returned by `readCallers`.
+    enum CallerMode {
+        /// No caller modification is currently active.
+        None
+    }
+
+    /// A modification applied to either `msg.sender` or `tx.origin`. Returned by `readCallers`.
+    enum CallerMode2 {
+        /// No caller modification is currently active.
+        None,
+        /// No caller modification is currently active2.
+        Some
+    }
+
+    function bar() public {}
+}
diff --git a/crates/fmt/testdata/EnumVariants/original.sol b/crates/fmt/testdata/EnumVariants/original.sol
new file mode 100644
index 000000000000..8e146ae0fb57
--- /dev/null
+++ b/crates/fmt/testdata/EnumVariants/original.sol
@@ -0,0 +1,23 @@
+interface I {
+    enum Empty {
+
+    }
+
+    /// A modification applied to either `msg.sender` or `tx.origin`. Returned by `readCallers`.
+    enum CallerMode
+    {/// No caller modification is currently active.
+        None
+    }
+
+    /// A modification applied to either `msg.sender` or `tx.origin`. Returned by `readCallers`.
+    enum CallerMode2
+    {/// No caller modification is currently active.
+        None,/// No caller modification is currently active2.
+
+        Some
+    }
+
+    function bar() public {
+
+    }
+}
\ No newline at end of file
diff --git a/crates/fmt/tests/formatter.rs b/crates/fmt/tests/formatter.rs
index 317cd28efd3e..5ec767b4ed57 100644
--- a/crates/fmt/tests/formatter.rs
+++ b/crates/fmt/tests/formatter.rs
@@ -229,6 +229,7 @@ test_directories! {
     EmitStatement,
     Repros,
     BlockComments,
+    EnumVariants,
 }
 
 test_dir!(SortedImports, TestConfig::skip_compare_ast_eq());