From c950bdc21c7c2826f6bdf9fcdf4e42f6f5fcbca4 Mon Sep 17 00:00:00 2001
From: Clark Fischer <439978+clarkf@users.noreply.github.com>
Date: Wed, 14 Feb 2024 06:22:32 +0000
Subject: [PATCH] Account for trailing comments in span handling (#527)

---
 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 0d71c512a..f015e715f 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#"