Add extend variants for tree-sitter textobjects

helix-core/src/movement.rs

@@ -389,6 +389,8 @@ fn reached_target(target: WordMotionTarget, prev_ch: char, next_ch: char) -> boo
     }
 }
 
+/// Finds the range of the next or previous textobject in the syntax sub-tree of `node`.
+/// Returns the range in the forwards direction.
 pub fn goto_treesitter_object(
     slice: RopeSlice,
     range: Range,
@@ -419,8 +421,8 @@ pub fn goto_treesitter_object(
                 .filter(|n| n.start_byte() > byte_pos)
                 .min_by_key(|n| n.start_byte())?,
             Direction::Backward => nodes
-                .filter(|n| n.start_byte() < byte_pos)
-                .max_by_key(|n| n.start_byte())?,
+                .filter(|n| n.end_byte() < byte_pos)
+                .max_by_key(|n| n.end_byte())?,
         };
 
         let len = slice.len_bytes();
@@ -434,7 +436,7 @@ pub fn goto_treesitter_object(
         let end_char = slice.byte_to_char(end_byte);
 
         // head of range should be at beginning
-        Some(Range::new(end_char, start_char))
+        Some(Range::new(start_char, end_char))
     };
     (0..count).fold(range, |range, _| get_range(range).unwrap_or(range))
 }

helix-core/src/selection.rs

@@ -122,7 +122,7 @@ impl Range {
         }
     }
 
-    // flips the direction of the selection
+    /// Flips the direction of the selection
     pub fn flip(&self) -> Self {
         Self {
             anchor: self.head,
@@ -131,6 +131,16 @@ impl Range {
         }
     }
 
+    /// Returns the selection if it goes in the direction of `direction`,
+    /// flipping the selection otherwise.
+    pub fn with_direction(self, direction: Direction) -> Self {
+        if self.direction() == direction {
+            self
+        } else {
+            self.flip()
+        }
+    }
+
     /// Check two ranges for overlap.
     #[must_use]
     pub fn overlaps(&self, other: &Self) -> bool {

helix-term/src/commands.rs

@@ -2159,10 +2159,7 @@ fn ensure_selections_forward(cx: &mut Context) {
     let selection = doc
         .selection(view.id)
         .clone()
-        .transform(|r| match r.direction() {
-            Direction::Forward => r,
-            Direction::Backward => r.flip(),
-        });
+        .transform(|r| r.with_direction(Direction::Forward));
 
     doc.set_selection(view.id, selection);
 }
@@ -4260,15 +4257,27 @@ fn goto_ts_object_impl(cx: &mut Context, object: &'static str, direction: Direct
             let root = syntax.tree().root_node();
 
             let selection = doc.selection(view.id).clone().transform(|range| {
-                movement::goto_treesitter_object(
+                let new_range = movement::goto_treesitter_object(
                     text,
                     range,
                     object,
                     direction,
                     root,
                     lang_config,
                     count,
-                )
+                );
+
+                if editor.mode == Mode::Select {
+                    let head = if new_range.head < range.anchor {
+                        new_range.anchor
+                    } else {
+                        new_range.head
+                    };
+
+                    Range::new(range.anchor, head)
+                } else {
+                    new_range.with_direction(direction)
+                }
             });
 
             doc.set_selection(view.id, selection);

helix-term/tests/test/movement.rs

@@ -85,3 +85,131 @@ async fn cursor_position_newly_opened_file() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn select_mode_tree_sitter_next_function_is_union_of_objects() -> anyhow::Result<()> {
+    test_with_config(
+        Args {
+            files: vec![(PathBuf::from("foo.rs"), Position::default())],
+            ..Default::default()
+        },
+        Config::default(),
+        (
+            helpers::platform_line(indoc! {"\
+                #[/|]#// Increments
+                fn inc(x: usize) -> usize { x + 1 }
+                /// Decrements
+                fn dec(x: usize) -> usize { x - 1 }
+            "})
+            .as_ref(),
+            "]fv]f",
+            helpers::platform_line(indoc! {"\
+                /// Increments
+                #[fn inc(x: usize) -> usize { x + 1 }
+                /// Decrements
+                fn dec(x: usize) -> usize { x - 1 }|]#
+            "})
+            .as_ref(),
+        ),
+    )
+    .await?;
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn select_mode_tree_sitter_prev_function_unselects_object() -> anyhow::Result<()> {
+    test_with_config(
+        Args {
+            files: vec![(PathBuf::from("foo.rs"), Position::default())],
+            ..Default::default()
+        },
+        Config::default(),
+        (
+            helpers::platform_line(indoc! {"\
+                /// Increments
+                #[fn inc(x: usize) -> usize { x + 1 }
+                /// Decrements
+                fn dec(x: usize) -> usize { x - 1 }|]#
+            "})
+            .as_ref(),
+            "v[f",
+            helpers::platform_line(indoc! {"\
+                /// Increments
+                #[fn inc(x: usize) -> usize { x + 1 }|]#
+                /// Decrements
+                fn dec(x: usize) -> usize { x - 1 }
+            "})
+            .as_ref(),
+        ),
+    )
+    .await?;
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn select_mode_tree_sitter_prev_function_goes_backwards_to_object() -> anyhow::Result<()> {
+    // Note: the anchor stays put and the head moves back.
+    test_with_config(
+        Args {
+            files: vec![(PathBuf::from("foo.rs"), Position::default())],
+            ..Default::default()
+        },
+        Config::default(),
+        (
+            helpers::platform_line(indoc! {"\
+                /// Increments
+                fn inc(x: usize) -> usize { x + 1 }
+                /// Decrements
+                fn dec(x: usize) -> usize { x - 1 }
+                /// Identity
+                #[fn ident(x: usize) -> usize { x }|]#
+            "})
+            .as_ref(),
+            "v[f",
+            helpers::platform_line(indoc! {"\
+                /// Increments
+                fn inc(x: usize) -> usize { x + 1 }
+                /// Decrements
+                #[|fn dec(x: usize) -> usize { x - 1 }
+                /// Identity
+                ]#fn ident(x: usize) -> usize { x }
+            "})
+            .as_ref(),
+        ),
+    )
+    .await?;
+
+    test_with_config(
+        Args {
+            files: vec![(PathBuf::from("foo.rs"), Position::default())],
+            ..Default::default()
+        },
+        Config::default(),
+        (
+            helpers::platform_line(indoc! {"\
+                /// Increments
+                fn inc(x: usize) -> usize { x + 1 }
+                /// Decrements
+                fn dec(x: usize) -> usize { x - 1 }
+                /// Identity
+                #[fn ident(x: usize) -> usize { x }|]#
+            "})
+            .as_ref(),
+            "v[f[f",
+            helpers::platform_line(indoc! {"\
+                /// Increments
+                #[|fn inc(x: usize) -> usize { x + 1 }
+                /// Decrements
+                fn dec(x: usize) -> usize { x - 1 }
+                /// Identity
+                ]#fn ident(x: usize) -> usize { x }
+            "})
+            .as_ref(),
+        ),
+    )
+    .await?;
+
+    Ok(())
+}