From 9934c4225633ff0524269902b8c574666cdfe0cb Mon Sep 17 00:00:00 2001
From: Theodor Thornhill <theo@thornhill.no>
Date: Tue, 17 Jan 2023 21:18:29 +0100
Subject: [PATCH] Add treesit-forward-sexp

* lisp/progmodes/java-ts-mode.el (java-ts-mode): Use
treesit-sexp-type-regexp.
* lisp/treesit.el (treesit-sexp-type-regexp): New defvar.
(treesit-forward-sexp): New command.
(treesit-major-mode-setup): Conditionally set forward-sexp-function.
* lisp/progmodes/ruby-ts-mode.el: Add some types to ruby-ts-mode.
---
 lisp/progmodes/java-ts-mode.el | 15 +++++++++++++++
 lisp/progmodes/ruby-ts-mode.el |  4 ++++
 lisp/treesit.el                | 17 +++++++++++++++++
 3 files changed, 36 insertions(+)

diff --git a/lisp/progmodes/java-ts-mode.el b/lisp/progmodes/java-ts-mode.el
index 83c437d307..03093e0980 100644
--- a/lisp/progmodes/java-ts-mode.el
+++ b/lisp/progmodes/java-ts-mode.el
@@ -328,6 +328,21 @@ java-ts-mode
                             "package_declaration"
                             "import_declaration")))
 
+  (setq-local treesit-sexp-type-regexp
+              (regexp-opt '("annotation"
+                            "parenthesized_expression"
+                            "argument_list"
+                            "identifier"
+                            "modifiers"
+                            "block"
+                            "body"
+                            "literal"
+                            "access"
+                            "reference"
+                            "_type"
+                            "true"
+                            "false")))
+
   ;; Font-lock.
   (setq-local treesit-font-lock-settings java-ts-mode--font-lock-settings)
   (setq-local treesit-font-lock-feature-list
diff --git a/lisp/progmodes/ruby-ts-mode.el b/lisp/progmodes/ruby-ts-mode.el
index 939c054b04..78186c78d5 100644
--- a/lisp/progmodes/ruby-ts-mode.el
+++ b/lisp/progmodes/ruby-ts-mode.el
@@ -984,6 +984,10 @@ ruby-ts-mode
   ;; Navigation.
   (setq-local treesit-defun-type-regexp ruby-ts--method-regex)
 
+  (setq-local treesit-sexp-type-regexp
+              (regexp-opt '("method"
+                            "if")))
+
   ;; AFAIK, Ruby can not nest methods
   (setq-local treesit-defun-prefer-top-level nil)
 
diff --git a/lisp/treesit.el b/lisp/treesit.el
index 69bfff21df..735df7a8f4 100644
--- a/lisp/treesit.el
+++ b/lisp/treesit.el
@@ -1622,6 +1622,21 @@ treesit-search-forward-goto
       (goto-char current-pos)))
     node))
 
+(defvar-local treesit-sexp-type-regexp nil
+  "A regexp that matches the node type of sexp nodes.
+
+A sexp node is a node that is bigger than punctuation, and
+delimits medium sized statements in the source code.  It is,
+however, smaller in scope than sentences.  This is used by
+`treesit-forward-sexp' and friends.")
+
+(defun treesit-forward-sexp (&optional arg)
+  (interactive "^p")
+  (or arg (setq arg 1))
+  (funcall
+   (if (> arg 0) #'treesit-end-of-thing #'treesit-beginning-of-thing)
+   treesit-sexp-type-regexp (abs arg)))
+
 (defun treesit-transpose-sexps (&optional arg)
   "Tree-sitter `transpose-sexps' function.
 Arg is the same as in `transpose-sexps'.
@@ -2287,6 +2302,8 @@ treesit-major-mode-setup
     (setq-local add-log-current-defun-function
                 #'treesit-add-log-current-defun))
 
+  (when treesit-sexp-type-regexp
+    (setq-local forward-sexp-function #'treesit-forward-sexp))
   (setq-local transpose-sexps-function #'treesit-transpose-sexps)
   (when treesit-sentence-type-regexp
     (setq-local forward-sentence-function #'treesit-forward-sentence))
-- 
2.34.1