fix: trailing whitespace issue when using codeblocks

formatter/src/collect_comments.rs

@@ -2,7 +2,7 @@ use std::collections::HashMap;
 
 use crop::Rope;
 
-use proc_macro2::{Span, TokenStream, TokenTree};
+use proc_macro2::{Span, TokenStream};
 
 use crate::get_text_beween_spans;
 
@@ -13,8 +13,7 @@ pub(crate) fn extract_whitespace_and_comments(
     let mut whitespace_and_comments = HashMap::new();
     let mut last_span: Option<Span> = None;
 
-    traverse_token_stream(tokens, &mut |token: &TokenTree| {
-        let span = token.span();
+    traverse_token_stream(tokens, &mut |span: Span| {
         if let Some(last_span) = last_span {
             if last_span.end().line != span.start().line {
                 let text = get_text_beween_spans(source, last_span.end(), span.start());
@@ -45,11 +44,15 @@ pub(crate) fn extract_whitespace_and_comments(
     whitespace_and_comments
 }
 
-fn traverse_token_stream(tokens: TokenStream, cb: &mut impl FnMut(&TokenTree)) {
+fn traverse_token_stream(tokens: TokenStream, cb: &mut impl FnMut(Span)) {
     for token in tokens {
         match token {
-            proc_macro2::TokenTree::Group(group) => traverse_token_stream(group.stream(), cb),
-            _ => cb(&token),
+            proc_macro2::TokenTree::Group(group) => {
+                cb(group.span_open());
+                traverse_token_stream(group.stream(), cb);
+                cb(group.span_close());
+            }
+            _ => cb(token.span()),
         }
     }
 }

formatter/src/formatter/attribute.rs

@@ -7,7 +7,7 @@ use super::ExpressionFormatter;
 
 impl Formatter<'_> {
     pub fn attribute(&mut self, attribute: &NodeAttribute) {
-        self.flush_comments(attribute.span().start().line - 1);
+        self.flush_comments(attribute.span().start().line - 1, false);
         match attribute {
             NodeAttribute::Attribute(k) => self.keyed_attribute(k),
             NodeAttribute::Block(b) => self.node_block(b),

formatter/src/formatter/element.rs

@@ -11,7 +11,7 @@ impl Formatter<'_> {
 
         if !is_void {
             self.children(&element.children, element.attributes().len());
-            self.flush_comments(element.close_tag.span().end().line - 1);
+            self.flush_comments(element.close_tag.span().end().line - 1, true);
             self.closing_tag(element)
         }
     }

formatter/src/formatter/expr.rs

@@ -109,7 +109,7 @@
 
     fn expr(&mut self, expr: &syn::Expr, formatter: Option<ExpressionFormatter>) {
         let span = expr.span();
-        self.flush_comments(span.start().line - 1);
+        self.flush_comments(span.start().line - 1, false);
         if let syn::Expr::Lit(ExprLit {
             lit: syn::Lit::Str(lit_str),
             ..
@@ -153,19 +153,21 @@
 
 #[cfg(test)]
 mod tests {
     use rstml::node::Node;
 
     use crate::formatter::*;
-    use crate::test_helpers::{element_from_string, format_element_from_string, format_with};
+    use crate::test_helpers::{
+        element, element_from_string, format_element_from_string, format_with,
+    };
 
     macro_rules! format_element {
         ($($tt:tt)*) => {{
             let settings = FormatterSettings {
                 max_width: 40,
                 ..Default::default()
             };
 
-            let element = element_from_string! { $($tt)* };
+            let element = element! { $($tt)* };
             format_with(settings,|formatter| {
                 formatter.node(&Node::Element(element));
             })
@@ -185,7 +187,7 @@
 
     #[test]
     fn multiline_string_as_child() {
-        let formatted = format_element! {r#"<div>
+        let formatted = format_element_from_string! {r#"<div>
                     "Lorem ipsum dolor sit amet, consectetur adipiscing elit,
                         sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
                                 Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
@@ -206,7 +208,7 @@
 
     #[test]
     fn string_whitespace_prefix() {
-        let formatted = format_element! {r#"<div>
+        let formatted = format_element_from_string! {r#"<div>
                     "    foo"
             </div>"#};
 
@@ -217,7 +219,7 @@
 
     #[test]
     fn multiline_string_whitespace_prefix() {
-        let formatted = format_element! {r#"<div>
+        let formatted = format_element_from_string! {r#"<div>
                     "        Lorem ipsum dolor sit amet, consectetur adipiscing elit,
                         sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
                                 Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
@@ -238,7 +240,7 @@
 
     #[test]
     fn multiline_unquoted_string_as_child() {
-        let formatted = format_element! {r#"<div>
+        let formatted = format_element_from_string! {r#"<div>
                     Lorem ipsum dolor sit amet, consectetur adipiscing elit,
                         sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
                                 Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
@@ -273,4 +275,42 @@
         <p>"\u{00A9}🦀"</p>
         "#);
     }
+
+    #[test]
+    fn codeblock_with_empty_lines() {
+        let formatted = format_element_from_string! { r#"
+                    <h2>
+                        {
+
+                        }
+                    </h2>
+            "#
+        };
+
+        insta::assert_snapshot!(formatted, @r###"
+        <h2>{}</h2>
+        "###);
+    }
+
+    #[test]
+    fn codeblock_body() {
+        let formatted = format_element_from_string! { r#"<h2>
+                        {match error_code {
+                            StatusCode::SERVICE_UNAVAILABLE => "custom error msg".to_string(),
+                            error_code => error_code.to_string(),
+                        }}
+                    </h2>"#
+        };
+
+        insta::assert_snapshot!(formatted, @r###"
+        <h2>
+            {match error_code {
+                StatusCode::SERVICE_UNAVAILABLE => {
+                    "custom error msg".to_string()
+                }
+                error_code => error_code.to_string(),
+            }}
+        </h2>
+        "###);
+    }
 }

formatter/src/formatter/mac.rs

@@ -84,7 +84,7 @@ impl Formatter<'_> {
         self.printer
             .cbox((parent_indent.tabs * self.settings.tab_spaces + parent_indent.spaces) as isize);
 
-        self.flush_comments(cx.span().start().line - 1);
+        self.flush_comments(cx.span().start().line - 1, false);
 
         let macro_word = format!("{}! {{", get_macro_full_path(view_mac.mac));
         self.printer.word(macro_word);

formatter/src/formatter/mod.rs

@@ -179,16 +179,29 @@ impl<'a> Formatter<'a> {
         }
     }
 
-    pub fn flush_comments(&mut self, line_index: usize) {
+    pub fn flush_comments(&mut self, line_index: usize, skip_trailing_whitespace: bool) {
         let last = self.line_offset.unwrap_or(0);
 
         let comments_or_empty_lines: Vec<_> = (last..=line_index)
             .filter_map(|l| self.whitespace_and_comments.remove(&l))
             .collect();
 
+        // If we need to skip trailing whitespace, calculate how many elements we need to take,
+        // until no comments are left in the vector
+        let take_n = if skip_trailing_whitespace {
+            comments_or_empty_lines
+                .iter()
+                .rev()
+                .position(Option::is_some)
+                .map(|i| comments_or_empty_lines.len() - i)
+        } else {
+            None
+        }
+        .unwrap_or(comments_or_empty_lines.len());
+
         let mut prev_is_empty_line = false;
 
-        for comment_or_empty in comments_or_empty_lines {
+        for comment_or_empty in comments_or_empty_lines.into_iter().take(take_n) {
             if let Some(comment) = comment_or_empty {
                 self.printer.word("// ");
                 self.printer.word(comment);

formatter/src/formatter/node.rs

@@ -5,7 +5,7 @@ use crate::{formatter::Formatter, get_text_beween_spans};
 
 impl Formatter<'_> {
     pub fn node(&mut self, node: &Node) {
-        self.flush_comments(node.span().start().line - 1);
+        self.flush_comments(node.span().start().line - 1, false);
 
         match node {
             Node::Element(ele) => self.element(ele),

formatter/src/source_file.rs

@@ -380,7 +380,7 @@ mod tests {
         "#};
 
         let result = format_file_source(source, &Default::default()).unwrap();
-        insta::assert_snapshot!(result, @r#"
+        insta::assert_snapshot!(result, @r###"
         // comment outside view macro
         fn main() {
             view! { cx,
@@ -416,7 +416,7 @@ mod tests {
         }
 
         // comment after view macro
-        "#);
+        "###);
     }
 
     #[test]