;;; ts-mode.el --- tree sitter support for TypeScript  -*- lexical-binding: t; -*-

;; Copyright (C) 2022 Free Software Foundation, Inc.

;; Author     : Theodor Thornhill <theo@thornhill.no>
;; Maintainer : Theodor Thornhill <theo@thornhill.no>
;; Created    : October 2022
;; Keywords   : typescript tsx languages tree-sitter

;; This file is part of GNU Emacs.

;; This program is free software; you can redistribute it and/or modify
;; it under the terms of the GNU General Public License as published by
;; the Free Software Foundation, either version 3 of the License, or
;; (at your option) any later version.

;; This program is distributed in the hope that it will be useful,
;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;; GNU General Public License for more details.

;; You should have received a copy of the GNU General Public License
;; along with this program.  If not, see <http://www.gnu.org/licenses/>.

;;; Code:

(require 'treesit)
(require 'rx)
(require 'js)

(defcustom ts-mode-indent-offset 2
  "Number of spaces for each indentation step in `ts-mode'."
  :type 'integer
  :safe 'integerp
  :group 'typescript)

(defvar ts-mode--syntax-table
  (let ((table (make-syntax-table)))
    ;; Taken from the cc-langs version
    (modify-syntax-entry ?_  "_"     table)
    (modify-syntax-entry ?$ "_"      table)
    (modify-syntax-entry ?\\ "\\"    table)
    (modify-syntax-entry ?+  "."     table)
    (modify-syntax-entry ?-  "."     table)
    (modify-syntax-entry ?=  "."     table)
    (modify-syntax-entry ?%  "."     table)
    (modify-syntax-entry ?<  "."     table)
    (modify-syntax-entry ?>  "."     table)
    (modify-syntax-entry ?&  "."     table)
    (modify-syntax-entry ?|  "."     table)
    (modify-syntax-entry ?` "\""     table)
    (modify-syntax-entry ?\240 "."   table)
    table)
  "Syntax table for `ts-mode'.")

(defvar ts-mode--indent-rules
  `((tsx
     ((node-is "}") parent-bol 0)
     ((node-is ")") parent-bol 0)
     ((node-is "]") parent-bol 0)
     ((node-is ">") parent-bol 0)
     ((node-is ".")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "ternary_expression")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "named_imports")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "statement_block")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "type_arguments")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "variable_declarator")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "arguments")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "array")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "formal_parameters")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "template_substitution")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "object_pattern")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "object")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "object_type")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "enum_body")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "arrow_function")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "parenthesized_expression")
      parent-bol ,ts-mode-indent-offset)

     ;; TSX
     ((parent-is "jsx_opening_element")
      parent ,ts-mode-indent-offset)
     ((node-is "jsx_closing_element") parent 0)
     ((parent-is "jsx_element")
      parent ,ts-mode-indent-offset)
     ((node-is "/") parent 0)
     ((parent-is "jsx_self_closing_element")
      parent ,ts-mode-indent-offset)
     (no-node parent-bol 0)))
  "Tree-sitter indent rules.")

(defvar ts-mode--keywords
  '("!" "abstract" "as" "async" "await" "break"
    "case" "catch" "class" "const" "continue" "debugger"
    "declare" "default" "delete" "do" "else" "enum"
    "export" "extends" "finally" "for" "from" "function"
    "get" "if" "implements" "import" "in" "instanceof" "interface"
    "keyof" "let" "namespace" "new" "of" "private" "protected"
    "public" "readonly" "return" "set" "static" "switch"
    "target" "throw" "try" "type" "typeof" "var" "void"
    "while" "with" "yield")
  "TypeScript keywords for tree-sitter font-locking.")

(defvar ts-mode--font-lock-minimum-settings
  (list
   :language 'tsx
   :override t
   :feature 'minimal
   `(
     ((identifier) @font-lock-constant-face
      (:match "^[A-Z_][A-Z_\\d]*$" @font-lock-constant-face))

     [,@ts-mode--keywords] @font-lock-keyword-face
     [(this) (super)] @font-lock-keyword-face

     [(true) (false) (null)] @font-lock-constant-face
     (regex pattern: (regex_pattern)) @font-lock-string-face
     (number) @font-lock-constant-face

     (string) @font-lock-string-face

     (template_string) @ts-mode--fontify-template-string
     (template_substitution ["${" "}"] @font-lock-builtin-face)

     (comment) @font-lock-comment-face))
  "Font lock settings for minimal TypeScript highlights.
This level will font-lock only keywords, strings, comments and
some identifiers.")

(defvar ts-mode--font-lock-medium-settings
  (list
   :language 'tsx
   :override t
   :feature 'medium
   '(
     (nested_type_identifier
      module: (identifier) @font-lock-type-face)

     (type_identifier) @font-lock-type-face

     (predefined_type) @font-lock-type-face

     (new_expression
      constructor: (identifier) @font-lock-type-face)

     (function
      name: (identifier) @font-lock-function-name-face)

     (function_declaration
      name: (identifier) @font-lock-function-name-face)

     (method_definition
      name: (property_identifier) @font-lock-function-name-face)

     (variable_declarator
      name: (identifier) @font-lock-variable-name-face)

     (enum_declaration (identifier) @font-lock-type-face)

     (enum_body (property_identifier) @font-lock-type-face)

     (enum_assignment name: (property_identifier) @font-lock-type-face)

     (assignment_expression
      left: [(identifier) @font-lock-variable-name-face
             (member_expression
              property: (property_identifier) @font-lock-variable-name-face)])

     (for_in_statement
      left: (identifier) @font-lock-variable-name-face)

     (arrow_function
      parameter: (identifier) @font-lock-variable-name-face)))
  "Font lock settings for medium TypeScript highlights.
This levels adds some flair to the more advanced patterns.")

(defvar ts-mode--font-lock-full-settings
  (list
   :language 'tsx
   :override t
   :feature 'full
   '(
     (variable_declarator
      name: (identifier) @font-lock-function-name-face
      value: [(function) (arrow_function)])

     (variable_declarator
      name: (array_pattern
             (identifier)
             (identifier) @font-lock-function-name-face)
      value: (array (number) (function)))

     (assignment_expression
      left: [(identifier) @font-lock-function-name-face
             (member_expression
              property: (property_identifier) @font-lock-function-name-face)]
      right: [(function) (arrow_function)])

     (call_expression
      function:
      [(identifier) @font-lock-function-name-face
       (member_expression
        property: (property_identifier) @font-lock-function-name-face)])

     (arrow_function
      parameters:
      [(_ (identifier) @font-lock-variable-name-face)
       (_ (_ (identifier) @font-lock-variable-name-face))
       (_ (_ (_ (identifier) @font-lock-variable-name-face)))])

     (pair key: (property_identifier) @font-lock-variable-name-face)

     (pair value: (identifier) @font-lock-variable-name-face)

     (pair
      key: (property_identifier) @font-lock-function-name-face
      value: [(function) (arrow_function)])

     (property_signature
      name: (property_identifier) @font-lock-variable-name-face)

     ((shorthand_property_identifier) @font-lock-variable-name-face)

     (pair_pattern
      key: (property_identifier) @font-lock-variable-name-face)

     ((shorthand_property_identifier_pattern)
      @font-lock-variable-name-face)

     (array_pattern (identifier) @font-lock-variable-name-face)

     (jsx_opening_element
      [(nested_identifier (identifier)) (identifier)]
      @font-lock-function-name-face)

     (jsx_closing_element
      [(nested_identifier (identifier)) (identifier)]
      @font-lock-function-name-face)

     (jsx_self_closing_element
      [(nested_identifier (identifier)) (identifier)]
      @font-lock-function-name-face)

     (jsx_attribute (property_identifier) @font-lock-constant-face)))
  "Font lock settings for full TypeScript highlights.
This level yields the richest experience, with support for
pattern matching and TSX.")

(defvar ts-mode--font-lock-settings
  (apply #'treesit-font-lock-rules
         (append ts-mode--font-lock-minimum-settings
                 ts-mode--font-lock-medium-settings
                 ts-mode--font-lock-full-settings))
  "Tree-sitter font-lock settings.")

(defun ts-mode--fontify-template-string (beg end node)
  "Fontify template string but not substitution inside it.
BEG, END, NODE refers to the template_string node."
  (ignore end)
  ;; Stolen from `js--fontify-template-string'
  (let ((child (treesit-node-child node 0)))
    (while child
      (if (equal (treesit-node-type child) "template_substitution")
          (put-text-property beg (treesit-node-start child)
                             'face 'font-lock-string-face)
        (put-text-property beg (treesit-node-end child)
                           'face 'font-lock-string-face))
      (setq beg (treesit-node-end child)
            child (treesit-node-next-sibling child)))))

(defvar ts-mode--defun-type-regexp
  (rx (or "class_declaration"
          "method_definition"
          "function_declaration"
          "lexical_declaration"))
  "Regular expression that matches type of defun nodes.
Used in `ts-mode--beginning-of-defun' and friends.")

(defun ts-mode--beginning-of-defun (&optional arg)
  "Tree-sitter `beginning-of-defun' function.
ARG is the same as in `beginning-of-defun."
  (let ((arg (or arg 1)))
    (if (> arg 0)
        ;; Go backward.
        (while (and (> arg 0)
                    (treesit-search-forward-goto
                     ts-mode--defun-type-regexp 'start nil t))
          (setq arg (1- arg)))
      ;; Go forward.
      (while (and (< arg 0)
                  (treesit-search-forward-goto
                   ts-mode--defun-type-regexp 'start))
        (setq arg (1+ arg))))))

(defun ts-mode--end-of-defun (&optional arg)
  "Tree-sitter `end-of-defun' function.
ARG is the same as in `end-of-defun."
  (let ((arg (or arg 1)))
    (if (< arg 0)
        ;; Go backward.
        (while (and (< arg 0)
                    (treesit-search-forward-goto
                     ts-mode--defun-type-regexp 'end nil t))
          (setq arg (1+ arg)))
      ;; Go forward.
      (while (and (> arg 0)
                  (treesit-search-forward-goto
                   ts-mode--defun-type-regexp 'end))
        (setq arg (1- arg))))))

;;;###autoload
(add-to-list 'auto-mode-alist '("\\.ts\\'" . ts-mode))

;;;###autoload
(add-to-list 'auto-mode-alist '("\\.tsx\\'" . ts-mode))

;;;###autoload
(define-derived-mode ts-mode prog-mode "TypeScript"
  "Major mode for editing TypeScript."
  :group 'typescript
  :syntax-table ts-mode--syntax-table

  (cond
   ((and (treesit-can-enable-p)
         (treesit-language-available-p 'tsx))
    ;; Comments
    (setq-local comment-start "// ")
    (setq-local comment-start-skip "\\(?://+\\|/\\*+\\)\\s *")
    (setq-local comment-end "")

    (setq-local treesit-simple-indent-rules ts-mode--indent-rules)
    (setq-local indent-line-function #'treesit-indent)

    (setq-local beginning-of-defun-function #'ts-mode--beginning-of-defun)
    (setq-local end-of-defun-function #'ts-mode--end-of-defun)

    (unless font-lock-defaults
      (setq font-lock-defaults '(nil t)))

    (setq-local treesit-font-lock-settings ts-mode--font-lock-settings)

    (setq treesit-font-lock-feature-list '((minimal medium full)))

    (treesit-font-lock-enable))
   (t
    (message "Tree sitter for TypeScript isn't available, defaulting to js-mode")
    (js-mode))))

(provide 'ts-mode)

;;; ts-mode.el ends here