From 3b750448561ad2a57bb6c5f348b18a96da527921 Mon Sep 17 00:00:00 2001
From: Michael Davis <mcarsondavis@gmail.com>
Date: Wed, 10 Jan 2024 17:55:29 -0500
Subject: [PATCH] Select subtree within injections in :tree-sitter-subtree

`: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 102ecb15d34e..0babe8c921ce 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 f1c3cb71cef9..b13af03a2030 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```");