From 26347556206bc70c6d386878094aca4745a6ff4a Mon Sep 17 00:00:00 2001
From: Jonathan Kelley <jkelleyrtp@gmail.com>
Date: Tue, 5 Jul 2022 02:06:54 -0400
Subject: [PATCH] fix: handle comments around attributes better

---
 packages/autofmt/src/buffer.rs           | 12 +++++++++---
 packages/autofmt/src/element.rs          | 13 ++++++++-----
 packages/autofmt/tests/samples/attrs.rsx | 14 ++++++++++++++
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/packages/autofmt/src/buffer.rs b/packages/autofmt/src/buffer.rs
index 2554c8ff4..fde161a63 100644
--- a/packages/autofmt/src/buffer.rs
+++ b/packages/autofmt/src/buffer.rs
@@ -136,8 +136,14 @@ impl Buffer {
         let mut total = 0;
 
         for attr in attributes {
-            if self.current_element_has_comments(attr.span()) {
-                return 100000;
+            if self.current_span_is_primary(attr.attr.start()) {
+                'line: for line in self.src[..attr.attr.start().start().line - 1].iter().rev() {
+                    match (line.trim().starts_with("//"), line.is_empty()) {
+                        (true, _) => return 100000,
+                        (_, true) => continue 'line,
+                        _ => break 'line,
+                    }
+                }
             }
 
             total += match &attr.attr {
@@ -182,7 +188,7 @@ impl Buffer {
     pub fn retrieve_formatted_expr(&mut self, expr: &Expr) -> &str {
         self.cached_formats
             .entry(Location::new(expr.span().start()))
-            .or_insert(prettyplease::unparse_expr(expr))
+            .or_insert_with(|| prettyplease::unparse_expr(expr))
             .as_str()
     }
 }
diff --git a/packages/autofmt/src/element.rs b/packages/autofmt/src/element.rs
index a778eb259..69966028a 100644
--- a/packages/autofmt/src/element.rs
+++ b/packages/autofmt/src/element.rs
@@ -222,12 +222,13 @@ impl Buffer {
         Ok(())
     }
 
-    pub fn current_element_has_comments(&self, location: Span) -> bool {
+    // make sure the comments are actually relevant to this element.
+    // test by making sure this element is the primary element on this line
+    pub fn current_span_is_primary(&self, location: Span) -> bool {
         let start = location.start();
-        let line_start = start.line;
+        let line_start = start.line - 1;
 
-        // make sure the comments are actually relevant to this element.
-        let this_line = self.src[line_start - 1].as_str();
+        let this_line = self.src[line_start].as_str();
 
         let beginning = if this_line.len() > start.column {
             this_line[..start.column].trim()
@@ -235,6 +236,8 @@ impl Buffer {
             ""
         };
 
+        // dbg!(beginning);
+
         beginning.is_empty()
     }
 
@@ -251,7 +254,7 @@ impl Buffer {
         }
 
         for child in children {
-            if self.current_element_has_comments(child.span()) {
+            if self.current_span_is_primary(child.span()) {
                 'line: for line in self.src[..child.span().start().line - 1].iter().rev() {
                     match (line.trim().starts_with("//"), line.is_empty()) {
                         (true, _) => return None,
diff --git a/packages/autofmt/tests/samples/attrs.rsx b/packages/autofmt/tests/samples/attrs.rsx
index 9363d802b..947456fbe 100644
--- a/packages/autofmt/tests/samples/attrs.rsx
+++ b/packages/autofmt/tests/samples/attrs.rsx
@@ -18,4 +18,18 @@ rsx! {
         // This here
         "hi"
     }
+
+
+    // Comment head
+    div {
+        class: "asd",
+        "Jon"
+    }
+
+    // Comment head
+    div {
+        // Collapse
+        class: "asd",
+        "Jon"
+    }
 }