From 2adaa25af30b30de3cfe9d7bdb2bbfb5e971694b Mon Sep 17 00:00:00 2001
From: tamasfe <me@tamasfe.dev>
Date: Tue, 25 Oct 2022 23:46:40 +0200
Subject: [PATCH] fix(fmt): switch formatting, closes #103

---
 crates/rhai-fmt/src/comments.rs | 60 ++++++++++++++++++++++++---------
 crates/rhai-fmt/src/expr.rs     | 12 ++++---
 testdata/valid/switch.rhai      |  2 +-
 3 files changed, 52 insertions(+), 22 deletions(-)

diff --git a/crates/rhai-fmt/src/comments.rs b/crates/rhai-fmt/src/comments.rs
index da12beb..77a6cc5 100644
--- a/crates/rhai-fmt/src/comments.rs
+++ b/crates/rhai-fmt/src/comments.rs
@@ -15,6 +15,9 @@
 //!   // 3. also with whitespace
 //! }
 //! ```
+//!
+//! This file also contains utilities for comments in
+//! other positions.
 
 #![allow(dead_code)]
 use rhai_rowan::syntax::{
@@ -94,25 +97,22 @@ impl<W: Write> Formatter<W> {
             .take_while(|e| matches!(e.kind(), WHITESPACE | COMMENT_BLOCK | COMMENT_LINE))
             .filter_map(SyntaxElement::into_token);
 
-        match ws_and_comments.next() {
-            Some(t) => {
-                if t.kind() == WHITESPACE && break_count(&t) > 0 {
-                    return Ok(info);
-                } else if t.kind() == WHITESPACE {
-                    if let Some(c) = ws_and_comments.next() {
-                        if c.kind() != WHITESPACE {
-                            self.space();
-                            self.word(c.static_text().trim())?;
-                            info.comment_added = true;
-                        }
+        if let Some(t) = ws_and_comments.next() {
+            if t.kind() == WHITESPACE && break_count(&t) > 0 {
+                return Ok(info);
+            } else if t.kind() == WHITESPACE {
+                if let Some(c) = ws_and_comments.next() {
+                    if c.kind() != WHITESPACE {
+                        self.space();
+                        self.word(c.static_text().trim())?;
+                        info.comment_added = true;
                     }
-                } else {
-                    self.space();
-                    self.word(t.static_text().trim())?;
-                    info.comment_added = true;
                 }
+            } else {
+                self.space();
+                self.word(t.static_text().trim())?;
+                info.comment_added = true;
             }
-            None => {}
         }
 
         Ok(info)
@@ -277,6 +277,34 @@ impl<W: Write> Formatter<W> {
 
         Ok(())
     }
+
+    /// Add comments that are before the node until another node
+    /// is encountered.
+    ///
+    /// Returns the amount of comments that were added.
+    pub(crate) fn add_standalone_comments_before(
+        &mut self,
+        expr: &SyntaxNode,
+    ) -> io::Result<usize> {
+        let mut comments_before = expr
+            .siblings_with_tokens(Direction::Prev)
+            .skip(1)
+            .take_while(|t| t.as_node().is_none())
+            .filter_map(SyntaxElement::into_token)
+            .filter(|t| matches!(t.kind(), COMMENT_LINE | COMMENT_BLOCK))
+            .collect::<Vec<_>>();
+
+        comments_before.reverse();
+
+        let count = comments_before.len();
+
+        for comment in comments_before {
+            self.word(comment.static_text().trim())?;
+            self.hardbreak();
+        }
+
+        Ok(count)
+    }
 }
 
 #[derive(Default)]
diff --git a/crates/rhai-fmt/src/expr.rs b/crates/rhai-fmt/src/expr.rs
index 0f55549..7ef4642 100644
--- a/crates/rhai-fmt/src/expr.rs
+++ b/crates/rhai-fmt/src/expr.rs
@@ -259,6 +259,7 @@ impl<S: Write> Formatter<S> {
         expr: rhai_rowan::ast::ExprSwitch,
     ) -> Result<(), io::Error> {
         self.word("switch ")?;
+
         if let Some(expr) = expr.expr() {
             self.fmt_expr(expr)?;
         }
@@ -269,11 +270,13 @@ impl<S: Write> Formatter<S> {
         if let Some(arm_list) = expr.switch_arm_list() {
             let arm_count = arm_list.arms().count();
             for (i, arm) in arm_list.arms().enumerate() {
-                if i != 0 {
-                    self.hardbreak();
-                }
+                let is_last = i + 1 == arm_count;
+
+                self.add_standalone_comments_before(&arm.syntax())?;
 
-                if let Some(pat) = arm.pattern_expr() {
+                if arm.discard_token().is_some() {
+                    self.word("_ ")?;
+                } else if let Some(pat) = arm.pattern_expr() {
                     self.fmt_expr(pat)?;
                 }
 
@@ -288,7 +291,6 @@ impl<S: Write> Formatter<S> {
                     self.fmt_expr(expr)?;
                 }
 
-                let is_last = i + 1 == arm_count;
                 self.trailing_comma(is_last)?;
             }
         }
diff --git a/testdata/valid/switch.rhai b/testdata/valid/switch.rhai
index 7729263..d349cbc 100644
--- a/testdata/valid/switch.rhai
+++ b/testdata/valid/switch.rhai
@@ -17,6 +17,6 @@ for item in arr {
         // Match another range
         0..100 => print(`A small odd number: ${item}`),
         // Default case
-        _ => print(`Something else: <${item}> is ${type_of(item)}`)
+        _  => print(`Something else: <${item}> is ${type_of(item)}`),
     }
 }