Simplify imenu setup for {cmake,dockerfile}-ts-modes

* lisp/progmodes/cmake-ts-mode.el (treesit-induce-sparse-tree,
treesit-node-child, treesit-node-start, cmake-ts-mode--imenu,
cmake-ts-mode--imenu-1): Remove.
(treesit-search-subtree): Declare.
(cmake-ts-mode--function-name): New function.
(cmake-ts-mode): Use it.

* lisp/progmodes/dockerfile-ts-mode.el (treesit-induce-sparse-tree,
treesit-node-start, dockerfile-ts-mode--imenu,
dockerfile-ts-mode--imenu-1): Remove.
(dockerfile-ts-mode--stage-name): New function.
(dockerfile-ts-mode): Use it.

2 files changed, 22 insertions, 71 deletions

diff --git a/lisp/progmodes/cmake-ts-mode.el b/lisp/progmodes/cmake-ts-mode.el
index d933e4ebb81..29c9e957d3c 100644
--- a/lisp/progmodes/cmake-ts-mode.el
+++ b/lisp/progmodes/cmake-ts-mode.el
@@ -32,10 +32,8 @@
 
 (declare-function treesit-parser-create "treesit.c")
 (declare-function treesit-query-capture "treesit.c")
-(declare-function treesit-induce-sparse-tree "treesit.c")
-(declare-function treesit-node-child "treesit.c")
-(declare-function treesit-node-start "treesit.c")
 (declare-function treesit-node-type "treesit.c")
+(declare-function treesit-search-subtree "treesit.c")
 
 (defcustom cmake-ts-mode-indent-offset 2
   "Number of spaces for each indentation step in `cmake-ts-mode'."
@@ -195,37 +193,14 @@ Check if a node type is available, then return the right font lock rules."
    '((ERROR) @font-lock-warning-face))
   "Tree-sitter font-lock settings for `cmake-ts-mode'.")
 
-(defun cmake-ts-mode--imenu ()
-  "Return Imenu alist for the current buffer."
-  (let* ((node (treesit-buffer-root-node))
-         (func-tree (treesit-induce-sparse-tree
-                     node "function_def" nil 1000))
-         (func-index (cmake-ts-mode--imenu-1 func-tree)))
-    (append
-     (when func-index `(("Function" . ,func-index))))))
-
-(defun cmake-ts-mode--imenu-1 (node)
-  "Helper for `cmake-ts-mode--imenu'.
-Find string representation for NODE and set marker, then recurse
-the subtrees."
-  (let* ((ts-node (car node))
-         (children (cdr node))
-         (subtrees (mapcan #'cmake-ts-mode--imenu-1
-                           children))
-         (name (when ts-node
-                 (pcase (treesit-node-type ts-node)
-                   ("function_def"
-                    (treesit-node-text
-                     (treesit-node-child (treesit-node-child ts-node 0) 2) t)))))
-         (marker (when ts-node
-                   (set-marker (make-marker)
-                               (treesit-node-start ts-node)))))
-    (cond
-     ((or (null ts-node) (null name)) subtrees)
-     (subtrees
-      `((,name ,(cons name marker) ,@subtrees)))
-     (t
-      `((,name . ,marker))))))
+(defun cmake-ts-mode--function-name (node)
+  "Return the function name of NODE.
+Return nil if there is no name or if NODE is not a function node."
+  (pcase (treesit-node-type node)
+    ("function_command"
+     (treesit-node-text
+      (treesit-search-subtree node "^argument$" nil nil 2)
+      t))))
 
 ;;;###autoload
 (define-derived-mode cmake-ts-mode prog-mode "CMake"
@@ -242,7 +217,8 @@ the subtrees."
     (setq-local comment-start-skip (rx "#" (* (syntax whitespace))))
 
     ;; Imenu.
-    (setq-local imenu-create-index-function #'cmake-ts-mode--imenu)
+    (setq-local treesit-simple-imenu-settings
+                `(("Function" "\\`function_command\\'" nil cmake-ts-mode--function-name)))
     (setq-local which-func-functions nil)
 
     ;; Indent.
diff --git a/lisp/progmodes/dockerfile-ts-mode.el b/lisp/progmodes/dockerfile-ts-mode.el
index d5b7f953e31..878335431af 100644
--- a/lisp/progmodes/dockerfile-ts-mode.el
+++ b/lisp/progmodes/dockerfile-ts-mode.el
@@ -31,10 +31,8 @@
 (eval-when-compile (require 'rx))
 
 (declare-function treesit-parser-create "treesit.c")
-(declare-function treesit-induce-sparse-tree "treesit.c")
 (declare-function treesit-node-child "treesit.c")
 (declare-function treesit-node-child-by-field-name "treesit.c")
-(declare-function treesit-node-start "treesit.c")
 (declare-function treesit-node-type "treesit.c")
 
 (defvar dockerfile-ts-mode--syntax-table
@@ -118,38 +116,15 @@ continuation to the previous entry."
    '((ERROR) @font-lock-warning-face))
   "Tree-sitter font-lock settings.")
 
-(defun dockerfile-ts-mode--imenu ()
-  "Return Imenu alist for the current buffer."
-  (let* ((node (treesit-buffer-root-node))
-         (stage-tree (treesit-induce-sparse-tree
-                      node "from_instruction"
-                      nil 1000))
-         (stage-index (dockerfile-ts-mode--imenu-1 stage-tree)))
-    (when stage-index `(("Stage" . ,stage-index)))))
-
-(defun dockerfile-ts-mode--imenu-1 (node)
-  "Helper for `dockerfile-ts-mode--imenu'.
-Find string representation for NODE and set marker, then recurse
-the subtrees."
-  (let* ((ts-node (car node))
-         (children (cdr node))
-         (subtrees (mapcan #'dockerfile-ts-mode--imenu-1
-                           children))
-         (name (when ts-node
-                 (pcase (treesit-node-type ts-node)
-                   ("from_instruction"
-                    (treesit-node-text
-                     (or (treesit-node-child-by-field-name ts-node "as")
-                         (treesit-node-child ts-node 1)) t)))))
-         (marker (when ts-node
-                   (set-marker (make-marker)
-                               (treesit-node-start ts-node)))))
-    (cond
-     ((or (null ts-node) (null name)) subtrees)
-     (subtrees
-      `((,name ,(cons name marker) ,@subtrees)))
-     (t
-      `((,name . ,marker))))))
+(defun dockerfile-ts-mode--stage-name (node)
+  "Return the stage name of NODE.
+Return nil if there is no name or if NODE is not a stage node."
+  (pcase (treesit-node-type node)
+    ("from_instruction"
+     (treesit-node-text
+      (or (treesit-node-child-by-field-name node "as")
+          (treesit-node-child node 1))
+      t))))
 
 ;;;###autoload
 (define-derived-mode dockerfile-ts-mode prog-mode "Dockerfile"
@@ -166,8 +141,8 @@ the subtrees."
     (setq-local comment-start-skip (rx "#" (* (syntax whitespace))))
 
     ;; Imenu.
-    (setq-local imenu-create-index-function
-                #'dockerfile-ts-mode--imenu)
+    (setq-local treesit-simple-imenu-settings
+                `(("Stage" "\\`from_instruction\\'" nil dockerfile-ts-mode--stage-name)))
     (setq-local which-func-functions nil)
 
     ;; Indent.