feat(fmt): Handle all possible comments in repeat

Author: felix-andreas

src/format.rs

@@ -429,7 +429,7 @@ fn traverse(
                 .skip(1)
                 .take(node.child_count() - 2)
                 .map(|child| {
-                    let prev_sibling = {
+                    let maybe_prev = {
                         let tmp = prev_sibling;
                         prev_sibling = Some(child);
                         tmp
@@ -448,21 +448,21 @@ fn traverse(
                         fmt(child)?
                     };
                     if is_fmt_skip_comment(&child)
-                        && prev_sibling
+                        && maybe_prev
                             .map(|prev| end_position(prev).row < child.start_position().row)
                             .unwrap_or(true)
                     {
                         fmt_skip = true;
                     }
                     if child.kind() == "comment" {
-                        return Ok(match prev_sibling {
+                        return Ok(match maybe_prev {
                             Some(prev) if same_line(prev, child) => {
                                 format!(" {formatted}")
                             }
-                            Some(prev_node) => format!(
+                            Some(prev) => format!(
                                 "{}{formatted}",
                                 line_ending.repeat(usize::clamp(
-                                    child.start_position().row - end_position(prev_node).row,
+                                    child.start_position().row - end_position(prev).row,
                                     1,
                                     2,
                                 ))
@@ -473,7 +473,7 @@ fn traverse(
                     if child.kind() == "comma" {
                         is_first_arg = false;
                         return Ok(
-                            if prev_sibling
+                            if maybe_prev
                                 .map(|prev| prev.kind() == "comment")
                                 .unwrap_or(false)
                             {
@@ -486,7 +486,7 @@ fn traverse(
                     let result = format!(
                         "{}{}",
                         if is_first_arg {
-                            if prev_sibling
+                            if maybe_prev
                                 .map(|prev| prev.kind() == "comment")
                                 .unwrap_or(false)
                             {
@@ -496,7 +496,7 @@ fn traverse(
                             }
                         } else if is_multiline {
                             line_ending.repeat(usize::clamp(
-                                prev_sibling
+                                maybe_prev
                                     .map(|prev| child.start_position().row - end_position(prev).row)
                                     .unwrap_or(1),
                                 1,
@@ -735,9 +735,9 @@ fn traverse(
             let mut out = String::with_capacity(node.end_byte() - node.start_byte());
             let mut indent_comments = false;
             tree::for_each_child(&mut node.walk(), |child, field_name| {
-                let prev_node = child.prev_sibling();
+                let maybe_prev = child.prev_sibling();
 
-                let prev_is_comment = prev_node
+                let prev_is_comment = maybe_prev
                     .map(|node| node.kind() == "comment")
                     .unwrap_or(false);
                 let next_is_comment = child
@@ -763,8 +763,8 @@ fn traverse(
                             }
                         }
                         "comment" => {
-                            if let Some(prev_node) = prev_node
-                                && same_line(prev_node, child)
+                            if let Some(prev) = maybe_prev
+                                && same_line(prev, child)
                             {
                                 out.push(' ');
                             } else {
@@ -869,9 +869,9 @@ fn traverse(
             let mut out = String::with_capacity(node.end_byte() - node.start_byte());
             let mut indent_comments = false;
             tree::for_each_child(&mut node.walk(), |child, field_name| {
-                let prev_node = child.prev_sibling();
+                let maybe_prev = child.prev_sibling();
 
-                let prev_is_comment = prev_node
+                let prev_is_comment = maybe_prev
                     .map(|node| node.kind() == "comment")
                     .unwrap_or(false);
 
@@ -885,8 +885,8 @@ fn traverse(
                             out.push_str("else")
                         }
                         "comment" => {
-                            if let Some(prev_node) = prev_node
-                                && same_line(prev_node, child)
+                            if let Some(prev) = maybe_prev
+                                && same_line(prev, child)
                             {
                                 out.push(' ');
                             } else {
@@ -1160,15 +1160,49 @@ fn traverse(
                 .collect::<Result<String, FormatError>>()?
         }
         "repeat_statement" => {
-            let body = field(node, "body")?;
-            format!(
-                "repeat {}",
-                if body.kind() == "braced_expression" {
-                    fmt_multiline(body, true)?
-                } else {
-                    wrap_with_braces(body)?
-                }
-            )
+            handles_comments = true;
+
+            let mut out = String::with_capacity(node.end_byte() - node.start_byte());
+            tree::for_each_child(&mut node.walk(), |child, field_name| {
+                let maybe_prev = child.prev_sibling();
+
+                let prev_is_comment = maybe_prev
+                    .map(|node| node.kind() == "comment")
+                    .unwrap_or(false);
+
+                match field_name {
+                    None => match child.kind() {
+                        "repeat" => out.push_str("repeat"),
+                        "comment" => {
+                            if let Some(prev) = maybe_prev
+                                && same_line(prev, child)
+                            {
+                                out.push(' ');
+                            } else if !prev_is_comment {
+                                out.push('\n');
+                            }
+                            out.push_str(&fmt(child)?);
+                            out.push('\n')
+                        }
+                        _ => unreachable!(),
+                    },
+                    Some(field_name) => {
+                        if !prev_is_comment {
+                            out.push(' ');
+                        }
+                        match field_name {
+                            "body" => out.push_str(&if child.kind() == "braced_expression" {
+                                fmt_multiline(child, true)?
+                            } else {
+                                wrap_with_braces(child)?
+                            }),
+                            _ => unreachable!(),
+                        }
+                    }
+                };
+                Ok::<_, FormatError>(())
+            })?;
+            out
         }
         "string" => {
             let maybe_string_content = field_optional(node, "content");
@@ -1236,18 +1270,18 @@ fn traverse(
             let mut out = String::with_capacity(node.end_byte() - node.start_byte());
             let mut indent_comments = false;
             tree::for_each_child(&mut node.walk(), |child, field_name| {
-                let prev_node = child.prev_sibling();
+                let maybe_prev = child.prev_sibling();
 
-                let prev_is_comment = prev_node
+                let prev_is_comment = maybe_prev
                     .map(|node| node.kind() == "comment")
                     .unwrap_or(false);
 
                 match field_name {
                     None => match child.kind() {
                         "while" => out.push_str("while"),
                         "comment" => {
-                            if let Some(prev_node) = prev_node
-                                && same_line(prev_node, child)
+                            if let Some(prev) = maybe_prev
+                                && same_line(prev, child)
                             {
                                 out.push(' ');
                             } else {

src/main.rs

@@ -67,13 +67,13 @@ enum Command {
     Fmt {
         /// R files to format
         files: Option<Vec<PathBuf>>,
-        /// Avoid writing any formatted files back; instead, exit with a non-zero status code if any files would have been modified, and zero otherwise
+        /// Exit with error if files would be modified without making changes
         #[clap(long, default_value_t = false)]
         check: bool,
-        /// Avoid writing any formatted files back; instead, exit with a non-zero status code and the difference between the current file and how the formatted file would look like
+        /// Show diff instead of modifying files; exit with error if changes needed
         #[clap(long, default_value_t = false)]
         diff: bool,
-        /// Stop if a comment can't be properly placed (exits with error if a comment would be moved to different location)
+        /// Stop if a comment can't be properly placed (exits with error if any comment would be moved to different location)
         #[clap(long, default_value_t = false)]
         stop_on_unhandled_comment: bool,
     },

tests/snapshots/test_format__repeat_statement-2.snap

@@ -0,0 +1,7 @@
+---
+source: tests/test_format.rs
+expression: "format_str(indoc!\n{\n    r#\"\n        repeat # repeat\n        # 1\n        {\n        }\n    \"#\n}).unwrap()"
+---
+repeat # repeat
+# 1
+{}

tests/snapshots/test_format__repeat_statement-3.snap

@@ -0,0 +1,12 @@
+---
+source: tests/test_format.rs
+expression: "format_str(indoc!\n{\n    r#\"\n        repeat # 1\n            body\n        repeat # 1\n        { body }\n        \"#\n}).unwrap()"
+---
+repeat # 1
+{
+  body
+}
+repeat # 1
+{
+  body
+}

tests/snapshots/test_format__repeat_statement.snap

@@ -1,13 +1,16 @@
 ---
-source: src/format.rs
-expression: "format_str(indoc!\n{\n    r#\"\n            repeat {\n                print(\"Hello, world!\")\n            }\n            repeat {\n            }\n            repeat { #foo\n            }\n            repeat #foo\n            { }\n        \"#\n})"
+source: tests/test_format.rs
+expression: "format_str(indoc!\n{\n    r#\"\n        repeat { body }\n        repeat {\n            body\n        }\n        repeat body\n        repeat\n            body\n    \"#\n}).unwrap()"
 ---
 repeat {
-  print("Hello, world!")
+  body
 }
-repeat {}
 repeat {
-  # foo
+  body
+}
+repeat {
+  body
+}
+repeat {
+  body
 }
-# foo
-repeat {}

tests/test_format.rs

@@ -559,16 +559,31 @@ fn program() {
 #[test]
 fn repeat_statement() {
     assert_fmt! {r#"
+        repeat { body }
         repeat {
-            print("Hello, world!")
-        }
-        repeat {
+            body
         }
-        repeat { #foo
+        repeat body
+        repeat
+            body
+    "#};
+
+    // check all possible comment positions
+    assert_fmt! {r#"
+        repeat # repeat
+        # 1
+        {
         }
-        repeat #foo
-        { }
     "#};
+
+    // comment after condition but no block
+    assert_fmt! {r#"
+        repeat # 1
+            body
+        repeat # 1
+        { body }
+        "#
+    };
 }
 
 #[test]