From 74421d47fd388732a93ac47074f7614812b98bce Mon Sep 17 00:00:00 2001
From: Jonathan Kelley <jkelleyrtp@gmail.com>
Date: Sat, 7 Jan 2023 14:09:32 -0800
Subject: [PATCH] fix comments being eaten in autofmt

---
 packages/autofmt/src/buffer.rs                | 10 ++++----
 packages/autofmt/src/element.rs               |  1 +
 packages/autofmt/tests/samples.rs             |  3 +++
 .../autofmt/tests/samples/commentshard.rsx    | 23 +++++++++++++++++++
 4 files changed, 33 insertions(+), 4 deletions(-)
 create mode 100644 packages/autofmt/tests/samples/commentshard.rsx

diff --git a/packages/autofmt/src/buffer.rs b/packages/autofmt/src/buffer.rs
index fd0b6b60d..95bc80673 100644
--- a/packages/autofmt/src/buffer.rs
+++ b/packages/autofmt/src/buffer.rs
@@ -126,8 +126,13 @@ impl Buffer {
 
     pub fn write_body_no_indent(&mut self, children: &[BodyNode]) -> Result {
         let last_child = children.len();
+        let iter = children.iter().peekable().enumerate();
+
+        for (idx, child) in iter {
+            if self.current_span_is_primary(child.span()) {
+                self.write_comments(child.span())?;
+            }
 
-        for (idx, child) in children.iter().enumerate() {
             match child {
                 // check if the expr is a short
                 BodyNode::RawExpr { .. } => {
@@ -138,9 +143,6 @@ impl Buffer {
                     }
                 }
                 _ => {
-                    if self.current_span_is_primary(child.span()) {
-                        self.write_comments(child.span())?;
-                    }
                     self.tabbed_line()?;
                     self.write_ident(child)?;
                 }
diff --git a/packages/autofmt/src/element.rs b/packages/autofmt/src/element.rs
index 310641121..3be25217c 100644
--- a/packages/autofmt/src/element.rs
+++ b/packages/autofmt/src/element.rs
@@ -124,6 +124,7 @@ impl Buffer {
                 if !children.is_empty() {
                     self.write_body_indented(children)?;
                 }
+
                 self.tabbed_line()?;
             }
         }
diff --git a/packages/autofmt/tests/samples.rs b/packages/autofmt/tests/samples.rs
index 1f9c359c1..19588fda0 100644
--- a/packages/autofmt/tests/samples.rs
+++ b/packages/autofmt/tests/samples.rs
@@ -30,3 +30,6 @@ twoway! ("key" => key);
 
 // Disabled because we can't handle comments on exprs yet
 twoway! ("multirsx" => multirsx);
+
+// Disabled because we can't handle comments on exprs yet
+twoway! ("commentshard" => commentshard);
diff --git a/packages/autofmt/tests/samples/commentshard.rsx b/packages/autofmt/tests/samples/commentshard.rsx
new file mode 100644
index 000000000..80c2de7a9
--- /dev/null
+++ b/packages/autofmt/tests/samples/commentshard.rsx
@@ -0,0 +1,23 @@
+rsx! {
+    // Comments
+    div {
+        // Comments
+        class: "asdasd",
+
+        // Comments
+        "hello world",
+
+        // Comments
+        expr1,
+
+        // Comments
+        expr2,
+
+        // Comments
+        // Comments
+        // Comments
+        // Comments
+        // Comments
+        expr3
+    }
+}