feature/tree-sitter 6fb6cb075f 06/10: Add tree-sitter imenu support for js-mode and ts-mode

[Yuan Fu]

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

    Add tree-sitter imenu support for js-mode and ts-mode
    
    js-mode's current imenu is pretty plain and incomplete, so I took the
    liberty to add a bit more flair to it.
    
    * lisp/progmodes/js.el (js--treesit-imenu-type-alist): New variable.
    (js--treesit-imenu-top-level-p)
    (js--treesit-imenu-label)
    (js--treesit-imenu-1)
    (js--treesit-imenu): New functions.
---
 lisp/progmodes/js.el      | 99 +++++++++++++++++++++++++++++++++++++++++++++++
 lisp/progmodes/ts-mode.el |  4 ++
 2 files changed, 103 insertions(+)

diff --git a/lisp/progmodes/js.el b/lisp/progmodes/js.el
index 1d53a624a7..8802ea0b5f 100644
--- a/lisp/progmodes/js.el
+++ b/lisp/progmodes/js.el
@@ -3601,6 +3601,99 @@ 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 "."))))
 
+(defun js--treesit-imenu-top-level-p (node)
+  "Return t if NODE is top-level, nil otherwise.
+Being top-level means there is no parent of NODE that has the
+same type."
+  (when node
+    (catch 'term
+      (let ((type (treesit-node-type node)))
+        (while (setq node (treesit-node-parent node))
+          (when (equal (treesit-node-type node) type)
+            (throw 'term nil))))
+      t)))
+
+;; 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.
+
+NODE is the root node of the tree returned by
+`treesit-induce-sparse-tree' (not a tree-sitter node, its car is
+a tree-sitter node).  Walk that tree and return an imenu alist.
+
+Return a list of ENTRY where
+
+ENTRY := (NAME . MARKER)
+       | (NAME . ((JUMP-LABEL . MARKER)
+                  ENTRY
+                  ...)
+
+NAME is the function/class's name, JUMP-LABEL is like \"*function
+definition*\"."
+  (let* ((ts-node (car node))
+         (children (cdr node))
+         (subtrees (mapcan #'js--treesit-imenu-1
+                           children))
+         (type (pcase (treesit-node-type ts-node)
+                 ("lexical_declaration" 'variable)
+                 ("class_declaration" 'class)
+                 ("method_definition" 'method)
+                 ("function_declaration" 'function)))
+         ;; The root of the tree could have a nil ts-node.
+         (name (when ts-node
+                 (let ((ts-node-1
+                        (if (eq type 'variable)
+                            (treesit-search-subtree
+                             ts-node "variable_declarator" nil nil 1)
+                          ts-node)))
+                   (treesit-node-text
+                    (treesit-node-child-by-field-name
+                     ts-node-1 "name")
+                    t))))
+         (marker (when ts-node
+                   (set-marker (make-marker)
+                               (treesit-node-start ts-node)))))
+    (cond
+     ((null ts-node)
+      subtrees)
+     ((and (eq type 'variable)
+           (not (js--treesit-imenu-top-level-p 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)))))))
+
+(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)))
+
 ;;; Main Function
 
 ;;;###autoload
@@ -3679,6 +3772,7 @@ This function can be used as a value in 
`which-func-functions'"
   (cond
    ;; Tree-sitter.
    ((treesit-ready-p 'js-mode 'javascript)
+    (treesit-parser-create 'javascript)
     ;; Indent.
     (setq-local treesit-simple-indent-rules js--treesit-indent-rules)
     ;; Navigation.
@@ -3690,6 +3784,11 @@ This function can be used as a value in 
`which-func-functions'"
     ;; Fontification.
     (setq-local treesit-font-lock-settings js--treesit-font-lock-settings)
     (setq-local treesit-font-lock-feature-list '((minimal) (moderate) (full)))
+    ;; Imenu
+    (setq-local imenu-create-index-function
+                #'js--treesit-imenu)
+    ;; Which-func (use imenu).
+    (setq-local which-func-functions nil)
     (treesit-major-mode-setup))
    ;; Elisp.
    (t
diff --git a/lisp/progmodes/ts-mode.el b/lisp/progmodes/ts-mode.el
index f2002f9ef7..28800e378a 100644
--- a/lisp/progmodes/ts-mode.el
+++ b/lisp/progmodes/ts-mode.el
@@ -278,6 +278,10 @@
     ;; Font-lock.
     (setq-local treesit-font-lock-settings ts-mode--font-lock-settings)
     (setq-local treesit-font-lock-feature-list '((minimal) (moderate) (full)))
+    ;; Imenu.
+    (setq-local imenu-create-index-function #'js--treesit-imenu)
+    ;; Which-func (use imenu).
+    (setq-local which-func-functions nil)
     (treesit-major-mode-setup))
    ;; Elisp.
    (t