From 3f9678b09faac9ee18d205ffe3de9aaa65bb1c6f Mon Sep 17 00:00:00 2001
From: Clark Fischer <clark.fischer@gmail.com>
Date: Tue, 2 Jan 2024 10:20:45 -0800
Subject: [PATCH] Account for trailing comments in span handling

Fixes https://github.com/tamasfe/taplo/issues/464

When using taplo rules to limit configuration to specific portions of a
document, trailing comments would break formatting.

`Node#text_ranges`, when considering tables, was returning ranges
excluding trailing comments. When the formatter walks the code and is
trying to determine which rules apply to a table, it checks whether the
table is a subset of the rule "span". When there is a trailing comment,
this check fails, and the rule will never apply.

Since the dom `Node` for tables contains only a reference to the table
header, we can get the text range of the entire table, including
trailing or interior comments, by traversing up to the parent and taking
its text range. If this is unavailable, the previous behavior is used.

Note that this modifies the behavior of `crate::dom::Node::text_ranges`,
which is part of the public API.

Signed-off-by: Clark Fischer <clark.fischer@gmail.com>
---
 crates/taplo/src/dom/node.rs        | 47 ++++++++++++++++++++++++++++-
 crates/taplo/src/tests/formatter.rs | 35 +++++++++++++++++++++
 2 files changed, 81 insertions(+), 1 deletion(-)

diff --git a/crates/taplo/src/dom/node.rs b/crates/taplo/src/dom/node.rs
index af2c51d76..f04132956 100644
--- a/crates/taplo/src/dom/node.rs
+++ b/crates/taplo/src/dom/node.rs
@@ -244,6 +244,11 @@ impl Node {
         Ok(all.into_iter())
     }
 
+    /// Determine the portions of the source that correspond to this
+    /// `Node`.
+    ///
+    /// Since a symbol may appear in multiple places, a single node
+    /// may correspond to multiple `TextRange`s.
     pub fn text_ranges(&self) -> impl ExactSizeIterator<Item = TextRange> {
         let mut ranges = Vec::with_capacity(1);
 
@@ -256,7 +261,20 @@ impl Node {
                     ranges.extend(entry.text_ranges());
                 }
 
-                if let Some(mut r) = v.syntax().map(|s| s.text_range()) {
+                // consider both v.syntax() and its parent node when
+                // broadening the range. v.syntax() is only the
+                // TableHeader, where its parent is the full Table
+                // syntax node. Take a range which encompasses both.
+                let syntax_range = match (
+                    v.syntax().map(|n| n.text_range()),
+                    v.syntax().and_then(|h| h.parent()).map(|n| n.text_range()),
+                ) {
+                    (Some(child), Some(parent)) => Some(child.cover(parent)),
+                    (Some(child), None) => Some(child),
+                    (None, Some(parent)) => Some(parent),
+                    _ => None,
+                };
+                if let Some(mut r) = syntax_range {
                     for range in &ranges {
                         r = r.cover(*range);
                     }
@@ -659,3 +677,30 @@ impl From<Invalid> for Node {
         Self::Invalid(v)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_dom_find_all_matches_includes_comments() {
+        let source = r#"
+[table]
+key = "value" # comment
+"#;
+        let dom = crate::parser::parse(source).into_dom();
+        let matches = dom
+            .find_all_matches("table".parse().unwrap(), false)
+            .unwrap()
+            .collect::<Vec<_>>();
+        assert_eq!(1, matches.len());
+        let (_, node) = matches.first().unwrap();
+
+        // Ensure that at least one of the resultant spans includes
+        // the comment.
+        let spans = node.text_ranges().map(|r| &source[r]).collect::<Vec<_>>();
+        assert!(
+            spans.iter().any(|s| s.contains('#')),
+            "expected one of {:?} to contain a comment",
+            &spans
+        );
+    }
+}
diff --git a/crates/taplo/src/tests/formatter.rs b/crates/taplo/src/tests/formatter.rs
index 8be146c30..d69f986a0 100644
--- a/crates/taplo/src/tests/formatter.rs
+++ b/crates/taplo/src/tests/formatter.rs
@@ -1025,6 +1025,41 @@ foo = [
     assert_format!(expected, &formatted);
 }
 
+/// See https://github.com/tamasfe/taplo/issues/464
+#[test]
+fn test_reorder_keys_trailing_comment() {
+    let src = r#"
+[mytable]
+a = "a"
+c = "c"
+b = "b" # ...
+"#;
+
+    let expected = r#"
+[mytable]
+a = "a"
+b = "b" # ...
+c = "c"
+"#;
+
+    let p = crate::parser::parse(src);
+    let formatted = crate::formatter::format_with_path_scopes(
+        p.into_dom(),
+        Default::default(),
+        &[],
+        vec![(
+            "mytable",
+            formatter::OptionsIncomplete {
+                reorder_keys: Some(true),
+                ..Default::default()
+            },
+        )],
+    )
+    .unwrap();
+
+    assert_format!(expected, &formatted);
+}
+
 #[test]
 fn test_single_comment_no_alignment() {
     let src = r#"