From 554b380c2b31c633796999ef5ca5c25437f78250 Mon Sep 17 00:00:00 2001
From: Michael Davis <mcarsondavis@gmail.com>
Date: Mon, 15 Jan 2024 01:34:38 -0500
Subject: [PATCH] Select subtree within injections in :tree-sitter-subtree
 (#9309)

`:tree-sitter-subtree` could previously only print subtrees of nodes
in the root injection layer. We can improve on that by finding the layer
that contains the given byte range and printing the subtree within that
layer. That gives more useful results when a selection is within an
injection layer.
---
 helix-core/src/syntax.rs         | 43 ++++++++++++++++++++++++++++++++
 helix-term/src/commands/typed.rs |  6 +----
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/helix-core/src/syntax.rs b/helix-core/src/syntax.rs
index 4e44c48666e24..83bd09b4dc0e9 100644
--- a/helix-core/src/syntax.rs
+++ b/helix-core/src/syntax.rs
@@ -1338,6 +1338,23 @@ impl Syntax {
         result
     }
 
+    pub fn descendant_for_byte_range(&self, start: usize, end: usize) -> Option<Node<'_>> {
+        let mut container_id = self.root;
+
+        for (layer_id, layer) in self.layers.iter() {
+            if layer.depth > self.layers[container_id].depth
+                && layer.contains_byte_range(start, end)
+            {
+                container_id = layer_id;
+            }
+        }
+
+        self.layers[container_id]
+            .tree()
+            .root_node()
+            .descendant_for_byte_range(start, end)
+    }
+
     // Commenting
     // comment_strings_for_pos
     // is_commented
@@ -1434,6 +1451,32 @@ impl LanguageLayer {
         self.tree = Some(tree);
         Ok(())
     }
+
+    /// Whether the layer contains the given byte range.
+    ///
+    /// If the layer has multiple ranges (i.e. combined injections), the
+    /// given range is considered contained if it is within the start and
+    /// end bytes of the first and last ranges **and** if the given range
+    /// starts or ends within any of the layer's ranges.
+    fn contains_byte_range(&self, start: usize, end: usize) -> bool {
+        let layer_start = self
+            .ranges
+            .first()
+            .expect("ranges should not be empty")
+            .start_byte;
+        let layer_end = self
+            .ranges
+            .last()
+            .expect("ranges should not be empty")
+            .end_byte;
+
+        layer_start <= start
+            && layer_end >= end
+            && self.ranges.iter().any(|range| {
+                let byte_range = range.start_byte..range.end_byte;
+                byte_range.contains(&start) || byte_range.contains(&end)
+            })
+    }
 }
 
 pub(crate) fn generate_edits(
diff --git a/helix-term/src/commands/typed.rs b/helix-term/src/commands/typed.rs
index f1c3cb71cef94..b13af03a20307 100644
--- a/helix-term/src/commands/typed.rs
+++ b/helix-term/src/commands/typed.rs
@@ -2117,11 +2117,7 @@ fn tree_sitter_subtree(
         let text = doc.text();
         let from = text.char_to_byte(primary_selection.from());
         let to = text.char_to_byte(primary_selection.to());
-        if let Some(selected_node) = syntax
-            .tree()
-            .root_node()
-            .descendant_for_byte_range(from, to)
-        {
+        if let Some(selected_node) = syntax.descendant_for_byte_range(from, to) {
             let mut contents = String::from("```tsq\n");
             helix_core::syntax::pretty_print_tree(&mut contents, selected_node)?;
             contents.push_str("\n```");