Subject: [PATCH] Simplify imenu setup for {cmake,dockerfile}-ts-modes

From: Randy Taylor <[email-masked]>
Date: Thu, 25 Jan 2024 03:30:11 +0000

X-Debbugs-CC: [email-masked]

This patch also fixes a bug in cmake-ts-mode's imenu creation, where function names would include their arguments.

For example:

function(name arg1 arg2)
endfunction()

This function would appear as "name arg1 arg2".

From cfc6cea70317ed1967e68f3ac4439382313be38b Mon Sep 17 00:00:00 2001
From: Randy Taylor <[email-masked]>
Date: Wed, 24 Jan 2024 21:39:45 -0500
Subject: [PATCH] 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.
---
 lisp/progmodes/cmake-ts-mode.el      | 46 +++++++--------------------
 lisp/progmodes/dockerfile-ts-mode.el | 47 +++++++---------------------
 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 @@ cmake-ts-mode--font-lock-settings
    '((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 @@ cmake-ts-mode
     (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 334f3064d98..f6587018513 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 @@ dockerfile-ts-mode--font-lock-settings
    '((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 @@ dockerfile-ts-mode
     (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.
-- 
2.43.0