(Hopefully) Finished zsh-ts-mode.el

elisp/zsh-ts-mode.el

@@ -8,6 +8,11 @@
 (require 'treesit)
 (require 'rx)
 
+(defgroup zsh-ts-mode nil
+  "Tree-sitter powered Zsh editing."
+  :group 'sh
+  :prefix "zsh-ts-mode-")
+
 ;;; #################
 ;;; # Fontification #
 ;;; #################
@@ -43,15 +48,26 @@ Fontify NODE, an argument to a Zsh \"let\" statement."
                                         'font-lock-variable-name-face))))))
 
 (defvar zsh-ts-mode--arithmetic-variable-query
-  (treesit-query-compile 'zsh '((word) @name))
+  (treesit-query-compile
-  "Query used in `zsh-ts-mode--fontify-arithmetic-variables'.")
+   'zsh '(((word) @use)
+          (binary_expression left: (variable_name) @assign
+                             operator: "=")))
+  "Query used in `zsh-ts-mode--arithmatic-variables'.")
+
+(defun zsh-ts-mode--arithmatic-variables (parent also-define)
+  "Return a list of arithmetic variables in the arithmetic_expansion PARENT.
+If ALSO-DEFINE is non-nil, also capture defined (e.g. x=2)."
+  (cl-loop for entry in (treesit-query-capture
+                         parent zsh-ts-mode--arithmetic-variable-query)
+           when (or also-define
+                    (eq (car entry) 'use))
+           collect (cdr entry)))
 
 (defun zsh-ts-mode--fontify-arithmetic-variables
     (node _override _start _end &rest _r)
   "Fontify all arithmetic variables below NODE.
 NODE should be an `arithmetic_expansion' node."
-  (dolist (child (treesit-query-capture
+  (dolist (child (zsh-ts-mode--arithmatic-variables node nil))
-                  node zsh-ts-mode--arithmetic-variable-query nil nil t))
     (let ((start (treesit-node-start child))
           (end (treesit-node-end child)))
       (font-lock-append-text-property start end 'face
@@ -104,13 +120,8 @@ command.")
    :feature 'variable-reference
    :language 'zsh
    :override t
-   '((expansion (simple_variable_name) @font-lock-variable-use-face)
+   '((simple_variable_name) @font-lock-variable-use-face
      (special_variable_name) @font-lock-variable-use-face
-     (expansion_with_modifier (simple_variable_name)
-                              @font-lock-variable-use-face)
-     (expansion_pattern (simple_variable_name) @font-lock-variable-use-face)
-     (expansion_substring (simple_variable_name) @font-lock-variable-use-face)
-     (variable_ref (simple_variable_name) @font-lock-variable-use-face)
      ;; variables in math expansions
      ((arithmetic_expansion :anchor "$((")
       @zsh-ts-mode--fontify-arithmetic-variables))
@@ -269,28 +280,41 @@ Return non-nil if NODE is not an anonymous function."
            name: ([(simple_variable_name)
                    (special_variable_name)]
                   @var))
+          (expansion_default
+           name: ([(simple_variable_name)
+                   (special_variable_name)]
+                  @var))
+          ((variable_name) @var)
           ((word) @func
            (:pred zsh-ts-mode--not-let-predicate @func))))
   "Query for `zsh-ts-mode' `xref-backend-identifier-at-point'.")
 
+(cl-defun zsh-ts-mode--extract-variable-from-face-at (&optional (pos (point)))
+  "Extract the variable under POS from its face."
+  (cl-flet ((var-face-p (pos)
+              (memq 'font-lock-variable-name-face
+                    (ensure-list (get-text-property pos 'face)))))
+    (when (var-face-p pos)
+      (let ((start (if (and (not (bobp)) (var-face-p (1- pos)))
+                       (or (previous-single-property-change pos 'face)
+                           (point-min))
+                     (point)))
+            (end (or (next-single-property-change (point) 'face)
+                     (point-max))))
+        (propertize
+         (buffer-substring-no-properties start end) 'type 'variable)))))
+
 (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
+  (or (zsh-ts-mode--extract-variable-from-face-at)
-                  (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)))
+        (if-let* ((func (alist-get 'func nodes)))
-            (propertize (treesit-node-text var t) 'type 'variable)
+            (propertize (treesit-node-text func t) 'type 'function)
-          (when-let* ((func (alist-get 'func nodes)))
+          (when-let* ((var (alist-get 'var nodes)))
-            (propertize (treesit-node-text func t) 'type 'function))))))
+            (propertize (treesit-node-text var t) 'type 'variable))))))
 
 (cl-defmethod xref-backend-identifier-completion-table ((_backend (eql zsh-ts)))
   "Return nil."
@@ -353,6 +377,8 @@ Return non-nil if NODE is not an anonymous function."
   (treesit-query-compile
    'zsh `((variable_assignment name: (variable_name) @non-let)
           (declaration_command argument: (word) @non-let)
+          (binary_expression left: (variable_name) @non-let
+                             operator: "=")
           (command name: ((command_name) @let.name
                           (:equal @let.name "let"))
                    argument: ((word) @let-arg
@@ -405,11 +431,24 @@ Return non-nil if NODE is not an anonymous function."
              when (save-excursion
                     (goto-char (treesit-node-start node))
                     (looking-at-p quote-ident))
-             collect (zsh-ts-mode--make-xref-for-node node))))
+             collect (zsh-ts-mode--make-xref-for-node node ))))
+
+(defvar zsh-ts-mode--string-query
+  (treesit-query-compile 'zsh '((string) @s))
+  "Query for `zsh-ts-mode--node-not-in-string-p'.")
+
+(defun zsh-ts-mode--node-in-string-p (node)
+  "Return non-nil if NODE is in a string."
+  (treesit-query-capture (treesit-buffer-root-node)
+                         zsh-ts-mode--string-query
+                         (treesit-node-start node)
+                         (treesit-node-start node)))
 
 (defvar zsh-ts-mode--variable-references-query
   (treesit-query-compile
-   'zsh '(([(simple_variable_name) (special_variable_name)] @var)))
+   'zsh '(([(simple_variable_name) (special_variable_name)] @var)
+          ((arithmetic_expansion :anchor "$((") @arith)
+          ((arithmetic_expansion :anchor "((") @arith-check)))
   "Query used by `zsh-ts-mode--xref-variable-references'.")
 
 (defun zsh-ts-mode--xref-variable-references (identifier)
@@ -418,14 +457,21 @@ Return non-nil if NODE is not an anonymous function."
         (quote-ident (regexp-quote identifier)))
     (set-text-properties 0 (length ident-copy) () ident-copy)
     (mapcan
-     #'(lambda (node)
+     #'(lambda (ent)
-         (when (save-excursion
+         (cl-destructuring-bind (tag . node) ent
-                 (goto-char (treesit-node-start node))
+           (cl-case tag
-                 (looking-at-p quote-ident))
+            (var (when (save-excursion
-           (list (zsh-ts-mode--make-xref-for-node node))))
+                         (goto-char (treesit-node-start node))
+                         (looking-at-p quote-ident))
+                   (list (zsh-ts-mode--make-xref-for-node node))))
+            (arith (mapcar #'zsh-ts-mode--make-xref-for-node
+                           (zsh-ts-mode--arithmatic-variables node nil)))
+            (arith-check
+             (unless (zsh-ts-mode--node-in-string-p node)
+               (mapcar #'zsh-ts-mode--make-xref-for-node
+                       (zsh-ts-mode--arithmatic-variables node nil)))))))
      (treesit-query-capture (treesit-buffer-root-node)
-                            zsh-ts-mode--variable-references-query
+                            zsh-ts-mode--variable-references-query))))
-                            nil nil t))))
 
 (cl-defmethod xref-backend-references ((_backend (eql zsh-ts)) identifier)
   "`zsh-ts-mode' implementation for finding xref references for IDENTIFIER."
@@ -443,11 +489,14 @@ Return non-nil if NODE is not an anonymous function."
    'zsh '(((simple_variable_name) @ident)
           ((special_variable_name) @ident)
           ((variable_name) @ident)
-          (command name: (command_name ( (word) @cname (:equal @cname "let")))
+          (command name: (command_name ((word) @cname (:equal @cname "let")))
                    argument: (word) @ident)
           (declaration_command argument: (word) @ident)
+          (command name: (command_name (word) @ident))
           ((variable_name) @ident)
-          (function_definition name: (word) @ident)))
+          (function_definition name: (word) @ident)
+          ((arithmetic_expansion :anchor "$((") @arith)
+          ((arithmetic_expansion :anchor "((") @arith-check)))
   "Query used by `xref-backend-apropos'.")
 
 (cl-defmethod xref-backend-apropos ((_backend (eql zsh-ts)) pattern)
@@ -456,9 +505,14 @@ Return non-nil if NODE is not an anonymous function."
     (dolist (ent (treesit-query-capture (treesit-buffer-root-node)
                                         zsh-ts-mode--xref-apropos-query))
       (cl-destructuring-bind (type . node) ent
-        (when (and (eq type 'ident)
+        (cond
-                   (string-match-p pattern (treesit-node-text node)))
+         ((eq type 'ident) (when (string-match-p pattern (treesit-node-text node))
-          (push (zsh-ts-mode--make-xref-for-node node) out))))
+                             (push (zsh-ts-mode--make-xref-for-node node) out)))
+         ((or (eq type 'arith)
+              (not (zsh-ts-mode--node-in-string-p node)))
+          (dolist (var-node (zsh-ts-mode--arithmatic-variables node nil))
+            (when (string-match-p pattern (treesit-node-text var-node))
+              (push (zsh-ts-mode--make-xref-for-node var-node) out)))))))
     out))
 
 ;;;###autoload
@@ -488,5 +542,19 @@ not written in Zsh."
     (add-hook 'xref-backend-functions #'zsh-ts-mode--xref-backend nil t)
     (treesit-major-mode-setup)))
 
+(derived-mode-add-parents 'zsh-ts-mode '(sh-mode))
+
+;;;###autoload
+(add-to-list 'auto-mode-alist
+             `(,(rx (or ".zsh"
+                        (seq (or "/" bos)
+                             (or ".zshrc" ".zprofile" ".zlogin"
+                                 ".zshenv")))
+                    eos)
+               . zsh-ts-mode))
+
+;;;###autoload
+(add-to-list 'interpreter-mode-alist '("zsh" . zsh-ts-mode))
+
 (provide 'zsh-ts-mode)
 ;;; zsh-ts-mode.el ends here

init.el

@@ -2134,7 +2134,7 @@ This is :around advice, so OLDFUN is the real function
   :config
   (my/setup-c-style-newline-keys typescript-ts-mode-map))
 
-;; shell-mode
+;; sh-mode
 (use-package sh-script
   :ensure nil
   :hook (sh-mode . my/-setup-sh-mode)
@@ -2143,6 +2143,10 @@ This is :around advice, so OLDFUN is the real function
   (defun my/-setup-sh-mode ()
     (add-hook 'completion-at-point-functions #'cape-file nil t)))
 
+;; zsh-ts-mode
+(use-package zsh-ts-mode
+  :ensure nil)
+
 ;; go mode
 (use-package go-mode
   :defer nil