Files
emacs-config/elisp/zsh-ts-mode.el

464 lines
18 KiB
EmacsLisp

;;; zsh-ts-mode.el --- TreeSitter based mode for editing Zsh sripts -*- lexical-binding: t; -*-
;;; 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."
(let ((ns (treesit-node-start node))
(ne (treesit-node-end node)))
(save-excursion
(goto-char ns)
(cond
((re-search-forward (rx (or "[" "=")) ne t)
(put-text-property ns (1- (point)) '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))
"Query used in `zsh-ts-mode--fontify-arithmetic-variables'.")
(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
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
'font-lock-variable-use-face))))
(defun zsh-ts-mode--other-keywords ()
"Return a list of `other' keywords for Zsh.
This just appends some more stuff to `sh-mode--treesit-other-keywords'."
(append (sh-mode--treesit-other-keywords) '("let")))
(defvar zsh-ts-mode--declaration-commands
'("typeset" "local" "integer" "float" "readonly" "export" "declare")
"Return a list of keywords that declare a parameter in Zsh.
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
:feature 'comment
:language 'zsh
'((comment) @font-lock-comment-face)
:feature 'function
:language 'zsh
'((function_definition name: (word) @font-lock-function-name-face))
:feature 'string
:language 'zsh
'(((raw_string) @font-lock-string-face)
((string :anchor ("\"" @font-lock-string-face)
[(string_content)
(arithmetic_expansion :anchor "((")]
:* @font-lock-string-face
("\"" @font-lock-string-face) :anchor)))
:feature 'string-interpolation
:language 'zsh
:override t
'((command_substitution :anchor ("$(" @sh-quoted-exec)
(")" @sh-quoted-exec) :anchor)
;; work around parser error of "(( ))" (no $) being an arithmetic expansion
(string (arithmetic_expansion :anchor ("$((" @sh-quoted-exec)
("))" @sh-quoted-exec) :anchor)))
:feature 'variable-reference
:language 'zsh
:override t
'((expansion (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_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))
:feature 'math-function-call
:language 'zsh
:override t
`((arithmetic_call name: (word) @font-lock-function-name-face))
:feature 'heredoc
:language 'zsh
'([(heredoc_start) (heredoc_body)] @sh-heredoc)
:feature 'keyword
:language 'zsh
`(;; keywords
[ ,@zsh-ts-mode--treesit-keywords ] @font-lock-keyword-face
;; reserved words
(command_name
((word) @font-lock-keyword-face
(:match ,(rx-to-string
`(seq bos
(or ,@(zsh-ts-mode--other-keywords))
eos))
@font-lock-keyword-face))))
:feature 'command
:language 'zsh
`(;; builtin commands
(command_name
((word) @font-lock-builtin-face
(:match ,(let ((builtins
(sh-feature sh-builtins)))
(rx-to-string
`(seq bos
(or ,@builtins)
eos)))
@font-lock-builtin-face)))
;; function/non-builtin command calls
(command_name (word) @font-lock-function-name-face))
:feature 'declaration-command
:language 'zsh
;; the parser doesn't treat "let" as a keyword
`([,@zsh-ts-mode--declaration-commands] @font-lock-keyword-face)
;; variable declarations in the form of a=b
;; e.g. local x=2
:feature 'decl-equal-variable
:language 'zsh
'((variable_name) @font-lock-variable-name-face)
;; variable declarations in the form of a
;; e.g. local y z
:feature 'decl-solo-variable
:language 'zsh
:override 'prepend
`((declaration_command
argument: ((_) @font-lock-variable-name-face
(:match ,(rx-to-string '(seq bos (not "-")))
@font-lock-variable-name-face))))
;; let variable declarations (handled differently by the parser)
:feature 'decl-let-variable
:language 'zsh
:override 'prepend
`((command
(command_name ((word) @ignore
(:equal "let" @ignore)))
argument: ((_) @zsh-ts-mode--fontify-let-arg
(:match ,(rx-to-string '(seq bos (not "-")))
@zsh-ts-mode--fontify-let-arg))))
:feature 'for-variable
:language 'zsh
`("for" variable: ((simple_variable_name) @font-lock-variable-name-face))
:feature 'constant
:language 'zsh
'((case_item value: (word) @font-lock-constant-face)
(file_descriptor) @font-lock-constant-face)
:feature 'operator
:language 'zsh
`([,@sh-mode--treesit-operators] @font-lock-operator-face)
:feature 'number
:language 'zsh
`(((word) @font-lock-number-face
(:match "\\`[0-9]+\\'" @font-lock-number-face))
((number) @font-lock-number-face))
:feature 'bracket
:language 'zsh
'((["(" ")" "((" "))" "[" "]" "[[" "]]" "{" "}"]) @font-lock-bracket-face)
:feature 'delimiter
:language 'zsh
'(([";" ";;"]) @font-lock-delimiter-face)
:feature 'misc-punctuation
: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))))
(defun zsh-ts-mode--xref-function-references (identifier)
"Find all references for the function IDENTIFIER in the current buffer."
(let ((nodes (treesit-query-capture
(treesit-buffer-root-node)
`((command name: (command_name (word) @name)))
nil nil t))
(ident-copy (copy-sequence identifier)))
(set-text-properties 0 (length ident-copy) nil ident-copy)
(cl-loop for node in nodes
collect (xref-make
(propertize ident-copy 'face
'font-lock-function-name-face)
(zsh-ts-mode--treesit-node-location node)))))
(defun zsh-ts-mode--xref-variable-references (identifier)
"Find all references for the variable IDENTIFIER in the current buffer."
())
(cl-defmethod xref-backend-references ((_backend (eql zsh-ts)) identifier)
"`zsh-ts-mode' implementation for finding xref references for IDENTIFIER."
(append
;; definitions are references
(xref-backend-definitions 'zsh-ts identifier)
(cl-case (get-text-property 0 'type identifier)
(function (zsh-ts-mode--xref-function-references identifier))
(variable (zsh-ts-mode--xref-variable-references identifier)))))
;;;###autoload
(define-derived-mode zsh-ts-mode sh-base-mode "Zsh"
;; This function is based mostly on `bash-ts-mode'.
"Major mode for editing Zsh shell scripts.
This mode automatically falls back to `sh-mode' if the buffer is
not written in Zsh."
:syntax-table sh-mode-syntax-table
(when (treesit-ready-p 'zsh)
(sh-set-shell "zsh")
(add-hook 'hack-local-variables-hook
#'sh-after-hack-local-variables nil t)
(treesit-parser-create 'zsh)
(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)
(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-simple-imenu-settings zsh-ts-mode--simple-imenu-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)
;;; zsh-ts-mode.el ends here