Partial support for DEFUN in c-ts-mode (bug#62825)

The DEFUN is hard to incorporate because it's made of two nodes rather
than one, and most tree-sitter functionalities assume a defun is one
node.  I fixed the indent-defun and add-log functionality, but
beginning/end-of-defun and imenu still don't recognize DEFUN.

* lisp/progmodes/c-ts-mode.el (c-ts-mode-emacs-devel): New variable.
(c-ts-mode--defun-name): Support DEFUN.
(c-ts-mode--defun-valid-p): Support DEFUN.
(c-ts-mode--emacs-defun-p)
(c-ts-mode--emacs-defun-at-point): New functions.
(c-ts-mode-indent-defun): Use c-ts-mode--emacs-defun-at-point.
(c-ts-mode--emacs-current-defun-name): New function.
(c-ts-mode, c++-ts-mode): Optionally setup custom defun-name function.

1 files changed, 101 insertions, 25 deletions

diff --git a/lisp/progmodes/c-ts-mode.el b/lisp/progmodes/c-ts-mode.el
index 2d54577cb09..f3c1dde10c0 100644
--- a/lisp/progmodes/c-ts-mode.el
+++ b/lisp/progmodes/c-ts-mode.el
@@ -192,6 +192,10 @@ To set the default indent style globally, use
            (c-ts-mode--get-indent-style
             (if (derived-mode-p 'c-ts-mode) 'c 'cpp))))))
 
+(defvar c-ts-mode-emacs-devel nil
+  "If the value is t, enable Emacs source-specific features.
+This needs to be set before enabling `c-ts-mode'.")
+
 ;;; Syntax table
 
 (defvar c-ts-mode--syntax-table
@@ -802,7 +806,14 @@ Return nil if NODE is not a defun node or doesn't have a name."
      ((or "struct_specifier" "enum_specifier"
           "union_specifier" "class_specifier"
           "namespace_definition")
-      (treesit-node-child-by-field-name node "name")))
+      (treesit-node-child-by-field-name node "name"))
+     ;; DEFUNs in Emacs source.
+     ("expression_statement"
+      (let* ((call-exp-1 (treesit-node-child node 0))
+             (call-exp-2 (treesit-node-child call-exp-1 0))
+             (arg-list (treesit-node-child call-exp-2 1))
+             (name (treesit-node-child arg-list 1 t)))
+        name)))
    t))
 
 ;;; Defun navigation
@@ -810,28 +821,29 @@ Return nil if NODE is not a defun node or doesn't have a name."
 (defun c-ts-mode--defun-valid-p (node)
   "Return non-nil if NODE is a valid defun node.
 Ie, NODE is not nested."
-  (not (or (and (member (treesit-node-type node)
-                        '("struct_specifier"
-                          "enum_specifier"
-                          "union_specifier"
-                          "declaration"))
-                ;; If NODE's type is one of the above, make sure it is
-                ;; top-level.
-                (treesit-node-top-level
-                 node (rx (or "function_definition"
-                              "type_definition"
-                              "struct_specifier"
+  (or (c-ts-mode--emacs-defun-p node)
+      (not (or (and (member (treesit-node-type node)
+                            '("struct_specifier"
                               "enum_specifier"
                               "union_specifier"
-                              "declaration"))))
+                              "declaration"))
+                    ;; If NODE's type is one of the above, make sure it is
+                    ;; top-level.
+                    (treesit-node-top-level
+                     node (rx (or "function_definition"
+                                  "type_definition"
+                                  "struct_specifier"
+                                  "enum_specifier"
+                                  "union_specifier"
+                                  "declaration"))))
 
-           (and (equal (treesit-node-type node) "declaration")
-                ;; If NODE is a declaration, make sure it is not a
-                ;; function declaration.
-                (equal (treesit-node-type
-                        (treesit-node-child-by-field-name
-                         node "declarator"))
-                       "function_declarator")))))
+               (and (equal (treesit-node-type node) "declaration")
+                    ;; If NODE is a declaration, make sure it is not a
+                    ;; function declaration.
+                    (equal (treesit-node-type
+                            (treesit-node-child-by-field-name
+                             node "declarator"))
+                           "function_declarator"))))))
 
 (defun c-ts-mode--defun-for-class-in-imenu-p (node)
   "Check if NODE is a valid entry for the Class subindex.
@@ -871,17 +883,73 @@ the semicolon.  This function skips the semicolon."
           (setq parent (treesit-node-parent smallest-node)))))
     (list node parent bol)))
 
+(defun c-ts-mode--emacs-defun-p (node)
+  "Return non-nil if NODE is a DEFUN in Emacs source files."
+  (and (equal (treesit-node-type node) "expression_statement")
+       (equal (treesit-node-text
+               (treesit-node-child-by-field-name
+                (treesit-node-child
+                 (treesit-node-child node 0) 0)
+                "function")
+               t)
+              "DEFUN")))
+
+(defun c-ts-mode--emacs-defun-at-point (&optional range)
+  "Return the current defun node.
+
+This function recognizes DEFUNs in Emacs source files.
+
+Note that for the case of a DEFUN, it is made of two separate
+nodes, one for the declaration and one for the body, this
+function returns the declaration node.
+
+If RANGE is non-nil, return (BEG . END) where BEG end END
+encloses the whole defun.  This solves the problem of only
+returning the declaration part for DEFUN."
+  (or (when-let ((node (treesit-defun-at-point)))
+        (if range
+            (cons (treesit-node-start node)
+                (treesit-node-end node))
+            node))
+      (and c-ts-mode-emacs-devel
+           (let ((candidate-1 ; For when point is in the DEFUN statement.
+                  (treesit-node-prev-sibling
+                   (treesit-node-top-level
+                    (treesit-node-at (point))
+                    "compound_statement")))
+                 (candidate-2 ; For when point is in the body.
+                  (treesit-node-top-level
+                   (treesit-node-at (point))
+                   "expression_statement")))
+             (when-let
+                 ((node (or (and (c-ts-mode--emacs-defun-p candidate-1)
+                                 candidate-1)
+                            (and (c-ts-mode--emacs-defun-p candidate-2)
+                                 candidate-2))))
+               (if range
+                   (cons (treesit-node-start node)
+                       (treesit-node-end
+                        (treesit-node-next-sibling node)))
+                   node))))))
+
 (defun c-ts-mode-indent-defun ()
   "Indent the current top-level declaration syntactically.
 
 `treesit-defun-type-regexp' defines what constructs to indent."
   (interactive "*")
   (when-let ((orig-point (point-marker))
-             (node (treesit-defun-at-point)))
-    (indent-region (treesit-node-start node)
-                   (treesit-node-end node))
+             (range (c-ts-mode--emacs-defun-at-point t)))
+    (indent-region (car range) (cdr range))
     (goto-char orig-point)))
 
+(defun c-ts-mode--emacs-current-defun-name ()
+  "Return the name of the current defun.
+This is used for `add-log-current-defun-function'.  This
+recognizes DEFUN in Emacs sources, in addition to normal function
+definitions."
+  (or (treesit-add-log-current-defun)
+      (c-ts-mode--defun-name (c-ts-mode--emacs-defun-at-point))))
+
 ;;; Modes
 
 (defvar-keymap c-ts-base-mode-map
@@ -996,7 +1064,11 @@ in your configuration."
     (setq-local treesit-font-lock-settings (c-ts-mode--font-lock-settings 'c))
     ;; Navigation.
     (setq-local treesit-defun-tactic 'top-level)
-    (treesit-major-mode-setup)))
+    (treesit-major-mode-setup)
+
+    (when c-ts-mode-emacs-devel
+      (setq-local add-log-current-defun-function
+                  #'c-ts-mode--emacs-current-defun-name))))
 
 ;;;###autoload
 (define-derived-mode c++-ts-mode c-ts-base-mode "C++"
@@ -1031,7 +1103,11 @@ recommended to enable `electric-pair-mode' with this mode."
                 (c-ts-mode--get-indent-style 'cpp))
     ;; Font-lock.
     (setq-local treesit-font-lock-settings (c-ts-mode--font-lock-settings 'cpp))
-    (treesit-major-mode-setup)))
+    (treesit-major-mode-setup)
+
+    (when c-ts-mode-emacs-devel
+      (setq-local add-log-current-defun-function
+                  #'c-ts-mode--emacs-current-defun-name))))
 
 (easy-menu-define c-ts-mode-menu (list c-ts-mode-map c++-ts-mode-map)
   "Menu for `c-ts-mode' and `c++-ts-mode'."