Merge pull request #726 from DioxusLabs/jk/fix-comments-on-exprs

fix: comments being eaten in autofmt

packages/autofmt/src/buffer.rs

@@ -1,33 +1,18 @@
-use std::{
-    collections::{HashMap, VecDeque},
-    fmt::{Result, Write},
-};
+//! The output buffer that supports some helpful methods
+//! These are separate from the input so we can lend references between the two
+//!
+//!
+//!
 
-use dioxus_rsx::{BodyNode, ElementAttr, ElementAttrNamed, IfmtInput};
-use proc_macro2::{LineColumn, Span};
-use syn::{spanned::Spanned, Expr};
+use std::fmt::{Result, Write};
 
-#[derive(Default, Debug)]
+use dioxus_rsx::IfmtInput;
+
+/// The output buffer that tracks indent and string
+#[derive(Debug, Default)]
 pub struct Buffer {
-    pub src: Vec<String>,
-    pub cached_formats: HashMap<Location, String>,
     pub buf: String,
     pub indent: usize,
-    pub comments: VecDeque<usize>,
-}
-
-#[derive(Clone, Copy, Hash, PartialEq, Eq, Debug)]
-pub struct Location {
-    pub line: usize,
-    pub col: usize,
-}
-impl Location {
-    pub fn new(start: LineColumn) -> Self {
-        Self {
-            line: start.line,
-            col: start.column,
-        }
-    }
 }
 
 impl Buffer {
@@ -62,160 +47,14 @@ impl Buffer {
         writeln!(self.buf)
     }
 
-    // Expects to be written directly into place
-    pub fn write_ident(&mut self, node: &BodyNode) -> Result {
-        match node {
-            BodyNode::Element(el) => self.write_element(el),
-            BodyNode::Component(component) => self.write_component(component),
-            BodyNode::Text(text) => self.write_text(text),
-            BodyNode::RawExpr(exp) => self.write_raw_expr(exp),
-            _ => Ok(()),
-        }
-    }
-
     pub fn write_text(&mut self, text: &IfmtInput) -> Result {
         write!(self.buf, "\"{}\"", text.source.as_ref().unwrap().value())
     }
-
-    pub fn consume(self) -> Option<String> {
-        Some(self.buf)
-    }
-
-    pub fn write_comments(&mut self, child: Span) -> Result {
-        // collect all comments upwards
-        let start = child.start();
-        let line_start = start.line - 1;
-
-        for (id, line) in self.src[..line_start].iter().enumerate().rev() {
-            if line.trim().starts_with("//") || line.is_empty() {
-                if id != 0 {
-                    self.comments.push_front(id);
-                }
-            } else {
-                break;
-            }
-        }
-
-        let mut last_was_empty = false;
-        while let Some(comment_line) = self.comments.pop_front() {
-            let line = &self.src[comment_line];
-            if line.is_empty() {
-                if !last_was_empty {
-                    self.new_line()?;
-                }
-                last_was_empty = true;
-            } else {
-                last_was_empty = false;
-                self.tabbed_line()?;
-                write!(self.buf, "{}", self.src[comment_line].trim())?;
-            }
-        }
-
-        Ok(())
-    }
-
-    // Push out the indent level and write each component, line by line
-    pub fn write_body_indented(&mut self, children: &[BodyNode]) -> Result {
-        self.indent += 1;
-
-        self.write_body_no_indent(children)?;
-
-        self.indent -= 1;
-        Ok(())
-    }
-
-    pub fn write_body_no_indent(&mut self, children: &[BodyNode]) -> Result {
-        let last_child = children.len();
-
-        for (idx, child) in children.iter().enumerate() {
-            match child {
-                // check if the expr is a short
-                BodyNode::RawExpr { .. } => {
-                    self.tabbed_line()?;
-                    self.write_ident(child)?;
-                    if idx != last_child - 1 {
-                        write!(self.buf, ",")?;
-                    }
-                }
-                _ => {
-                    if self.current_span_is_primary(child.span()) {
-                        self.write_comments(child.span())?;
-                    }
-                    self.tabbed_line()?;
-                    self.write_ident(child)?;
-                }
-            }
-        }
-
-        Ok(())
-    }
-
-    pub(crate) fn is_short_attrs(&mut self, attributes: &[ElementAttrNamed]) -> usize {
-        let mut total = 0;
-
-        for attr in attributes {
-            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 {
-                ElementAttr::AttrText { value, name } => {
-                    value.source.as_ref().unwrap().value().len() + name.span().line_length() + 3
-                }
-                ElementAttr::AttrExpression { name, value } => {
-                    value.span().line_length() + name.span().line_length() + 3
-                }
-                ElementAttr::CustomAttrText { value, name } => {
-                    value.source.as_ref().unwrap().value().len() + name.value().len() + 3
-                }
-                ElementAttr::CustomAttrExpression { name, value } => {
-                    name.value().len() + value.span().line_length() + 3
-                }
-                ElementAttr::EventTokens { tokens, name } => {
-                    let location = Location::new(tokens.span().start());
-
-                    let len = if let std::collections::hash_map::Entry::Vacant(e) =
-                        self.cached_formats.entry(location)
-                    {
-                        let formatted = prettyplease::unparse_expr(tokens);
-                        let len = if formatted.contains('\n') {
-                            10000
-                        } else {
-                            formatted.len()
-                        };
-                        e.insert(formatted);
-                        len
-                    } else {
-                        self.cached_formats[&location].len()
-                    };
-
-                    len + name.span().line_length() + 3
-                }
-            };
-        }
-
-        total
-    }
-
-    pub fn retrieve_formatted_expr(&mut self, expr: &Expr) -> &str {
-        self.cached_formats
-            .entry(Location::new(expr.span().start()))
-            .or_insert_with(|| prettyplease::unparse_expr(expr))
-            .as_str()
-    }
 }
 
-trait SpanLength {
-    fn line_length(&self) -> usize;
-}
-impl SpanLength for Span {
-    fn line_length(&self) -> usize {
-        self.end().line - self.start().line
+impl std::fmt::Write for Buffer {
+    fn write_str(&mut self, s: &str) -> std::fmt::Result {
+        self.buf.push_str(s);
+        Ok(())
     }
 }

packages/autofmt/src/component.rs

@@ -1,4 +1,4 @@
-use crate::{buffer::Location, Buffer};
+use crate::{writer::Location, Writer};
 use dioxus_rsx::*;
 use quote::ToTokens;
 use std::fmt::{Result, Write};
@@ -19,7 +19,7 @@ enum ShortOptimization {
     NoOpt,
 }
 
-impl Buffer {
+impl Writer {
     pub fn write_component(
         &mut self,
         Component {
@@ -28,6 +28,7 @@ impl Buffer {
             children,
             manual_props,
             prop_gen_args,
+            ..
         }: &Component,
     ) -> Result {
         self.write_component_name(name, prop_gen_args)?;
@@ -82,46 +83,46 @@ impl Buffer {
         match opt_level {
             ShortOptimization::Empty => {}
             ShortOptimization::Oneliner => {
-                write!(self.buf, " ")?;
+                write!(self.out, " ")?;
 
                 self.write_component_fields(fields, manual_props, true)?;
 
                 if !children.is_empty() && !fields.is_empty() {
-                    write!(self.buf, ", ")?;
+                    write!(self.out, ", ")?;
                 }
 
                 for child in children {
                     self.write_ident(child)?;
                 }
 
-                write!(self.buf, " ")?;
+                write!(self.out, " ")?;
             }
 
             ShortOptimization::PropsOnTop => {
-                write!(self.buf, " ")?;
+                write!(self.out, " ")?;
                 self.write_component_fields(fields, manual_props, true)?;
 
                 if !children.is_empty() && !fields.is_empty() {
-                    write!(self.buf, ",")?;
+                    write!(self.out, ",")?;
                 }
 
                 self.write_body_indented(children)?;
-                self.tabbed_line()?;
+                self.out.tabbed_line()?;
             }
 
             ShortOptimization::NoOpt => {
                 self.write_component_fields(fields, manual_props, false)?;
 
                 if !children.is_empty() && !fields.is_empty() {
-                    write!(self.buf, ",")?;
+                    write!(self.out, ",")?;
                 }
 
                 self.write_body_indented(children)?;
-                self.tabbed_line()?;
+                self.out.tabbed_line()?;
             }
         }
 
-        write!(self.buf, "}}")?;
+        write!(self.out, "}}")?;
         Ok(())
     }
 
@@ -133,16 +134,16 @@ impl Buffer {
         let mut name = name.to_token_stream().to_string();
         name.retain(|c| !c.is_whitespace());
 
-        write!(self.buf, "{name}")?;
+        write!(self.out, "{name}")?;
 
         if let Some(generics) = generics {
             let mut written = generics.to_token_stream().to_string();
             written.retain(|c| !c.is_whitespace());
 
-            write!(self.buf, "{}", written)?;
+            write!(self.out, "{}", written)?;
         }
 
-        write!(self.buf, " {{")?;
+        write!(self.out, " {{")?;
 
         Ok(())
     }
@@ -157,18 +158,18 @@ impl Buffer {
 
         while let Some(field) = field_iter.next() {
             if !sameline {
-                self.indented_tabbed_line()?;
+                self.out.indented_tabbed_line()?;
             }
 
             let name = &field.name;
             match &field.content {
                 ContentField::ManExpr(exp) => {
                     let out = prettyplease::unparse_expr(exp);
-                    write!(self.buf, "{}: {}", name, out)?;
+                    write!(self.out, "{}: {}", name, out)?;
                 }
                 ContentField::Formatted(s) => {
                     write!(
-                        self.buf,
+                        self.out,
                         "{}: \"{}\"",
                         name,
                         s.source.as_ref().unwrap().value()
@@ -178,27 +179,27 @@ impl Buffer {
                     let out = prettyplease::unparse_expr(exp);
                     let mut lines = out.split('\n').peekable();
                     let first = lines.next().unwrap();
-                    write!(self.buf, "{}: {}", name, first)?;
+                    write!(self.out, "{}: {}", name, first)?;
                     for line in lines {
-                        self.new_line()?;
-                        self.indented_tab()?;
-                        write!(self.buf, "{}", line)?;
+                        self.out.new_line()?;
+                        self.out.indented_tab()?;
+                        write!(self.out, "{}", line)?;
                     }
                 }
             }
 
             if field_iter.peek().is_some() || manual_props.is_some() {
-                write!(self.buf, ",")?;
+                write!(self.out, ",")?;
 
                 if sameline {
-                    write!(self.buf, " ")?;
+                    write!(self.out, " ")?;
                 }
             }
         }
 
         if let Some(exp) = manual_props {
             if !sameline {
-                self.indented_tabbed_line()?;
+                self.out.indented_tabbed_line()?;
             }
             self.write_manual_props(exp)?;
         }
@@ -258,10 +259,10 @@ impl Buffer {
 
         let first_line = lines.next().unwrap();
 
-        write!(self.buf, "..{first_line}")?;
+        write!(self.out, "..{first_line}")?;
         for line in lines {
-            self.indented_tabbed_line()?;
-            write!(self.buf, "{line}")?;
+            self.out.indented_tabbed_line()?;
+            write!(self.out, "{line}")?;
         }
 
         Ok(())

packages/autofmt/src/element.rs

@@ -1,35 +1,56 @@
-use crate::Buffer;
+use crate::Writer;
 use dioxus_rsx::*;
 use proc_macro2::Span;
-use std::{fmt::Result, fmt::Write};
-use syn::{spanned::Spanned, Expr};
+use std::{
+    fmt::Result,
+    fmt::{self, Write},
+};
+use syn::{spanned::Spanned, token::Brace, Expr};
 
 #[derive(Debug)]
 enum ShortOptimization {
-    // Special because we want to print the closing bracket immediately
+    /// Special because we want to print the closing bracket immediately
+    ///
+    /// IE
+    /// `div {}` instead of `div { }`
     Empty,
 
-    // Special optimization to put everything on the same line
+    /// Special optimization to put everything on the same line and add some buffer spaces
+    ///
+    /// IE
+    ///
+    /// `div { "asdasd" }` instead of a multiline variant
     Oneliner,
 
-    // Optimization where children flow but props remain fixed on top
+    /// Optimization where children flow but props remain fixed on top
     PropsOnTop,
 
-    // The noisiest optimization where everything flows
+    /// The noisiest optimization where everything flows
     NoOpt,
 }
 
-impl Buffer {
-    pub fn write_element(
-        &mut self,
-        Element {
+/*
+// whitespace
+div {
+    // some whitespace
+    class: "asdasd"
+
+    // whjiot
+    asdasd // whitespace
+}
+*/
+
+impl Writer {
+    pub fn write_element(&mut self, el: &Element) -> Result {
+        let Element {
             name,
             key,
             attributes,
             children,
             _is_static,
-        }: &Element,
-    ) -> Result {
+            brace,
+        } = el;
+
         /*
             1. Write the tag
             2. Write the key
@@ -37,7 +58,7 @@ impl Buffer {
             4. Write the children
         */
 
-        write!(self.buf, "{name} {{")?;
+        write!(self.out, "{name} {{")?;
 
         // decide if we have any special optimizations
         // Default with none, opt the cases in one-by-one
@@ -70,6 +91,9 @@ impl Buffer {
         // If there's nothing at all, empty optimization
         if attributes.is_empty() && children.is_empty() && key.is_none() {
             opt_level = ShortOptimization::Empty;
+
+            // Write comments if they exist
+            self.write_todo_body(brace)?;
         }
 
         // multiline handlers bump everything down
@@ -80,55 +104,56 @@ impl Buffer {
         match opt_level {
             ShortOptimization::Empty => {}
             ShortOptimization::Oneliner => {
-                write!(self.buf, " ")?;
+                write!(self.out, " ")?;
 
                 self.write_attributes(attributes, key, true)?;
 
                 if !children.is_empty() && (!attributes.is_empty() || key.is_some()) {
-                    write!(self.buf, ", ")?;
+                    write!(self.out, ", ")?;
                 }
 
                 for (id, child) in children.iter().enumerate() {
                     self.write_ident(child)?;
                     if id != children.len() - 1 && children.len() > 1 {
-                        write!(self.buf, ", ")?;
+                        write!(self.out, ", ")?;
                     }
                 }
 
-                write!(self.buf, " ")?;
+                write!(self.out, " ")?;
             }
 
             ShortOptimization::PropsOnTop => {
                 if !attributes.is_empty() || key.is_some() {
-                    write!(self.buf, " ")?;
+                    write!(self.out, " ")?;
                 }
                 self.write_attributes(attributes, key, true)?;
 
                 if !children.is_empty() && (!attributes.is_empty() || key.is_some()) {
-                    write!(self.buf, ",")?;
+                    write!(self.out, ",")?;
                 }
 
                 if !children.is_empty() {
                     self.write_body_indented(children)?;
                 }
-                self.tabbed_line()?;
+                self.out.tabbed_line()?;
             }
 
             ShortOptimization::NoOpt => {
                 self.write_attributes(attributes, key, false)?;
 
                 if !children.is_empty() && (!attributes.is_empty() || key.is_some()) {
-                    write!(self.buf, ",")?;
+                    write!(self.out, ",")?;
                 }
 
                 if !children.is_empty() {
                     self.write_body_indented(children)?;
                 }
-                self.tabbed_line()?;
+
+                self.out.tabbed_line()?;
             }
         }
 
-        write!(self.buf, "}}")?;
+        write!(self.out, "}}")?;
 
         Ok(())
     }
@@ -143,39 +168,39 @@ impl Buffer {
 
         if let Some(key) = key {
             if !sameline {
-                self.indented_tabbed_line()?;
+                self.out.indented_tabbed_line()?;
             }
             write!(
-                self.buf,
+                self.out,
                 "key: \"{}\"",
                 key.source.as_ref().unwrap().value()
             )?;
             if !attributes.is_empty() {
-                write!(self.buf, ",")?;
+                write!(self.out, ",")?;
                 if sameline {
-                    write!(self.buf, " ")?;
+                    write!(self.out, " ")?;
                 }
             }
         }
 
         while let Some(attr) = attr_iter.next() {
-            self.indent += 1;
+            self.out.indent += 1;
             if !sameline {
                 self.write_comments(attr.attr.start())?;
             }
-            self.indent -= 1;
+            self.out.indent -= 1;
 
             if !sameline {
-                self.indented_tabbed_line()?;
+                self.out.indented_tabbed_line()?;
             }
 
             self.write_attribute(attr)?;
 
             if attr_iter.peek().is_some() {
-                write!(self.buf, ",")?;
+                write!(self.out, ",")?;
 
                 if sameline {
-                    write!(self.buf, " ")?;
+                    write!(self.out, " ")?;
                 }
             }
         }
@@ -187,19 +212,19 @@ impl Buffer {
         match &attr.attr {
             ElementAttr::AttrText { name, value } => {
                 write!(
-                    self.buf,
+                    self.out,
                     "{name}: \"{value}\"",
                     value = value.source.as_ref().unwrap().value()
                 )?;
             }
             ElementAttr::AttrExpression { name, value } => {
                 let out = prettyplease::unparse_expr(value);
-                write!(self.buf, "{}: {}", name, out)?;
+                write!(self.out, "{}: {}", name, out)?;
             }
 
             ElementAttr::CustomAttrText { name, value } => {
                 write!(
-                    self.buf,
+                    self.out,
                     "\"{name}\": \"{value}\"",
                     name = name.value(),
                     value = value.source.as_ref().unwrap().value()
@@ -208,7 +233,7 @@ impl Buffer {
 
             ElementAttr::CustomAttrExpression { name, value } => {
                 let out = prettyplease::unparse_expr(value);
-                write!(self.buf, "\"{}\": {}", name.value(), out)?;
+                write!(self.out, "\"{}\": {}", name.value(), out)?;
             }
 
             ElementAttr::EventTokens { name, tokens } => {
@@ -220,17 +245,17 @@ impl Buffer {
                 // a one-liner for whatever reason
                 // Does not need a new line
                 if lines.peek().is_none() {
-                    write!(self.buf, "{}: {}", name, first)?;
+                    write!(self.out, "{}: {}", name, first)?;
                 } else {
-                    writeln!(self.buf, "{}: {}", name, first)?;
+                    writeln!(self.out, "{}: {}", name, first)?;
 
                     while let Some(line) = lines.next() {
-                        self.indented_tab()?;
-                        write!(self.buf, "{}", line)?;
+                        self.out.indented_tab()?;
+                        write!(self.out, "{}", line)?;
                         if lines.peek().is_none() {
-                            write!(self.buf, "")?;
+                            write!(self.out, "")?;
                         } else {
-                            writeln!(self.buf)?;
+                            writeln!(self.out)?;
                         }
                     }
                 }
@@ -254,8 +279,6 @@ impl Buffer {
             ""
         };
 
-        // dbg!(beginning);
-
         beginning.is_empty()
     }
 
@@ -268,6 +291,10 @@ impl Buffer {
         if children.is_empty() {
             // todo: allow elements with comments but no children
             // like div { /* comment */ }
+            // or
+            // div {
+            //  // some helpful
+            // }
             return Some(0);
         }
 
@@ -342,6 +369,35 @@ impl Buffer {
             }
         }
     }
+
+    /// empty everything except for some comments
+    fn write_todo_body(&mut self, brace: &Brace) -> fmt::Result {
+        let span = brace.span.span();
+        let start = span.start();
+        let end = span.end();
+
+        if start.line == end.line {
+            return Ok(());
+        }
+
+        writeln!(self.out)?;
+
+        for idx in start.line..end.line {
+            let line = &self.src[idx];
+            if line.trim().starts_with("//") {
+                for _ in 0..self.out.indent + 1 {
+                    write!(self.out, "    ")?
+                }
+                writeln!(self.out, "{}", line.trim()).unwrap();
+            }
+        }
+
+        for _ in 0..self.out.indent {
+            write!(self.out, "    ")?
+        }
+
+        Ok(())
+    }
 }
 
 fn get_expr_length(expr: &Expr) -> Option<usize> {

packages/autofmt/src/expr.rs

@@ -1,45 +1,23 @@
 //! pretty printer for rsx!
 use std::fmt::{Result, Write};
 
-use crate::Buffer;
+use crate::Writer;
 
-impl Buffer {
+impl Writer {
     pub fn write_raw_expr(&mut self, exp: &syn::Expr) -> Result {
         /*
         We want to normalize the expr to the appropriate indent level.
         */
 
-        // in a perfect world, just fire up the rust pretty printer
-        // pretty_print_rust_code_as_if_it_were_rustfmt()
-
         use syn::spanned::Spanned;
         let placement = exp.span();
         let start = placement.start();
         let end = placement.end();
-        // let num_spaces_desired = (self.indent * 4) as isize;
-
-        // print comments
-        // let mut queued_comments = vec![];
-        // let mut offset = 2;
-        // loop {
-        //     let line = &self.src[start.line - offset];
-        //     if line.trim_start().starts_with("//") {
-        //         queued_comments.push(line);
-        //     } else {
-        //         break;
-        //     }
-
-        //     offset += 1;
-        // }
-        // let had_comments = !queued_comments.is_empty();
-        // for comment in queued_comments.into_iter().rev() {
-        //     writeln!(self.buf, "{}", comment)?;
-        // }
 
         // if the expr is on one line, just write it directly
         if start.line == end.line {
             write!(
-                self.buf,
+                self.out,
                 "{}",
                 &self.src[start.line - 1][start.column - 1..end.column].trim()
             )?;
@@ -50,7 +28,7 @@ impl Buffer {
         // This involves unshifting the first line if it's aligned
         let first_line = &self.src[start.line - 1];
         write!(
-            self.buf,
+            self.out,
             "{}",
             &first_line[start.column - 1..first_line.len()].trim()
         )?;
@@ -66,7 +44,7 @@ impl Buffer {
         };
 
         for (id, line) in self.src[start.line..end.line].iter().enumerate() {
-            writeln!(self.buf)?;
+            writeln!(self.out)?;
             // trim the leading whitespace
             let line = match id {
                 x if x == (end.line - start.line) - 1 => &line[..end.column],
@@ -75,52 +53,17 @@ impl Buffer {
 
             if offset < 0 {
                 for _ in 0..-offset {
-                    write!(self.buf, " ")?;
+                    write!(self.out, " ")?;
                 }
 
-                write!(self.buf, "{}", line)?;
+                write!(self.out, "{}", line)?;
             } else {
                 let offset = offset as usize;
                 let right = &line[offset..];
-                write!(self.buf, "{}", right)?;
+                write!(self.out, "{}", right)?;
             }
         }
 
-        // let first = &self.src[start.line - 1];
-        // let num_spaces_real = first.chars().take_while(|c| c.is_whitespace()).count() as isize;
-        // let offset = num_spaces_real - num_spaces_desired;
-
-        // for (row, line) in self.src[start.line - 1..end.line].iter().enumerate() {
-        //     let line = match row {
-        //         0 => &line[start.column - 1..],
-        //         a if a == (end.line - start.line) => &line[..end.column - 1],
-        //         _ => line,
-        //     };
-
-        //     writeln!(self.buf)?;
-        //     // trim the leading whitespace
-        //     if offset < 0 {
-        //         for _ in 0..-offset {
-        //             write!(self.buf, " ")?;
-        //         }
-
-        //         write!(self.buf, "{}", line)?;
-        //     } else {
-        //         let offset = offset as usize;
-        //         let right = &line[offset..];
-        //         write!(self.buf, "{}", right)?;
-        //     }
-        // }
-
         Ok(())
     }
 }
-
-// :(
-// fn pretty_print_rust_code_as_if_it_were_rustfmt(code: &str) -> String {
-//     let formatted = prettyplease::unparse_expr(exp);
-//     for line in formatted.lines() {
-//         write!(self.buf, "{}", line)?;
-//         self.new_line()?;
-//     }
-// }

packages/autofmt/src/lib.rs

@@ -1,13 +1,14 @@
 use dioxus_rsx::CallBody;
 
-use crate::buffer::*;
 use crate::util::*;
+use crate::writer::*;
 
 mod buffer;
 mod component;
 mod element;
 mod expr;
 mod util;
+mod writer;
 
 /// A modification to the original file to be applied by an IDE
 ///
@@ -101,10 +102,9 @@ pub fn fmt_file(contents: &str) -> Vec<FormattedBlock> {
 }
 
 pub fn write_block_out(body: CallBody) -> Option<String> {
-    let mut buf = Buffer {
+    let mut buf = Writer {
         src: vec!["".to_string()],
-        indent: 0,
-        ..Buffer::default()
+        ..Writer::default()
     };
 
     // Oneliner optimization
@@ -120,12 +120,13 @@ pub fn write_block_out(body: CallBody) -> Option<String> {
 pub fn fmt_block(block: &str, indent_level: usize) -> Option<String> {
     let body = syn::parse_str::<dioxus_rsx::CallBody>(block).ok()?;
 
-    let mut buf = Buffer {
+    let mut buf = Writer {
         src: block.lines().map(|f| f.to_string()).collect(),
-        indent: indent_level,
-        ..Buffer::default()
+        ..Writer::default()
     };
 
+    buf.out.indent = indent_level;
+
     // Oneliner optimization
     if buf.is_short_children(&body.roots).is_some() {
         buf.write_ident(&body.roots[0]).unwrap();
@@ -134,8 +135,8 @@ pub fn fmt_block(block: &str, indent_level: usize) -> Option<String> {
     }
 
     // writing idents leaves the final line ended at the end of the last ident
-    if buf.buf.contains('\n') {
-        buf.new_line().unwrap();
+    if buf.out.buf.contains('\n') {
+        buf.out.new_line().unwrap();
     }
 
     buf.consume()

packages/autofmt/src/writer.rs

@@ -0,0 +1,188 @@
+use dioxus_rsx::{BodyNode, ElementAttr, ElementAttrNamed};
+use proc_macro2::{LineColumn, Span};
+use std::{
+    collections::{HashMap, VecDeque},
+    fmt::{Result, Write},
+};
+use syn::{spanned::Spanned, Expr};
+
+use crate::buffer::Buffer;
+
+#[derive(Debug, Default)]
+pub struct Writer {
+    pub src: Vec<String>,
+    pub cached_formats: HashMap<Location, String>,
+    pub comments: VecDeque<usize>,
+    pub out: Buffer,
+}
+
+#[derive(Clone, Copy, Hash, PartialEq, Eq, Debug)]
+pub struct Location {
+    pub line: usize,
+    pub col: usize,
+}
+impl Location {
+    pub fn new(start: LineColumn) -> Self {
+        Self {
+            line: start.line,
+            col: start.column,
+        }
+    }
+}
+
+impl Writer {
+    // Expects to be written directly into place
+    pub fn write_ident(&mut self, node: &BodyNode) -> Result {
+        match node {
+            BodyNode::Element(el) => self.write_element(el),
+            BodyNode::Component(component) => self.write_component(component),
+            BodyNode::Text(text) => self.out.write_text(text),
+            BodyNode::RawExpr(exp) => self.write_raw_expr(exp),
+            _ => Ok(()),
+        }
+    }
+
+    pub fn consume(self) -> Option<String> {
+        Some(self.out.buf)
+    }
+
+    pub fn write_comments(&mut self, child: Span) -> Result {
+        // collect all comments upwards
+        let start = child.start();
+        let line_start = start.line - 1;
+
+        for (id, line) in self.src[..line_start].iter().enumerate().rev() {
+            if line.trim().starts_with("//") || line.is_empty() {
+                if id != 0 {
+                    self.comments.push_front(id);
+                }
+            } else {
+                break;
+            }
+        }
+
+        let mut last_was_empty = false;
+        while let Some(comment_line) = self.comments.pop_front() {
+            let line = &self.src[comment_line];
+            if line.is_empty() {
+                if !last_was_empty {
+                    self.out.new_line()?;
+                }
+                last_was_empty = true;
+            } else {
+                last_was_empty = false;
+                self.out.tabbed_line()?;
+                write!(self.out, "{}", self.src[comment_line].trim())?;
+            }
+        }
+
+        Ok(())
+    }
+
+    // Push out the indent level and write each component, line by line
+    pub fn write_body_indented(&mut self, children: &[BodyNode]) -> Result {
+        self.out.indent += 1;
+
+        self.write_body_no_indent(children)?;
+
+        self.out.indent -= 1;
+        Ok(())
+    }
+
+    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())?;
+            }
+
+            match child {
+                // check if the expr is a short
+                BodyNode::RawExpr { .. } => {
+                    self.out.tabbed_line()?;
+                    self.write_ident(child)?;
+                    if idx != last_child - 1 {
+                        write!(self.out, ",")?;
+                    }
+                }
+                _ => {
+                    self.out.tabbed_line()?;
+                    self.write_ident(child)?;
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    pub(crate) fn is_short_attrs(&mut self, attributes: &[ElementAttrNamed]) -> usize {
+        let mut total = 0;
+
+        for attr in attributes {
+            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 {
+                ElementAttr::AttrText { value, name } => {
+                    value.source.as_ref().unwrap().value().len() + name.span().line_length() + 3
+                }
+                ElementAttr::AttrExpression { name, value } => {
+                    value.span().line_length() + name.span().line_length() + 3
+                }
+                ElementAttr::CustomAttrText { value, name } => {
+                    value.source.as_ref().unwrap().value().len() + name.value().len() + 3
+                }
+                ElementAttr::CustomAttrExpression { name, value } => {
+                    name.value().len() + value.span().line_length() + 3
+                }
+                ElementAttr::EventTokens { tokens, name } => {
+                    let location = Location::new(tokens.span().start());
+
+                    let len = if let std::collections::hash_map::Entry::Vacant(e) =
+                        self.cached_formats.entry(location)
+                    {
+                        let formatted = prettyplease::unparse_expr(tokens);
+                        let len = if formatted.contains('\n') {
+                            10000
+                        } else {
+                            formatted.len()
+                        };
+                        e.insert(formatted);
+                        len
+                    } else {
+                        self.cached_formats[&location].len()
+                    };
+
+                    len + name.span().line_length() + 3
+                }
+            };
+        }
+
+        total
+    }
+
+    pub fn retrieve_formatted_expr(&mut self, expr: &Expr) -> &str {
+        self.cached_formats
+            .entry(Location::new(expr.span().start()))
+            .or_insert_with(|| prettyplease::unparse_expr(expr))
+            .as_str()
+    }
+}
+
+trait SpanLength {
+    fn line_length(&self) -> usize;
+}
+impl SpanLength for Span {
+    fn line_length(&self) -> usize {
+        self.end().line - self.start().line
+    }
+}

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);

packages/autofmt/tests/samples/commentshard.rsx

@@ -0,0 +1,42 @@
+rsx! {
+    // Comments
+    div {
+        // Comments
+        class: "asdasd",
+
+        // Comments
+        "hello world"
+
+        // Comments
+        expr1,
+
+        // Comments
+        expr2,
+
+        // Comments
+        // Comments
+        // Comments
+        // Comments
+        // Comments
+        expr3,
+
+        div {
+            // todo some work in here
+        }
+
+        div {
+            // todo some work in here
+            // todo some work in here
+            //
+            // todo some work in here
+        }
+
+        div {
+            // todo some work in here
+            class: "hello world",
+
+            // todo some work in here
+            class: "hello world"
+        }
+    }
+}

packages/rsx/src/component.rs

@@ -19,8 +19,7 @@ use syn::{
     ext::IdentExt,
     parse::{Parse, ParseBuffer, ParseStream},
     spanned::Spanned,
-    token, AngleBracketedGenericArguments, Error, Expr, Ident, LitStr, PathArguments, Result,
-    Token,
+    AngleBracketedGenericArguments, Error, Expr, Ident, LitStr, PathArguments, Result, Token,
 };
 
 #[derive(PartialEq, Eq, Clone, Debug, Hash)]
@@ -30,6 +29,7 @@ pub struct Component {
     pub fields: Vec<ComponentField>,
     pub children: Vec<BodyNode>,
     pub manual_props: Option<Expr>,
+    pub brace: syn::token::Brace,
 }
 
 impl Component {
@@ -88,13 +88,9 @@ impl Parse for Component {
 
         // if we see a `{` then we have a block
         // else parse as a function-like call
-        if stream.peek(token::Brace) {
-            syn::braced!(content in stream);
-        } else {
-            syn::parenthesized!(content in stream);
-        }
+        let brace = syn::braced!(content in stream);
 
-        let mut body = Vec::new();
+        let mut fields = Vec::new();
         let mut children = Vec::new();
         let mut manual_props = None;
 
@@ -104,7 +100,7 @@ impl Parse for Component {
                 content.parse::<Token![..]>()?;
                 manual_props = Some(content.parse::<Expr>()?);
             } else if content.peek(Ident) && content.peek2(Token![:]) && !content.peek3(Token![:]) {
-                body.push(content.parse::<ComponentField>()?);
+                fields.push(content.parse::<ComponentField>()?);
             } else {
                 children.push(content.parse::<BodyNode>()?);
             }
@@ -117,9 +113,10 @@ impl Parse for Component {
         Ok(Self {
             name,
             prop_gen_args,
-            fields: body,
+            fields,
             children,
             manual_props,
+            brace,
         })
     }
 }

packages/rsx/src/element.rs

@@ -17,6 +17,7 @@ pub struct Element {
     pub attributes: Vec<ElementAttrNamed>,
     pub children: Vec<BodyNode>,
     pub _is_static: bool,
+    pub brace: syn::token::Brace,
 }
 
 impl Parse for Element {
@@ -25,7 +26,7 @@ impl Parse for Element {
 
         // parse the guts
         let content: ParseBuffer;
-        syn::braced!(content in stream);
+        let brace = syn::braced!(content in stream);
 
         let mut attributes: Vec<ElementAttrNamed> = vec![];
         let mut children: Vec<BodyNode> = vec![];
@@ -152,6 +153,7 @@ impl Parse for Element {
             name: el_name,
             attributes,
             children,
+            brace,
             _is_static: false,
         })
     }