Add treesit-aggregated-simple-imenu-settings

Now we support setting up Imenu for multiple languages

* doc/lispref/modes.texi: Update manual.
* lisp/treesit.el:
(treesit-aggregated-simple-imenu-settings): New variable.
(treesit--imenu-merge-entries): New function.
(treesit--generate-simple-imenu): This was previously
treesit-simple-imenu.
(treesit-simple-imenu): Support
treesit-aggregated-simple-imenu-settings.
(treesit-major-mode-setup): Recognize
treesit-aggregated-simple-imenu-settings.
* test/src/treesit-tests.el (treesit-imenu): New test.

3 files changed, 113 insertions, 14 deletions

diff --git a/doc/lispref/modes.texi b/doc/lispref/modes.texi
index 73edb688c85..f227bdc635f 100644
--- a/doc/lispref/modes.texi
+++ b/doc/lispref/modes.texi
@@ -3109,6 +3109,15 @@ instead.
 automatically sets up Imenu if this variable is non-@code{nil}.
 @end defvar
 
+@defvar treesit-aggregated-simple-imenu-settings
+This variable allows major modes to configure Imenu for multiple
+languages.  Its value is an alist mapping language symbols to Imenu
+settings described in @var{treesit-simple-imenu-settings}.
+
+If both this variable and @var{treesit-simple-imenu-settings} is
+non-@code{nil}, Emacs uses this variable for setting up Imenu.
+@end defvar
+
 @node Outline Minor Mode
 @section Outline Minor Mode
 
diff --git a/lisp/treesit.el b/lisp/treesit.el
index 2cf7bccdeed..464b7e688be 100644
--- a/lisp/treesit.el
+++ b/lisp/treesit.el
@@ -3123,6 +3123,31 @@ node and returns the name of that defun node.  If NAME-FN is nil,
 `treesit-major-mode-setup' automatically sets up Imenu if this
 variable is non-nil.")
 
+;; `treesit-simple-imenu-settings' doesn't support multiple languages,
+;; and we need to add multi-lang support for Imenu.  One option is to
+;; extend treesit-simple-imenu-settings to specify language, either by
+;; making it optionally an alist (just like
+;; `treesit-aggregated-simple-imenu-settings'), or add a fifth element
+;; to each setting.  But either way makes borrowing Imenu settings from
+;; other modes difficult: with the alist approach, you'd need to check
+;; whether other mode uses a plain list or an alist; with the fifth
+;; element approach, again, you need to check if each setting has the
+;; fifth element, and add it if not.
+;;
+;; OTOH, with `treesit-aggregated-simple-imenu-settings', borrowing
+;; Imenu settings is easy: if `treesit-aggregated-simple-imenu-settings'
+;; is non-nil, copy everything over; if `treesit-simple-imenu-settings'
+;; is non-nil, copy the settings and put them under a language symbol.
+(defvar treesit-aggregated-simple-imenu-settings nil
+  "Settings that configure `treesit-simple-imenu' for multi-language modes.
+
+The value should be an alist of (LANG . SETTINGS), where LANG is a
+language symbol, and SETTINGS has the same form as
+`treesit-simple-imenu-settings'.
+
+When both this variable and `treesit-simple-imenu-settings' are non-nil,
+this variable takes priority.")
+
 (defun treesit--simple-imenu-1 (node pred name-fn)
   "Given a sparse tree, create an Imenu index.
 
@@ -3170,20 +3195,69 @@ ENTRY.  MARKER marks the start of each tree-sitter node."
      ;; Leaf node, return a (list of) plain index entry.
      (t (list (cons name marker))))))
 
+(defun treesit--imenu-merge-entries (entries)
+  "Merge ENTRIES by category.
+
+ENTRIES is a list of (CATEGORY . SUB-ENTRIES...).  Merge them so there's
+no duplicate CATEGORY.  CATEGORY's are strings.  The merge is stable,
+meaning the order of elements are kept."
+  (let ((return-entries nil))
+    (dolist (entry entries)
+      (let* ((category (car entry))
+             (sub-entries (cdr entry))
+             (existing-entries
+              (alist-get category return-entries nil nil #'equal)))
+        (if (not existing-entries)
+            (push entry return-entries)
+          (setf (alist-get category return-entries nil nil #'equal)
+                (append existing-entries sub-entries)))))
+    (nreverse return-entries)))
+
+(defun treesit--generate-simple-imenu (node settings)
+  "Return an Imenu index for NODE with SETTINGS.
+
+NODE usually should be a root node of a parser.  SETTINGS is described
+by `treesit-simple-imenu-settings'."
+  (mapcan (lambda (setting)
+            (pcase-let ((`(,category ,regexp ,pred ,name-fn)
+                         setting))
+              (when-let* ((tree (treesit-induce-sparse-tree
+                                 node regexp))
+                          (index (treesit--simple-imenu-1
+                                  tree pred name-fn)))
+                (if category
+                    (list (cons category index))
+                  index))))
+          settings))
+
 (defun treesit-simple-imenu ()
   "Return an Imenu index for the current buffer."
-  (let ((root (treesit-buffer-root-node)))
-    (mapcan (lambda (setting)
-              (pcase-let ((`(,category ,regexp ,pred ,name-fn)
-                           setting))
-                (when-let* ((tree (treesit-induce-sparse-tree
-                                   root regexp))
-                            (index (treesit--simple-imenu-1
-                                    tree pred name-fn)))
-                  (if category
-                      (list (cons category index))
-                    index))))
-            treesit-simple-imenu-settings)))
+  (if (not treesit-aggregated-simple-imenu-settings)
+      (treesit--generate-simple-imenu
+       (treesit-parser-root-node treesit-primary-parser)
+       treesit-simple-imenu-settings)
+    ;; Use `treesit-aggregated-simple-imenu-settings'.  Remove languages
+    ;; that doesn't have any Imenu entries.
+    (seq-filter
+     #'cdr
+     (mapcar
+      (lambda (entry)
+        (let* ((lang (car entry))
+               (settings (cdr entry))
+               (global-parser (car (treesit-parser-list nil lang)))
+               (local-parsers
+                (treesit-parser-list nil lang 'embedded)))
+          (cons (treesit-language-display-name lang)
+                ;; No one says you can't have both global and local
+                ;; parsers for the same language.  E.g., Rust uses
+                ;; local parsers for the same language to handle
+                ;; macros.
+                (treesit--imenu-merge-entries
+                 (mapcan (lambda (parser)
+                           (treesit--generate-simple-imenu
+                            (treesit-parser-root-node parser) settings))
+                         (cons global-parser local-parsers))))))
+      treesit-aggregated-simple-imenu-settings))))
 
 ;;; Outline minor mode
 
@@ -3321,7 +3395,8 @@ and `end-of-defun-function'.
 If `treesit-defun-name-function' is non-nil, set up
 `add-log-current-defun'.
 
-If `treesit-simple-imenu-settings' is non-nil, set up Imenu.
+If `treesit-simple-imenu-settings' or
+`treesit-aggregated-simple-imenu-settings' is non-nil, set up Imenu.
 
 If either `treesit-outline-predicate' or `treesit-simple-imenu-settings'
 are non-nil, and Outline minor mode settings don't already exist, setup
@@ -3395,7 +3470,8 @@ before calling this function."
     (setq-local forward-sentence-function #'treesit-forward-sentence))
 
   ;; Imenu.
-  (when treesit-simple-imenu-settings
+  (when (or treesit-aggregated-simple-imenu-settings
+            treesit-simple-imenu-settings)
     (setq-local imenu-create-index-function
                 #'treesit-simple-imenu))
 
diff --git a/test/src/treesit-tests.el b/test/src/treesit-tests.el
index 50f205421d7..43102fc97e0 100644
--- a/test/src/treesit-tests.el
+++ b/test/src/treesit-tests.el
@@ -1270,6 +1270,20 @@ This tests bug#60355."
     (should node)
     (should (equal (treesit-node-text node) "2"))))
 
+;;; Imenu
+
+(ert-deftest treesit-imenu ()
+  "Test imenu functions."
+  (should (equal (treesit--imenu-merge-entries
+                  '(("Function" . (f1 f2))
+                    ("Function" . (f3 f4 f5))
+                    ("Class" . (c1 c2 c3))
+                    ("Variables" . (v1 v2))
+                    ("Class" . (c4))))
+                 '(("Function" . (f1 f2 f3 f4 f5))
+                   ("Class" . (c1 c2 c3 c4))
+                   ("Variables" . (v1 v2))))))
+
 
 ;; TODO
 ;; - Functions in treesit.el