feat(fmt): shiny new comment-preserving formatter! (#38)

Fixes: #34

Author: zkat

src/document.rs

@@ -185,15 +185,10 @@ impl KdlDocument {
         self.trailing = None;
     }
 
-    /// Auto-formats this Document.
-    ///
-    /// Note: This currently removes comments as well.
+    /// Auto-formats this Document, making everything nice while preserving
+    /// comments.
     pub fn fmt(&mut self) {
-        self.leading = None;
-        self.trailing = None;
-        for node in &mut self.nodes {
-            node.fmt();
-        }
+        self.fmt_impl(0);
     }
 }
 
@@ -204,6 +199,18 @@ impl Display for KdlDocument {
 }
 
 impl KdlDocument {
+    pub(crate) fn fmt_impl(&mut self, indent: usize) {
+        if let Some(s) = self.leading.as_mut() {
+            crate::fmt::fmt_leading(s, indent);
+        }
+        if let Some(s) = self.trailing.as_mut() {
+            crate::fmt::fmt_trailing(s);
+        }
+        for node in &mut self.nodes {
+            node.fmt_impl(indent);
+        }
+    }
+
     pub(crate) fn stringify(
         &self,
         f: &mut std::fmt::Formatter<'_>,
@@ -361,6 +368,67 @@ baz
         );
     }
 
+    #[test]
+    fn fmt() -> miette::Result<()> {
+        let mut doc: KdlDocument = r#"
+
+        /* x */ foo    1 "bar"=0xDEADbeef {
+    child1     1  ;
+
+ // child 2 comment
+
+        child2 2 // comment
+
+               child3    "
+
+   string\t" \
+{
+       /*
+
+
+       multiline*/
+                                    inner1    \
+                    r"value" \
+                    ;
+
+        inner2      \ //comment
+        {
+            inner3
+        }
+    }
+               }
+
+        // trailing comment here
+
+        "#
+        .parse()?;
+
+        KdlDocument::fmt(&mut doc);
+
+        print!("{}", doc);
+        assert_eq!(
+            doc.to_string(),
+            r#"/* x */
+foo 1 bar=0xdeadbeef {
+    child1 1
+    // child 2 comment
+    child2 2 // comment
+    child3 "\n\n   string\t" {
+        /*
+
+
+       multiline*/
+        inner1 r"value"
+        inner2 {
+            inner3
+        }
+    }
+}
+// trailing comment here"#
+        );
+        Ok(())
+    }
+
     #[test]
     fn parse_examples() -> miette::Result<()> {
         include_str!("../examples/kdl-schema.kdl").parse::<KdlDocument>()?;

src/fmt.rs

@@ -0,0 +1,30 @@
+pub(crate) fn fmt_leading(leading: &mut String, indent: usize) {
+    if leading.is_empty() {
+        return;
+    }
+    let comments = crate::parser::parse(leading.trim(), crate::parser::leading_comments)
+        .expect("invalid leading text");
+    let mut result = String::new();
+    for line in comments {
+        let trimmed = line.trim();
+        if !trimmed.is_empty() {
+            result.push_str(&format!("{:indent$}{}\n", "", trimmed, indent = indent));
+        }
+    }
+    result.push_str(&format!("{:indent$}", "", indent = indent));
+    *leading = result;
+}
+
+pub(crate) fn fmt_trailing(decor: &mut String) {
+    if decor.is_empty() {
+        return;
+    }
+    *decor = decor.trim().to_string();
+    let mut result = String::new();
+    let comments = crate::parser::parse(decor, crate::parser::trailing_comments)
+        .expect("invalid trailing text");
+    for comment in comments {
+        result.push_str(comment);
+    }
+    *decor = result;
+}

src/lib.rs

@@ -114,6 +114,7 @@ pub use value::*;
 mod document;
 mod entry;
 mod error;
+mod fmt;
 mod identifier;
 mod node;
 mod nom_compat;

src/node.rs

@@ -340,14 +340,7 @@ impl KdlNode {
 
     /// Auto-formats this node and its contents.
     pub fn fmt(&mut self) {
-        self.leading = None;
-        self.trailing = None;
-        for entry in &mut self.entries {
-            entry.fmt();
-        }
-        if let Some(children) = &mut self.children {
-            children.fmt();
-        }
+        self.fmt_impl(0);
     }
 }
 
@@ -428,6 +421,41 @@ impl Display for KdlNode {
 }
 
 impl KdlNode {
+    pub(crate) fn fmt_impl(&mut self, indent: usize) {
+        if let Some(s) = self.leading.as_mut() {
+            crate::fmt::fmt_leading(s, indent);
+        }
+        if let Some(s) = self.trailing.as_mut() {
+            crate::fmt::fmt_trailing(s);
+            if s.starts_with(';') {
+                s.remove(0);
+            }
+            if let Some(c) = s.chars().next() {
+                if !c.is_whitespace() {
+                    s.insert(0, ' ');
+                }
+            }
+            s.push('\n');
+        }
+        self.before_children = None;
+        self.name.clear_fmt();
+        if let Some(ty) = self.ty.as_mut() {
+            ty.clear_fmt()
+        }
+        for entry in &mut self.entries {
+            entry.fmt();
+        }
+        if let Some(children) = self.children.as_mut() {
+            children.fmt_impl(indent + 4);
+            if let Some(leading) = children.leading.as_mut() {
+                leading.push('\n');
+            }
+            if let Some(trailing) = children.trailing.as_mut() {
+                trailing.push_str(format!("{:indent$}", "", indent = indent).as_str());
+            }
+        }
+    }
+
     pub(crate) fn stringify(
         &self,
         f: &mut std::fmt::Formatter<'_>,

src/parser.rs

@@ -112,6 +112,33 @@ pub(crate) fn identifier(input: &str) -> IResult<&str, KdlIdentifier, KdlParseEr
     alt((plain_identifier, quoted_identifier))(input)
 }
 
+pub(crate) fn leading_comments(input: &str) -> IResult<&str, Vec<&str>, KdlParseError<&str>> {
+    terminated(
+        many0(preceded(opt(many0(alt((newline, unicode_space)))), comment)),
+        opt(many0(alt((newline, unicode_space, eof)))),
+    )(input)
+}
+
+pub(crate) fn trailing_comments(mut input: &str) -> IResult<&str, Vec<&str>, KdlParseError<&str>> {
+    let mut comments = vec![];
+    loop {
+        let (inp, _) = opt(many0(alt((newline, unicode_space, tag("\\")))))(input)?;
+        let (inp, comment) = opt(comment)(inp)?;
+        if let Some(comment) = comment {
+            comments.push(comment);
+        }
+        let (inp, _) = opt(many0(alt((newline, unicode_space, tag("\\"), tag(";")))))(inp)?;
+        let (inp, end) = opt(eof)(inp)?;
+        if end.is_some() {
+            return Ok((inp, comments));
+        }
+        if input == inp {
+            panic!("invalid trailing text");
+        }
+        input = inp;
+    }
+}
+
 fn plain_identifier(input: &str) -> IResult<&str, KdlIdentifier, KdlParseError<&str>> {
     let start = input;
     let (input, name) = recognize(preceded(
@@ -643,6 +670,13 @@ mod comment_tests {
     fn slashdash() {
         assert_eq!(comment("/-foo 1 2"), Ok(("", "/-foo 1 2")));
     }
+
+    #[test]
+    fn surrounding() {
+        // assert_eq!(trailing_comments("// foo"), Ok(("", vec!["// foo"])));
+        // assert_eq!(trailing_comments("/* foo */"), Ok(("", vec!["/* foo */"])));
+        // assert_eq!(trailing_comments("/* foo */ \n // bar"), Ok(("", vec!["/* foo */", "// bar"])));
+    }
 }
 
 #[cfg(test)]

src/value.rs

@@ -165,22 +165,22 @@ impl Display for KdlValue {
 
 impl KdlValue {
     fn write_raw_string(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "r")?;
         let raw = self.as_string().unwrap();
         let mut consecutive = 0usize;
         let mut maxhash = 0usize;
         for char in raw.chars() {
             if char == '#' {
                 consecutive += 1;
             } else if char == '"' {
-                maxhash = maxhash.max(consecutive);
+                maxhash = maxhash.max(consecutive + 1);
             } else {
                 consecutive = 0;
             }
         }
-        write!(f, "{}", "#".repeat(maxhash + 1))?;
+        write!(f, "r")?;
+        write!(f, "{}", "#".repeat(maxhash))?;
         write!(f, "\"{}\"", raw)?;
-        write!(f, "{}", "#".repeat(maxhash + 1))?;
+        write!(f, "{}", "#".repeat(maxhash))?;
         Ok(())
     }
 }