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'.")
 
+
+;;; 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