1 files changed, 44 insertions, 42 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())
     })
 }