Add improved tree-sitter navigation

This new set of functions (and tests) should eliminate
defun-navigation bugs and limitations we currently have.  This commit
doesn't change any existing bahavior: treesit-beginning/end-of-defun
and friends are unchanged.  The plan is to later switch gear and
replace the current functions with the new ones introduced in this
change.

This is a relatively big change, but I've setup a comprehensive test,
and it should fix current bugs, so I think it's ok to put it on the
release branch.

The gist of the new navigation is to use treesit--defuns-around to
find the previous sibling defun, next sibling defun, and the parent
defun, then use this information to move to previous/next
beginning/end of defun in treesit--navigate-defun.

I also added comprehensive testing that tests all four possible
operations (prev-beg, next-beg, prev-end, next-end) starting at all
possible positions (between two sibling defuns, inside a sibling
defun, etc).

* lisp/treesit.el (treesit-defun-type-regexp): Expand definition to
allow (REGEXP . FILTER).  Old functions don't support this, but it
should be fine since we are soon replacing them.

(treesit-defun-tactic)
(treesit-defun-skipper): New variables.

(treesit-default-defun-skipper)
(treesit--defuns-around)
(treesit--top-level-defun)
(treesit--navigate-defun): New functions.

* test/src/treesit-tests.el (treesit--ert-insert-and-parse-marker)
(treesit--ert-collect-positions)
(treesit--ert-test-defun-navigation): New helper functions.

(treesit--ert-defun-navigation-python-program)
(treesit--ert-defun-navigation-js-program)
(treesit--ert-defun-navigation-bash-program)
(treesit--ert-defun-navigation-nested-master): New variables.

(treesit-defun-navigation-nested-1)
(treesit-defun-navigation-nested-2)
(treesit-defun-navigation-nested-3)
(treesit-defun-navigation-top-level): New tests.

1 files changed, 207 insertions, 0 deletions

diff --git a/lisp/treesit.el b/lisp/treesit.el
index f176664bfde..3d7ae7031ef 100644
--- a/lisp/treesit.el
+++ b/lisp/treesit.el
@@ -1569,8 +1569,25 @@ BACKWARD and ALL are the same as in `treesit-search-forward'."
   "A regexp that matches the node type of defun nodes.
 For example, \"(function|class)_definition\".
 
+Sometimes not all nodes matched by the regexp are valid defuns.
+In that case, set this variable to a cons cell of the
+form (REGEXP . FILTER), where FILTER is a function that takes a
+node (the matched node) and returns t if node is valid, or nil
+for invalid node.
+
 This is used by `treesit-beginning-of-defun' and friends.")
 
+(defvar-local treesit-defun-tactic 'nested
+  "Determines how does Emacs treat nested defuns.
+If the value is `top-level', Emacs only move across top-level
+defuns, if the value is `nested', Emacs recognizes nested defuns.")
+
+(defvar-local treesit-defun-skipper #'treesit-default-defun-skipper
+  "A function called after tree-sitter navigation moved a step.
+It is called with no arguments.  By default, this function tries
+to move to the beginning of a line, either by moving to the empty
+newline after a defun, or the beginning of a defun.")
+
 (defvar-local treesit-defun-prefer-top-level nil
   "When non-nil, Emacs prefers top-level defun.
 
@@ -1639,6 +1656,196 @@ ARG is the same as in `beginning-of-defun'."
     (when top
       (goto-char (treesit-node-end top)))))
 
+(defun treesit-default-defun-skipper ()
+  "Skips spaces after navigating a defun.
+This fucntion tries to move to the beginning of a line, either by
+moving to the empty newline after a defun, or to the beginning of
+the current line if the beginning of the defun is indented."
+  (cond ((and (looking-at (rx (* (or " " "\\t")) "\n"))
+              (not (looking-at (rx bol))))
+         (goto-char (match-end 0)))
+        ((save-excursion
+           (skip-chars-backward " \t")
+           (eq (point) (line-beginning-position)))
+         (goto-char (line-beginning-position)))))
+
+;; prev-sibling:
+;; 1. end-of-node before pos
+;; 2. highest such node
+;;
+;; next-sibling:
+;; 1. beg-of-node after pos
+;; 2. highest such node
+;;
+;; parent:
+;; 1. node covers pos
+;; 2. smallest such node
+(defun treesit--defuns-around (pos regexp &optional pred)
+  "Return the previous, next, and parent defun around POS.
+
+Return a list of (PREV NEXT PARENT), where PREV and NEXT are
+previous and next sibling defuns around POS, and PARENT is the
+parent defun surrouding POS.  All of three could be nil if no
+sound defun exists.
+
+REGEXP and PRED are the same as in `treesit-defun-type-regexp'."
+  (let* ((node (treesit-node-at pos))
+         ;; NODE-BEFORE/AFTER = NODE when POS is completely in NODE,
+         ;; but if not, that means point could be in between two
+         ;; defun, in that case we want to use a node that's actually
+         ;; before/after point.
+         (node-before (if (>= (treesit-node-start node) pos)
+                          (treesit-search-forward-goto node "" t t t)
+                        node))
+         (node-after (if (<= (treesit-node-end node) pos)
+                         (treesit-search-forward-goto node "" nil nil t)
+                       node))
+         (result (list nil nil nil))
+         (pred (or pred (lambda (_) t))))
+    ;; 1. Find previous and next sibling defuns.
+    (cl-loop
+     for idx from 0 to 1
+     for node in (list node-before node-after)
+     for backward in '(t nil)
+     for pos-pred in (list (lambda (n) (<= (treesit-node-end n) pos))
+                           (lambda (n) (>= (treesit-node-start n) pos)))
+     ;; If point is inside a defun, our process below will never
+     ;; return a next/prev sibling outside of that defun, effectively
+     ;; any prev/next sibling is locked inside the smallest defun
+     ;; covering point, which is the correct behavior.  That's because
+     ;; when there exists a defun that covers point,
+     ;; `treesit-search-forward' will first reach that defun, after
+     ;; that we only go upwards in the tree, so other defuns outside
+     ;; of the covering defun is never reached.  (Don't use
+     ;; `treesit-search-forward-goto' as it breaks when NODE-AFTER is
+     ;; the last token of a parent defun: it will skip the parent
+     ;; defun because it wants to ensure progress.)
+     do (cl-loop for cursor = (when node
+                                (save-excursion
+                                  (treesit-search-forward
+                                   node regexp backward backward)))
+                 then (treesit-node-parent cursor)
+                 while cursor
+                 if (and (string-match-p
+                          regexp (treesit-node-type cursor))
+                         (funcall pred cursor)
+                         (funcall pos-pred cursor))
+                 do (setf (nth idx result) cursor)))
+    ;; 2. Find the parent defun.
+    (setf (nth 2 result)
+          (cl-loop for cursor = (or (nth 0 result)
+                                    (nth 1 result)
+                                    node)
+                   then (treesit-node-parent cursor)
+                   while cursor
+                   if (and (string-match-p
+                            regexp (treesit-node-type cursor))
+                           (funcall pred cursor)
+                           (not (member cursor result)))
+                   return cursor))
+    result))
+
+(defun treesit--top-level-defun (node regexp &optional pred)
+  "Return the top-level parent defun of NODE.
+REGEXP and PRED are the same as in `treesit-defun-type-regexp'."
+  (let* ((pred (or pred (lambda (_) t))))
+    ;; `treesit-search-forward-goto' will make sure the matched node
+    ;; is before POS.
+    (cl-loop for cursor = node
+             then (treesit-node-parent cursor)
+             while cursor
+             if (and (string-match-p
+                      regexp (treesit-node-type cursor))
+                     (funcall pred cursor))
+             do (setq node cursor))
+    node))
+
+(defun treesit--navigate-defun (pos arg side &optional recursing)
+  "Navigate defun ARG steps from POS.
+
+If ARG is positive, move forward that many steps, if negative,
+move backward.  If SIDE is `beg', stop at the beginning of a
+defun, if SIDE is `end', stop at the end.
+
+This function doesn't actaully move point, it just returns the
+position it would move to.  If there aren't enough defuns to move
+across, return nil.
+
+RECURSING is an internal parameter, if non-nil, it means this
+function is called recursively."
+  (pcase-let*
+      ((counter (abs arg))
+       (`(,regexp . ,pred)
+        (if (consp treesit-defun-type-regexp)
+            treesit-defun-type-regexp
+          (cons treesit-defun-type-regexp nil)))
+       ;; Move POS to the beg/end of NODE.  If NODE is nil, terminate.
+       ;; Return the position we moved to.
+       (advance (lambda (node)
+                  (let ((dest (pcase side
+                                ('beg (treesit-node-start node))
+                                ('end (treesit-node-end node)))))
+                    (if (null dest)
+                        (throw 'term nil)
+                      dest)))))
+    (catch 'term
+      (while (> counter 0)
+        (pcase-let
+            ((`(,prev ,next ,parent)
+              (treesit--defuns-around pos regexp pred)))
+          ;; When PARENT is nil, nested and top-level are the same,
+          ;; there there is a PARENT, make PARENT to be the top-level
+          ;; parent and pretend there is no nested PREV and NEXT.
+          (when (and (eq treesit-defun-tactic 'top-level)
+                     parent)
+            (setq parent (treesit--top-level-defun
+                          parent regexp pred)
+                  prev nil
+                  next nil))
+          ;; Move...
+          (if (> arg 0)
+              ;; ...forward.
+              (if (and (eq side 'beg)
+                       ;; Should we skip the defun (recurse)?
+                       (cond (next (not recursing)) ; [1] (see below)
+                             (parent t) ; [2]
+                             (t nil)))
+                  ;; Special case: go to next beg-of-defun.  Set POS
+                  ;; to the end of next/parent defun, and run one more
+                  ;; step.  If there is a next defun, step over it, so
+                  ;; we only need to recurse once, so we don't need to
+                  ;; recurse if we are already recursing [1]. If there
+                  ;; is no next but a parent, keep stepping out
+                  ;; (recursing) until we got out of the parents until
+                  ;; (1) there is a next sibling defun, or (2) no more
+                  ;; parents [2].
+                  (setq pos
+                        (or (treesit--navigate-defun
+                             (treesit-node-end (or next parent))
+                             1 'beg t)
+                            (throw 'term nil)))
+                ;; Normal case.
+                (setq pos (funcall advance (or next parent))))
+            ;; ...backward.
+            (if (and (eq side 'end)
+                     (cond (prev (not recursing))
+                           (parent t)
+                           (t nil)))
+                ;; Special case: go to prev end-of-defun.
+                (setq pos
+                      (or (treesit--navigate-defun
+                           (treesit-node-start (or prev parent))
+                           -1 'end t)
+                          (throw 'term nil)))
+              ;; Normal case.
+              (setq pos (funcall advance (or prev parent)))))
+          ;; A successful step! Decrement counter.
+          (cl-decf counter))))
+    ;; Counter equal to 0 means we successfully stepped ARG steps.
+    (if (eq counter 0)
+        pos
+      nil)))
+
 ;;; Activating tree-sitter
 
 (defun treesit-ready-p (language &optional quiet)