Reimplement tree motions in terms of syntax::TreeCursor

This uses the new TreeCursor type from the parent commit to reimplement
the tree-sitter motions (`A-p/o/i/n`). Other tree-sitter related
features like textobjects are not touched with this change and will
need a different, unrelated approach to solve.

3 files changed, 107 insertions, 50 deletions

diff --git a/helix-core/src/object.rs b/helix-core/src/object.rs
index d2d4fe70..0df105f1 100644
--- a/helix-core/src/object.rs
+++ b/helix-core/src/object.rs
@@ -1,42 +1,52 @@
-use crate::{Range, RopeSlice, Selection, Syntax};
-use tree_sitter::Node;
+use crate::{syntax::TreeCursor, Range, RopeSlice, Selection, Syntax};
 
 pub fn expand_selection(syntax: &Syntax, text: RopeSlice, selection: Selection) -> Selection {
-    select_node_impl(syntax, text, selection, |mut node, from, to| {
-        while node.start_byte() == from && node.end_byte() == to {
-            node = node.parent()?;
+    let cursor = &mut syntax.walk();
+
+    selection.transform(|range| {
+        let from = text.char_to_byte(range.from());
+        let to = text.char_to_byte(range.to());
+
+        let byte_range = from..to;
+        cursor.reset_to_byte_range(from, to);
+
+        while cursor.node().byte_range() == byte_range {
+            if !cursor.goto_parent() {
+                break;
+            }
         }
-        Some(node)
+
+        let node = cursor.node();
+        let from = text.byte_to_char(node.start_byte());
+        let to = text.byte_to_char(node.end_byte());
+
+        Range::new(to, from).with_direction(range.direction())
     })
 }
 
 pub fn shrink_selection(syntax: &Syntax, text: RopeSlice, selection: Selection) -> Selection {
-    select_node_impl(syntax, text, selection, |descendant, _from, _to| {
-        descendant.child(0).or(Some(descendant))
+    select_node_impl(syntax, text, selection, |cursor| {
+        cursor.goto_first_child();
     })
 }
 
-pub fn select_sibling<F>(
-    syntax: &Syntax,
-    text: RopeSlice,
-    selection: Selection,
-    sibling_fn: &F,
-) -> Selection
-where
-    F: Fn(Node) -> Option<Node>,
-{
-    select_node_impl(syntax, text, selection, |descendant, _from, _to| {
-        find_sibling_recursive(descendant, sibling_fn)
+pub fn select_next_sibling(syntax: &Syntax, text: RopeSlice, selection: Selection) -> Selection {
+    select_node_impl(syntax, text, selection, |cursor| {
+        while !cursor.goto_next_sibling() {
+            if !cursor.goto_parent() {
+                break;
+            }
+        }
     })
 }
 
-fn find_sibling_recursive<F>(node: Node, sibling_fn: F) -> Option<Node>
-where
-    F: Fn(Node) -> Option<Node>,
-{
-    sibling_fn(node).or_else(|| {
-        node.parent()
-            .and_then(|node| find_sibling_recursive(node, sibling_fn))
+pub fn select_prev_sibling(syntax: &Syntax, text: RopeSlice, selection: Selection) -> Selection {
+    select_node_impl(syntax, text, selection, |cursor| {
+        while !cursor.goto_prev_sibling() {
+            if !cursor.goto_parent() {
+                break;
+            }
+        }
     })
 }
 
@@ -44,33 +54,25 @@ fn select_node_impl<F>(
     syntax: &Syntax,
     text: RopeSlice,
     selection: Selection,
-    select_fn: F,
+    motion: F,
 ) -> Selection
 where
-    F: Fn(Node, usize, usize) -> Option<Node>,
+    F: Fn(&mut TreeCursor),
 {
-    let tree = syntax.tree();
+    let cursor = &mut syntax.walk();
 
     selection.transform(|range| {
         let from = text.char_to_byte(range.from());
         let to = text.char_to_byte(range.to());
 
-        let node = match tree
-            .root_node()
-            .descendant_for_byte_range(from, to)
-            .and_then(|node| select_fn(node, from, to))
-        {
-            Some(node) => node,
-            None => return range,
-        };
+        cursor.reset_to_byte_range(from, to);
 
+        motion(cursor);
+
+        let node = cursor.node();
         let from = text.byte_to_char(node.start_byte());
         let to = text.byte_to_char(node.end_byte());
 
-        if range.head < range.anchor {
-            Range::new(to, from)
-        } else {
-            Range::new(from, to)
-        }
+        Range::new(from, to).with_direction(range.direction())
     })
 }
diff --git a/helix-term/src/commands.rs b/helix-term/src/commands.rs
index 291ab1fe..ff72cee1 100644
--- a/helix-term/src/commands.rs
+++ b/helix-term/src/commands.rs
@@ -26,7 +26,6 @@ use helix_core::{
     syntax::{BlockCommentToken, LanguageServerFeature},
     text_annotations::TextAnnotations,
     textobject,
-    tree_sitter::Node,
     unicode::width::UnicodeWidthChar,
     visual_offset_from_block, Deletion, LineEnding, Position, Range, Rope, RopeGraphemes,
     RopeReader, RopeSlice, Selection, SmallVec, Tendril, Transaction,
@@ -4775,18 +4774,17 @@ fn shrink_selection(cx: &mut Context) {
     cx.editor.apply_motion(motion);
 }
 
-fn select_sibling_impl<F>(cx: &mut Context, sibling_fn: &'static F)
+fn select_sibling_impl<F>(cx: &mut Context, sibling_fn: F)
 where
-    F: Fn(Node) -> Option<Node>,
+    F: Fn(&helix_core::Syntax, RopeSlice, Selection) -> Selection + 'static,
 {
-    let motion = |editor: &mut Editor| {
+    let motion = move |editor: &mut Editor| {
         let (view, doc) = current!(editor);
 
         if let Some(syntax) = doc.syntax() {
             let text = doc.text().slice(..);
             let current_selection = doc.selection(view.id);
-            let selection =
-                object::select_sibling(syntax, text, current_selection.clone(), sibling_fn);
+            let selection = sibling_fn(syntax, text, current_selection.clone());
             doc.set_selection(view.id, selection);
         }
     };
@@ -4794,11 +4792,11 @@ where
 }
 
 fn select_next_sibling(cx: &mut Context) {
-    select_sibling_impl(cx, &|node| Node::next_sibling(&node))
+    select_sibling_impl(cx, object::select_next_sibling)
 }
 
 fn select_prev_sibling(cx: &mut Context) {
-    select_sibling_impl(cx, &|node| Node::prev_sibling(&node))
+    select_sibling_impl(cx, object::select_prev_sibling)
 }
 
 fn move_node_bound_impl(cx: &mut Context, dir: Direction, movement: Movement) {
diff --git a/helix-term/tests/test/movement.rs b/helix-term/tests/test/movement.rs
index 4ebaae85..1c25032c 100644
--- a/helix-term/tests/test/movement.rs
+++ b/helix-term/tests/test/movement.rs
@@ -635,3 +635,60 @@ async fn test_surround_delete() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn tree_sitter_motions_work_across_injections() -> anyhow::Result<()> {
+    test_with_config(
+        AppBuilder::new().with_file("foo.html", None),
+        (
+            "<script>let #[|x]# = 1;</script>",
+            "<A-o>",
+            "<script>let #[|x = 1]#;</script>",
+        ),
+    )
+    .await?;
+
+    // When the full injected layer is selected, expand_selection jumps to
+    // a more shallow layer.
+    test_with_config(
+        AppBuilder::new().with_file("foo.html", None),
+        (
+            "<script>#[|let x = 1;]#</script>",
+            "<A-o>",
+            "#[|<script>let x = 1;</script>]#",
+        ),
+    )
+    .await?;
+
+    test_with_config(
+        AppBuilder::new().with_file("foo.html", None),
+        (
+            "<script>let #[|x = 1]#;</script>",
+            "<A-i>",
+            "<script>let #[|x]# = 1;</script>",
+        ),
+    )
+    .await?;
+
+    test_with_config(
+        AppBuilder::new().with_file("foo.html", None),
+        (
+            "<script>let #[|x]# = 1;</script>",
+            "<A-n>",
+            "<script>let x #[|=]# 1;</script>",
+        ),
+    )
+    .await?;
+
+    test_with_config(
+        AppBuilder::new().with_file("foo.html", None),
+        (
+            "<script>let #[|x]# = 1;</script>",
+            "<A-p>",
+            "<script>#[|let]# x = 1;</script>",
+        ),
+    )
+    .await?;
+
+    Ok(())
+}