feature/tree-sitter 908aab6144 2/2: Improve tree-sitter imenu for c-mode and js-mode

[Yuan Fu]

branch: feature/tree-sitter
commit 908aab6144bf9f62dd5b2c2f1776ed177fa6d18f
Author: Yuan Fu <casouri@gmail.com>
Commit: Yuan Fu <casouri@gmail.com>

    Improve tree-sitter imenu for c-mode and js-mode
    
    Instead of a flat list, now categorize imenu entries into categories
    like "Function", "Variable", "Class", etc.
    
    * lisp/progmodes/c-ts-mode.el (c-ts-mode--imenu-1): Handle more edge
    cases.
    (c-ts-mode--imenu): Categorize menu entries.
    * lisp/progmodes/js.el (js--treesit-imenu-type-alist): Remove variable.
    (js--treesit-imenu-label): Remove function.
    (js--treesit-imenu-1): Use the name alone for labels.
    (js--treesit-imenu): Categorize menu entries.
---
 lisp/progmodes/c-ts-mode.el | 42 ++++++++++++++++++++++++++++++++----------
 lisp/progmodes/js.el        | 44 ++++++++++++++------------------------------
 2 files changed, 46 insertions(+), 40 deletions(-)

diff --git a/lisp/progmodes/c-ts-mode.el b/lisp/progmodes/c-ts-mode.el
index 09241b7ab9..5617ea7d7c 100644
--- a/lisp/progmodes/c-ts-mode.el
+++ b/lisp/progmodes/c-ts-mode.el
@@ -340,23 +340,40 @@ the subtrees."
          (subtrees (mapcan #'c-ts-mode--imenu-1 (cdr node)))
          (name (when ts-node
                  (or (treesit-node-text
-                      (or (treesit-node-child-by-field-name
-                           ts-node "declarator")
+                      (pcase (treesit-node-type ts-node)
+                        ("function_definition"
+                         (treesit-node-child-by-field-name
                           (treesit-node-child-by-field-name
-                           ts-node "name"))
+                           ts-node "declarator")
+                          "declarator"))
+                        ("declaration"
+                         (let ((child (treesit-node-child ts-node -1 t)))
+                           (pcase (treesit-node-type child)
+                             ("identifier" child)
+                             (_ (treesit-node-child-by-field-name
+                                 child "declarator")))))
+                        ("struct_specifier"
+                         (treesit-node-child-by-field-name
+                          ts-node "name")))
                       t)
                      "Unnamed node")))
          (marker (when ts-node
                    (set-marker (make-marker)
                                (treesit-node-start ts-node)))))
-    ;; A struct_specifier could be inside a parameter list or another
-    ;; struct definition.  In those cases we don't include it.
+    ;; A struct_specifier could be inside a parameter list, another
+    ;; struct definition, a variable declaration, a function
+    ;; declaration.  In those cases we don't include it.
     (cond
      ((string-match-p
-       (rx (or "parameter" "field") "_declaration")
+       (rx (or "parameter_declaration" "field_declaration"
+               "declaration" "function_definition"))
        (or (treesit-node-type (treesit-node-parent ts-node))
            ""))
       nil)
+     ((and (equal (treesit-node-type ts-node) "declaration")
+           (not (equal (treesit-node-type (treesit-node-parent ts-node))
+                       "translation_unit")))
+      nil)
      ((null ts-node) subtrees)
      (subtrees
       `((,name ,(cons name marker) ,@subtrees)))
@@ -366,10 +383,15 @@ the subtrees."
 (defun c-ts-mode--imenu ()
   "Return Imenu alist for the current buffer."
   (let* ((node (treesit-buffer-root-node))
-         (tree (treesit-induce-sparse-tree
-                node (rx (or "function_definition"
-                             "struct_specifier")))))
-    (c-ts-mode--imenu-1 tree)))
+         (func-tree (treesit-induce-sparse-tree
+                     node "^function_definition$"))
+         (var-tree (treesit-induce-sparse-tree
+                    node "^declaration$"))
+         (struct-tree (treesit-induce-sparse-tree
+                       node "^struct_specifier$")))
+    `(("Struct" . ,(c-ts-mode--imenu-1 struct-tree))
+      ("Variable" . ,(c-ts-mode--imenu-1 var-tree))
+      ("Function" . ,(c-ts-mode--imenu-1 func-tree)))))
 
 ;;;###autoload
 (define-derived-mode c-ts-mode--base-mode prog-mode "C"
diff --git a/lisp/progmodes/js.el b/lisp/progmodes/js.el
index 8ec7e7c675..ac2b6a27b7 100644
--- a/lisp/progmodes/js.el
+++ b/lisp/progmodes/js.el
@@ -3631,23 +3631,6 @@ This function can be used as a value in 
`which-func-functions'"
              do (setq node (treesit-node-parent node))
              finally return  (string-join name-list "."))))
 
-;; Keep this private since we might later change it or generalize it.
-(defvar js--treesit-imenu-type-alist
-  '((variable . "V ")
-    (function . "F ")
-    (class . "C ")
-    (method . "M "))
-  "Maps imenu label types to their \"symbol\".
-Symbols are prefixed to each label in imenu (see
-`js--treesit-imenu-label').")
-
-(defun js--treesit-imenu-label (type name)
-  "Format label for imenu.
-TYPE can be `variable', `function', `class', `method'.
-NAME is a string."
-  (format "%s%s" (alist-get type js--treesit-imenu-type-alist)
-          name))
-
 (defun js--treesit-imenu-1 (node)
   "Given a sparse tree, create an imenu alist.
 
@@ -3695,23 +3678,24 @@ definition*\"."
            (treesit-node-top-level ts-node))
       nil)
      (subtrees
-      (let ((parent-label (js--treesit-imenu-label type name))
-            (jump-label ""))
-        `((,parent-label
-           ,(cons jump-label marker)
-           ,@subtrees))))
-     (t (let ((label (js--treesit-imenu-label type name)))
-          (list (cons label marker)))))))
+      `((,name
+         ,(cons "" marker)
+         ,@subtrees)))
+     (t (list (cons name marker))))))
 
 (defun js--treesit-imenu ()
   "Return Imenu alist for the current buffer."
   (let* ((node (treesit-buffer-root-node))
-         (tree (treesit-induce-sparse-tree
-                node (rx (or "class_declaration"
-                             "method_definition"
-                             "function_declaration"
-                             "lexical_declaration")))))
-    (js--treesit-imenu-1 tree)))
+         (class-tree (treesit-induce-sparse-tree
+                      node (rx (or "class_declaration"
+                                   "method_definition"))))
+         (func-tree (treesit-induce-sparse-tree
+                     node "function_declaration"))
+         (var-tree (treesit-induce-sparse-tree
+                    node "lexical_declaration")))
+    `(("Class" . ,(js--treesit-imenu-1 class-tree))
+      ("Varieable" . ,(js--treesit-imenu-1 var-tree))
+      ("Function" . ,(js--treesit-imenu-1 func-tree)))))
 
 ;;; Main Function