Work on zsh-ts-mode.el

elisp/zsh-ts-mode.el

@@ -2,10 +2,19 @@
 ;;; Commentary:
 ;;; Code:
 
+(require 'cl-lib)
+(require 'eieio)
 (require 'sh-script)
 (require 'treesit)
 (require 'rx)
 
+;;; #################
+;;; # Fontification #
+;;; #################
+(defvar zsh-ts-mode--previous-let-arg-query
+  (treesit-query-compile 'zsh '((word) @arg))
+  "Query used in `zsh-ts-mode--fontify-let-arg'.")
+
 (defun zsh-ts-mode--fontify-let-arg (node _override _start _end &rest _r)
   "Internal function used from `zsh-ts-mode--treesit-settings'.
 Fontify NODE, an argument to a Zsh \"let\" statement."
@@ -13,10 +22,25 @@ Fontify NODE, an argument to a Zsh \"let\" statement."
         (ne (treesit-node-end node)))
     (save-excursion
       (goto-char ns)
-      (if (re-search-forward (rx (or "=" "[")) ne t)
+      (cond
-          (put-text-property ns (1- (point)) 'face
+       ((re-search-forward (rx (or "[" "=")) ne t)
-                             'font-lock-variable-name-face)
+        (put-text-property ns (1- (point)) 'face
-        (put-text-property ns ne 'face 'font-lock-variable-name-face)))))
+                           'font-lock-variable-name-face)
+        (unless (eql (char-before) ?\[)
+          (ignore-errors
+            (with-restriction (point) ne
+              (when (looking-at (rx (? "-") (+ (any (?0 . ?9)))))
+                (put-text-property (point-min) (point-max) 'face
+                                   'font-lock-number-face))))))
+       ;; the parser is broken and parses:
+       ;; let ans=${?}
+       ;; as
+       ;; command:let word:ans= word:${}
+       ((null (treesit-query-capture (treesit-node-parent node)
+                                     zsh-ts-mode--previous-let-arg-query
+                                     (1- ns) ns))
+        (font-lock-append-text-property ns ne 'face
+                                        'font-lock-variable-name-face))))))
 
 (defvar zsh-ts-mode--arithmetic-variable-query
   (treesit-query-compile 'zsh '((word) @name))
@@ -27,9 +51,9 @@ Fontify NODE, an argument to a Zsh \"let\" statement."
   "Fontify all arithmetic variables below NODE.
 NODE should be an `arithmetic_expansion' node."
   (dolist (child (treesit-query-capture
-                  node zsh-ts-mode--arithmetic-variable-query))
+                  node zsh-ts-mode--arithmetic-variable-query nil nil t))
-    (let ((start (treesit-node-start (cdr child)))
+    (let ((start (treesit-node-start child))
-          (end (treesit-node-end (cdr child))))
+          (end (treesit-node-end child)))
       (font-lock-append-text-property start end 'face
                                       'font-lock-variable-use-face))))
 
@@ -44,6 +68,10 @@ This just appends some more stuff to `sh-mode--treesit-other-keywords'."
 Note that \"let\" isn't included here as the parser treats it as a normal
 command.")
 
+(defvar zsh-ts-mode--treesit-keywords
+  (append sh-mode--treesit-keywords '("until"))
+  "List of keywords defined in the grammar.")
+
 ;; from `bash-ts-mode'
 (defvar zsh-ts-mode--treesit-settings
   (treesit-font-lock-rules
@@ -98,7 +126,7 @@ command.")
    :feature 'keyword
    :language 'zsh
    `(;; keywords
-     [ ,@sh-mode--treesit-keywords ] @font-lock-keyword-face
+     [ ,@zsh-ts-mode--treesit-keywords ] @font-lock-keyword-face
      ;; reserved words
      (command_name
       ((word) @font-lock-keyword-face
@@ -187,6 +215,197 @@ command.")
    :language 'zsh
    '((["$"]) @font-lock-misc-punctuation-face)))
 
+;;; #########
+;;; # IMenu #
+;;; #########
+(defvar zsh-ts-mode--defun-name-query
+  (treesit-query-compile 'zsh '((function_definition
+                                 name: ((word) @name))))
+  "Query used by `zsh-ts-mode--defun-name'.")
+
+(defun zsh-ts-mode--defun-name (node)
+  "If NODE is a defun node, return its name.
+This behaves according to `treesit-defun-name-function'."
+  (when-let ((node)
+             (res (treesit-query-capture node zsh-ts-mode--defun-name-query
+                                         nil nil t)))
+    (treesit-node-text (car res) t)))
+
+(defun zsh-ts-mode--imenu-predicate (node)
+  "Predicate function for `zsh-ts-mode--simple-imenu-settings'.
+Return non-nil if NODE is not an anonymous function."
+  (zsh-ts-mode--defun-name node))
+
+(defvar zsh-ts-mode--simple-imenu-settings
+  '(("Function" "function_definition" zsh-ts-mode--imenu-predicate nil))
+  "`zsh-ts-mode' settings for `treesit-simple-imenu-settings'.")
+
+;;; ########
+;;; # xref #
+;;; ########
+(defun zsh-ts-mode--xref-backend ()
+  "Return `zsh-ts'."
+  'zsh-ts)
+
+(defun zsh-ts-mode--not-let-predicate (node)
+  "Return nil if NODE is a let node, non-nil otherwise."
+  (not (equal (treesit-node-text node) "let")))
+
+(defvar zsh-ts-mode--identifier-query
+  (treesit-query-compile
+   'zsh `((expansion
+           name: ([(simple_variable_name)
+                   (special_variable_name)]
+                  @var))
+          (expansion_with_modifier
+           name: ([(simple_variable_name)
+                   (special_variable_name)]
+                  @var))
+          ((word) @func
+           (:pred zsh-ts-mode--not-let-predicate @func))))
+  "Query for `zsh-ts-mode' `xref-backend-identifier-at-point'.")
+
+(cl-defmethod xref-backend-identifier-at-point ((_backend (eql zsh-ts)))
+  "`zsh-ts-mode' implementation for xref's identifier at point function."
+
+  (or (when (memq 'font-lock-variable-name-face
+                  (ensure-list (get-text-property (point) 'face)))
+        (let ((start (previous-single-property-change (point) 'face))
+              (end (next-single-property-change (point) 'face)))
+          (propertize
+           (buffer-substring-no-properties (if start (1+ start) (point-min))
+                                           (or end (point-max)))
+           'type 'variable)))
+      (when-let* ((nodes (treesit-query-capture (treesit-buffer-root-node)
+                                                zsh-ts-mode--identifier-query
+                                                (point) (1+ (point)))))
+        (if-let* ((var (alist-get 'var nodes)))
+            (propertize (treesit-node-text var t) 'type 'variable)
+          (when-let* ((func (alist-get 'func nodes)))
+            (propertize (treesit-node-text func t) 'type 'function))))))
+
+(defun zsh-ts-mode--treesit-node-location (node)
+  "Return an xref location pointing to NODE."
+  (save-excursion
+    (goto-char (treesit-node-start node))
+    (make-xref-file-location :file (buffer-file-name)
+                             :line (line-number-at-pos)
+                             :column (- (point) (pos-bol)))))
+
+(defun zsh-ts-mode--make-function-definition-xref (identifier node)
+  "Make an xref object pointing to IDENTIFIER.
+NODE is IDENTIFIER's tree-sitter node."
+  (xref-make (format "%s %s()"
+                     (propertize "function"
+                                 'face 'font-lock-keyword-face)
+                     (propertize identifier
+                                 'face 'font-lock-function-name-face))
+             (zsh-ts-mode--treesit-node-location node)))
+
+(defvar zsh-ts-mode--xref-function-definitions-query
+  (treesit-query-compile 'zsh '((function_definition
+                                 name: (word) @name)))
+  "Function query for `zsh-ts-mode--xref-function-definitions'.")
+
+(defun zsh-ts-mode--xref-function-definitions (identifier)
+  "Find all definitions of function IDENTIFIER in the current buffer."
+  (cl-loop for node in (treesit-query-capture
+                        (treesit-buffer-root-node)
+                        zsh-ts-mode--xref-function-definitions-query
+                        nil nil t)
+           when (equal identifier (treesit-node-text node))
+           collect (zsh-ts-mode--make-function-definition-xref identifier node)
+           into out
+           finally return (if (length= out 1)
+                              (car out)
+                            out)))
+
+(defvar zsh-ts-mode--xref-variable-definitions-query
+  (treesit-query-compile
+   'zsh `((variable_assignment name: (variable_name) @non-let)
+          (declaration_command argument: (word) @non-let)
+          (command name: ((command_name) @let.name
+                          (:equal @let.name "let"))
+                   argument: ((word) @let-arg
+                              (:match ,(rx-to-string '(seq bos (not "-")))
+                                      @let-arg)))))
+  "Variable query for `zsh-ts-mode--xref-variable-definitions'.")
+
+(defvar zsh-ts-mode--non-let-declaration-type-query
+  (treesit-query-compile 'zsh '((declaration_command :anchor _ @name)))
+  "Helper query for `zsh-ts-mode--get-non-let-assignment-type'.")
+
+(defun zsh-ts-mode--get-assignment-parent (node)
+  "Get the node that is a parent of NODE that is a variable assignment.
+This only looks up two levels."
+  (cond ((equal (treesit-node-type
+                 (treesit-node-parent node))
+                "declaration_command")
+         (treesit-node-parent node))
+        ((equal (treesit-node-type
+                 (treesit-node-parent (treesit-node-parent node)))
+                "declaration_command")
+         (treesit-node-parent (treesit-node-parent node)))))
+
+(defun zsh-ts-mode--get-non-let-assignment-type (node)
+  "Return the assignment command for NODE."
+  (when-let* ((type-node (zsh-ts-mode--get-assignment-parent node)))
+    (propertize (treesit-node-text
+                 (car (treesit-query-capture
+                       type-node
+                       zsh-ts-mode--non-let-declaration-type-query
+                       nil nil t))
+                 t)
+                'face 'font-lock-keyword-face)))
+
+(defun zsh-ts-mode--make-xref-for-non-let-assignment (node)
+  "Create an xref for an \"other\" assignment for NODE.
+An \"other\" assignment is an assignment using local, typeset, etc."
+  (let ((type (zsh-ts-mode--get-non-let-assignment-type node)))
+    (xref-make (format "%s%s" (if type (concat type " ") "")
+                       (propertize (treesit-node-text node t)
+                                   'face 'font-lock-variable-name-face))
+               (zsh-ts-mode--treesit-node-location node))))
+
+(defun zsh-ts-mode--make-xref-for-let-assignment (node)
+  "Create an xref node for a list assingment.
+NODE is the argument to let (possibly with a = in it)."
+  (save-excursion
+    (goto-char (treesit-node-start node))
+    (xref-make
+     (format "%s %s" (propertize "let" 'face 'font-lock-keyword-face)
+             (propertize (buffer-substring-no-properties
+                          (treesit-node-start node)
+                          (if (re-search-forward (rx (or "=" "["))
+                                                 (treesit-node-end node)
+                                                 t)
+                              (1- (point))
+                            (treesit-node-end node)))))
+     (zsh-ts-mode--treesit-node-location node))))
+
+(defun zsh-ts-mode--xref-variable-definitions (identifier)
+  "Find all definitions of variable IDENTIFIER in the current buffer."
+  (let (out)
+    (dolist (ent (treesit-query-capture
+                  (treesit-buffer-root-node)
+                  zsh-ts-mode--xref-variable-definitions-query))
+      (cl-destructuring-bind (tag . node) ent
+        (cl-case tag
+          (non-let
+           (when (equal (treesit-node-text node) identifier)
+             (push (zsh-ts-mode--make-xref-for-non-let-assignment node) out)))
+          (let-arg
+           (when (string-match-p (rx bos (literal identifier) (or eos "="))
+                                 (treesit-node-text node))
+             (push (zsh-ts-mode--make-xref-for-let-assignment node) out))))))
+    (if (length= out 1) (car out) out)))
+
+(cl-defmethod xref-backend-definitions ((_backend (eql zsh-ts)) identifier)
+  "`zsh-ts-mode' implementation for finding xref definitions for IDENTIFIER."
+  (cl-case (get-text-property 0 'type identifier)
+    (function (zsh-ts-mode--xref-function-definitions identifier))
+    (variable (zsh-ts-mode--xref-variable-definitions identifier))))
+
 ;;;###autoload
 (define-derived-mode zsh-ts-mode sh-base-mode "Zsh"
   ;; This function is based mostly on `bash-ts-mode'.
@@ -199,15 +418,18 @@ not written in Zsh."
     (add-hook 'hack-local-variables-hook
               #'sh-after-hack-local-variables nil t)
     (treesit-parser-create 'zsh)
-    (setq-local treesit-font-lock-feature-list
+    (setq-local treesit-font-lock-settings zsh-ts-mode--treesit-settings
+                treesit-font-lock-feature-list
                 '((comment function)
-                  (command declaration-command keyword string math-function-call)
+                  (command declaration-command keyword string
+                   math-function-call)
                   (constant heredoc number string-interpolation
                    decl-equal-variable decl-solo-variable decl-let-variable
                    for-variable variable-reference)
                   (bracket delimiter misc-punctuation operator)))
-    (setq-local treesit-font-lock-settings
+    (setq-local treesit-simple-imenu-settings zsh-ts-mode--simple-imenu-settings
-                zsh-ts-mode--treesit-settings)
+                treesit-defun-name-function #'zsh-ts-mode--defun-name
+                imenu-create-index-function #'treesit-simple-imenu)
     (treesit-major-mode-setup)))
 
 (provide 'zsh-ts-mode)