From 966fbc59849701dc56f70af406444c351bc9f54a Mon Sep 17 00:00:00 2001
From: Gokul Soumya <gokulps15@gmail.com>
Date: Thu, 3 Feb 2022 13:41:25 +0530
Subject: [PATCH] Add tree-sitter based function, class navigation

---
 helix-core/src/movement.rs | 52 +++++++++++++++++++++++++++++++++++++-
 helix-term/src/commands.rs | 52 ++++++++++++++++++++++++++++++++++++++
 helix-term/src/keymap.rs   |  6 +++++
 3 files changed, 109 insertions(+), 1 deletion(-)

diff --git a/helix-core/src/movement.rs b/helix-core/src/movement.rs
index 47fe6827..e9cd299b 100644
--- a/helix-core/src/movement.rs
+++ b/helix-core/src/movement.rs
@@ -1,6 +1,7 @@
 use std::iter;
 
 use ropey::iter::Chars;
+use tree_sitter::{Node, QueryCursor};
 
 use crate::{
     chars::{categorize_char, char_is_line_ending, CharCategory},
@@ -9,7 +10,10 @@ use crate::{
         next_grapheme_boundary, nth_next_grapheme_boundary, nth_prev_grapheme_boundary,
         prev_grapheme_boundary,
     },
-    pos_at_coords, Position, Range, RopeSlice,
+    pos_at_coords,
+    syntax::LanguageConfiguration,
+    textobject::TextObject,
+    Position, Range, RopeSlice,
 };
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
@@ -305,6 +309,52 @@ fn reached_target(target: WordMotionTarget, prev_ch: char, next_ch: char) -> boo
     }
 }
 
+pub fn goto_treesitter_object(
+    slice: RopeSlice,
+    range: Range,
+    object_name: &str,
+    dir: Direction,
+    slice_tree: Node,
+    lang_config: &LanguageConfiguration,
+    _count: usize,
+) -> Range {
+    let get_range = move || -> Option<Range> {
+        let byte_pos = slice.char_to_byte(range.cursor(slice));
+
+        let capture_name = format!("{}.{}", object_name, TextObject::Around);
+        let mut cursor = QueryCursor::new();
+        let nodes = lang_config.textobject_query()?.capture_nodes(
+            &capture_name,
+            slice_tree,
+            slice,
+            &mut cursor,
+        )?;
+
+        let node = match dir {
+            Direction::Forward => nodes
+                .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())?,
+        };
+
+        let len = slice.len_bytes();
+        let start_byte = node.start_byte();
+        let end_byte = node.end_byte();
+        if start_byte >= len || end_byte >= len {
+            return None;
+        }
+
+        let start_char = slice.byte_to_char(start_byte);
+        let end_char = slice.byte_to_char(end_byte);
+
+        // head of range should be at beginning
+        Some(Range::new(end_char, start_char))
+    };
+    get_range().unwrap_or(range)
+}
+
 #[cfg(test)]
 mod test {
     use ropey::Rope;
diff --git a/helix-term/src/commands.rs b/helix-term/src/commands.rs
index 1454a93f..10c18a85 100644
--- a/helix-term/src/commands.rs
+++ b/helix-term/src/commands.rs
@@ -398,6 +398,12 @@ impl MappableCommand {
         surround_delete, "Surround delete",
         select_textobject_around, "Select around object",
         select_textobject_inner, "Select inside object",
+        goto_next_function, "Goto next function",
+        goto_prev_function, "Goto previous function",
+        goto_next_class, "Goto next class",
+        goto_prev_class, "Goto previous class",
+        goto_next_parameter, "Goto next parameter",
+        goto_prev_parameter, "Goto previous parameter",
         dap_launch, "Launch debug target",
         dap_toggle_breakpoint, "Toggle breakpoint",
         dap_continue, "Continue program execution",
@@ -5907,6 +5913,52 @@ fn scroll_down(cx: &mut Context) {
     scroll(cx, cx.count(), Direction::Forward);
 }
 
+fn goto_ts_object_impl(cx: &mut Context, object: &str, direction: Direction) {
+    let count = cx.count();
+    let (view, doc) = current!(cx.editor);
+    let text = doc.text().slice(..);
+    let range = doc.selection(view.id).primary();
+
+    let new_range = match doc.language_config().zip(doc.syntax()) {
+        Some((lang_config, syntax)) => movement::goto_treesitter_object(
+            text,
+            range,
+            object,
+            direction,
+            syntax.tree().root_node(),
+            lang_config,
+            count,
+        ),
+        None => range,
+    };
+
+    doc.set_selection(view.id, Selection::single(new_range.anchor, new_range.head));
+}
+
+fn goto_next_function(cx: &mut Context) {
+    goto_ts_object_impl(cx, "function", Direction::Forward)
+}
+
+fn goto_prev_function(cx: &mut Context) {
+    goto_ts_object_impl(cx, "function", Direction::Backward)
+}
+
+fn goto_next_class(cx: &mut Context) {
+    goto_ts_object_impl(cx, "class", Direction::Forward)
+}
+
+fn goto_prev_class(cx: &mut Context) {
+    goto_ts_object_impl(cx, "class", Direction::Backward)
+}
+
+fn goto_next_parameter(cx: &mut Context) {
+    goto_ts_object_impl(cx, "parameter", Direction::Forward)
+}
+
+fn goto_prev_parameter(cx: &mut Context) {
+    goto_ts_object_impl(cx, "parameter", Direction::Backward)
+}
+
 fn select_textobject_around(cx: &mut Context) {
     select_textobject(cx, textobject::TextObject::Around);
 }
diff --git a/helix-term/src/keymap.rs b/helix-term/src/keymap.rs
index fc90da33..f414f797 100644
--- a/helix-term/src/keymap.rs
+++ b/helix-term/src/keymap.rs
@@ -572,11 +572,17 @@ impl Default for Keymaps {
             "[" => { "Left bracket"
                 "d" => goto_prev_diag,
                 "D" => goto_first_diag,
+                "f" => goto_prev_function,
+                "c" => goto_prev_class,
+                "p" => goto_prev_parameter,
                 "space" => add_newline_above,
             },
             "]" => { "Right bracket"
                 "d" => goto_next_diag,
                 "D" => goto_last_diag,
+                "f" => goto_next_function,
+                "c" => goto_next_class,
+                "p" => goto_next_parameter,
                 "space" => add_newline_below,
             },