From eca3ccff76b9f45dd2b266215764fe7630dee884 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 4e44c4866..83bd09b4d 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 f1c3cb71c..b13af03a2 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```");