Code navigation for sh-mode with Tree-sitter

["João Paulo Labegalini de Carvalho" <jaopaulolc@gmail.com>]

[-- Attachment #1.1: Type: text/plain, Size: 1930 bytes --]

Hi all,

I am working on basic navigation for sh-mode. My idea is that in sh-mode
invoking C-M-a or C-M-e moves point to the beginning/end of the
surrounding function if the point was inside of said function or to the
beginning/end of the next/previous top-level function otherwise.

I got the functions to do so in the attached patch (I did not include the
C-M-a with a negative argument yet).

If the functions are defined and used in isolation they work as intended.
However, because `end-of-defun' calls `beginning-of-defun-raw' (which,
AFAIK, uses `beginning-of-defun-function') it causes the C-M-e to not work
when the point is not inside of a function.

What I think is happening is that, when `end-of-defun' calls
`beginning-of-defun-raw' it brings point to the beginning of a top-level
function, so the subsequent call to `end-of-defun' moves point to the start
location, does it make it seems as the point did not move.

I have also noticed that calling the private function I defined directly
via M-: brings the point to a slightly different location. For example:

<---- point before C-M-e
A {
} <--- point after M-: (sh-mode--treesit-end-of-defun)
<---- pint after C-M-e

Calling sh-mode--treesit-end-of-defun brings the point right after the
closing curly brace and with C-M-e the point is positioned at the next line.

That might be due to additional processing done inside `end-of-defun'.

So, would it be ok to override `beginning-of-defun' and `end-of-defun' and
thus have bash-ts-mode use the functions I wrote directly? If so, how could
I do that? If not, how can I fix the "miss-behavior"?

I would appreciate all feedback and suggestions.

-- 
João Paulo L. de Carvalho
Ph.D Computer Science |  IC-UNICAMP | Campinas , SP - Brazil
Postdoctoral Research Fellow | University of Alberta | Edmonton, AB - Canada
joao.carvalho@ic.unicamp.br
joao.carvalho@ualberta.ca

[-- Attachment #1.2: Type: text/html, Size: 2423 bytes --]

[-- Attachment #2: 0001-Basic-navigation-for-sh-mode.patch --]
[-- Type: text/x-patch, Size: 3759 bytes --]

From ec7e99cc93fb8f11d94e4c02c558cbfd756237a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20P=2E=20L=2E=20de=20Carvalho?=
 <jaopaulolc@gmail.com>
Date: Sat, 3 Dec 2022 12:55:27 -0700
Subject: [PATCH] Basic navigation for sh-mode

---
 lisp/progmodes/sh-script.el | 63 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 62 insertions(+), 1 deletion(-)

diff --git a/lisp/progmodes/sh-script.el b/lisp/progmodes/sh-script.el
index 408ebfc045..5885261c01 100644
--- a/lisp/progmodes/sh-script.el
+++ b/lisp/progmodes/sh-script.el
@@ -1619,7 +1619,13 @@ bash-ts-mode
                   ( bracket delimiter misc-punctuation operator)))
     (setq-local treesit-font-lock-settings
                 sh-mode--treesit-settings)
-    (treesit-major-mode-setup)))
+    (setq-local treesit-defun-prefer-top-level t)
+    (setq-local treesit-defun-type-regexp "function_definition")
+    (treesit-major-mode-setup)
+    (setq-local beginning-of-defun-function
+                #'sh-mode--treesit-beginning-of-defun)
+    (setq-local end-of-defun-function
+                #'sh-mode--treesit-end-of-defun)))
 
 (advice-add 'bash-ts-mode :around #'sh--redirect-bash-ts-mode
             ;; Give it lower precedence than normal advice, so other
@@ -3364,5 +3370,60 @@ sh-mode--treesit-settings
    '((["$"]) @font-lock-misc-punctuation-face))
   "Tree-sitter font-lock settings for `sh-mode'.")
 
+\f
+;;; Tree-sitter navigation
+
+(defmacro get-fun-or-nil (node)
+  `((treesit-parent-until ,node
+                        (lambda (p)
+                          (string= (treesit-node-type p)
+                                   "function_definition"))))
+
+(defun sh-mode--treesit-beginning-of-defun (&optional arg)
+  "Tree-sitter `beginning-of-defun' function.
+ARG is the same as in `beginning-of-defun'."
+  (let ((arg (or arg 1))
+        (target-node nil)
+        (node-at-point nil)
+        (function treesit-defun-type-regexp))
+    (if (> arg 0)
+        ;; Go backward.
+        (while (> arg 0)
+          (setq node-at-point (treesit-node-at (point)))
+          (if (string= (treesit-node-type node-at-point) "function")
+              (setq node-at-point (treesit-node-parent node-at-point)))
+          (setq target-node (get-fun-or-nil node-at-point))
+          (unless target-node
+            (setq maybe-target-node (treesit-search-forward node-at-point
+                                                            function
+                                                            t))
+            (setq target-node (or (treesit-node-top-level maybe-target-node)
+                                  maybe-target-node)))
+          (when target-node
+            (goto-char (treesit-node-start target-node)))
+          (setq arg (1- arg)))
+      ;; Go forward.
+      (while (< arg 0)
+        (setq arg (1+ arg))))))
+
+(defun sh-mode--treesit-end-of-defun ()
+  "Tree-sitter `end-of-defun' function."
+  (let ((node-at-point nil)
+        (target-node nil)
+        (function treesit-defun-type-regexp))
+    (setq node-at-point (treesit-node-at (point)))
+    (setq target-node (get-fun-or-nil node-at-point))
+    (unless target-node
+      (setq target-node (treesit-search-forward node-at-point
+                                                function))
+      (when (and target-node
+                 (treesit-parent-until target-node
+                                       (lambda (p)
+                                         (string= (treesit-node-type p)
+                                                  function))))
+        (setq target-node (treesit-node-top-level target-node))))
+    (when target-node
+      (goto-char (treesit-node-end target-node)))))
+
 (provide 'sh-script)
 ;;; sh-script.el ends here
-- 
2.31.1