tree-edit.el

;;; tree-edit.el --- A library for structural refactoring and editing -*- lexical-binding: t; -*-
;;
;; Copyright (C) Ethan Leba <https://github.com/ethan-leba>
;;
;; Author: Ethan Leba <ethanleba5@gmail.com>
;; Version: 0.1.0
;; Homepage: https://github.com/ethan-leba/tree-edit
;; Package-Requires: ((emacs "29.1") (dash "2.19") (reazon "0.4.0") (s "0.0.0"))
;; SPDX-License-Identifier: GPL-3.0-or-later
;;
;; This file is not part of GNU Emacs.
;;
;;; Commentary:
;;
;; Provides a set of functions for structural editing or refactoring in any
;; language supported by tree-sitter.
;;
;; The interface for this package is moderately stable, but no promises ;)
;;
;; See `evil-tree-edit' if you're looking for a complete editing package.
;;
;;; Code:
;;* Requires
(require 'treesit)
(require 'mode-local)
(require 'dash)
(require 'reazon)
(require 's)

;;* Parser-local utilities
(defun tree-edit--lookup-parser-value (var)
  "Lookup the value of VAR in the current parser language."
  (alist-get
   (thread-last
     (treesit-buffer-root-node)
     (treesit-node-parser)
     (treesit-parser-language))
   var))

(defmacro tree-edit--def-parser-local-var (var &rest args)
  "Define a parser-local variables named VAR with optional ARGS."
  (declare (indent defun))
  `(progn
     (defvar ,var nil ,@args)
     (defun ,var ()
       "Return the value of the parser-local variable VAR."
       (tree-edit--lookup-parser-value ,var))))

(defmacro tree-edit--set-parser-local-vars (parser &rest args)
  "Set parser-local variables in PARSER using ARGS.

ARGS should be an alternating list of VAR VALUE VAR VALUE etc."
  `(let (var val (args ',args))
     (while args
       (setq var (car args)
             val (cadr args)
             args (cddr args))
       (set var (cons (cons ,parser val) (symbol-value var))))))

;;* Internal variables
(tree-edit--def-parser-local-var tree-edit-grammar
  "The grammar rules generated by tree-sitter. Set by mode-local grammar file.")
(tree-edit--def-parser-local-var tree-edit--supertypes
  "A mapping from type to supertype, i.e. if_statement is a statement. Set by mode-local grammar file.")
(tree-edit--def-parser-local-var tree-edit--subtypes
  "A mapping from type to subtype, i.e. statement is subtyped by if_statement. Set by mode-local grammar file.")
(tree-edit--def-parser-local-var tree-edit--containing-types
  "A mapping from a type to all possible types that can exist as it's children. Set by mode-local grammar file.")
(tree-edit--def-parser-local-var tree-edit--alias-map
  "A mapping from a type to a mapping from original name to aliased name. Set by mode-local grammar file.")
(tree-edit--def-parser-local-var tree-edit--hidden-node-types
  "Nodes which are hidden by tree-sitter. Set by mode-local grammar file.

Unfortunately tree-sitter allows certain nodes to be hidden from
the syntax tree, which will throw off tree-edit's parser. The
best we can do for now is pretend that these nodes don't exist at
all.

https://tree-sitter.github.io/tree-sitter/creating-parsers#hiding-rules")
(defvar tree-edit--query-cursors nil
  "A mapping from language to query cursor for performance.")
(defvar tree-edit--string-parse-buffer nil
  "The buffer being used to store a node's text when parsing a string. This is a bit of a hack. Meant to be shadowed with let.")
;; TODO: This should be mode local
(defvar tree-edit--type-cache (make-hash-table :test #'equal)
  "Caches the type and split node for a given piece of text.

This cache both as an optimization for text we copied through
tree-edit, and works around certain edge cases.")
(defvar tree-edit-parse-comments t
  "If non-nil, allow abritrary 'comment' nodes when parsing.

This should be enabled when using parser for insertions, but
seeing the comments is unnecessary when previewing the parser.")
(tree-edit--def-parser-local-var tree-edit-node-insertion-override
  "A mapping from type to function, overriding `tree-edit--valid-insertions'.

Set by mode-local grammar file.

This should only be used in one of the following cases:

1. Performance, i.e. blocks or compound statements, where there
is no syntax between nodes, and the amount of nodes present can
be very high.

2. Nodes containing tree-sitter externals, so the JSON grammar
may not truly capture what qualifies as a valid node.")
(tree-edit--def-parser-local-var tree-edit-node-deletion-override
  "A mapping from type to function, overriding `tree-edit--valid-deletions'.

Set by mode-local grammar file.

This should only be used in one of the following cases:

1. Performance, i.e. blocks or compound statements, where there
is no syntax between nodes, and the amount of nodes present can
be very high.

2. Nodes containing tree-sitter externals, so the JSON grammar
may not truly capture what qualifies as a valid node.")
(tree-edit--def-parser-local-var tree-edit-node-replacement-override
  "A mapping from type to function, overriding `tree-edit--valid-replacement-p'.

Set by mode-local grammar file.

This should only be used in one of the following cases:

1. Performance, i.e. blocks or compound statements, where there
is no syntax between nodes, and the amount of nodes present can
be very high.

2. Nodes containing tree-sitter externals, so the JSON grammar
may not truly capture what qualifies as a valid node.")
(tree-edit--def-parser-local-var tree-edit-significant-node-types
  "List of nodes that are considered significant, like methods or classes. Set by mode-local grammar file.")
(tree-edit--def-parser-local-var tree-edit-syntax-snippets
  "Snippets for constructing nodes. Set by mode-local grammar file.

Must be an alist of node type (as a symbol) to list, where the list can
contain any string or a symbol referencing another node type in the alist.

The syntax snippets are intended to mirror the grammar constructions
of the tree-sitter grammar, but is currently not validated by tree-edit.")
(tree-edit--def-parser-local-var tree-edit-nodes
  "Nodes that a user can create via tree-edit. Set by mode-local grammar file.

Must be a list of plists, with the following properties:

Properties
  :type           the node's type
  :key            the keybinding for the given node
  :name           human readable name for which-key, defaults to
                  :type if left unset
  :node-override  overrides syntax snippets for the verb
  :wrap-override  overrides syntax snippets for the verb when wrapping")
(tree-edit--def-parser-local-var tree-edit-query-nodes
  "Nodes that users can query for jumping. Set by mode-local grammar file.

Must be a list of plists, with the following properties:

Properties
  :type           the node's type, or a list of types
  :key            the keybinding for the given node
  :name           human readable name for which-key, defaults to
                  :type if left unset")
(tree-edit--def-parser-local-var tree-edit-whitespace-rules
  "Rules for formatting nodes. Set by mode-local grammar file.

Must by an alist of node type to a pair of lists, where the car
is the whitespace rules before the node, and the cdr is after.

The following keywords are valid whitespace rules:

  :newline      insert a newline before the next text
  :indent       increase the indentation by 4 for the next newline
  :dedent       decrease the indentation by 4 for the next newline")
(tree-edit--def-parser-local-var tree-edit-placeholder-node-type
  "Node considered a placeholder. Set by mode-local grammar file.

Typically an identifier, but can conceivably be any type of node.")
(tree-edit--def-parser-local-var tree-edit-indentation-level
  "How many spaces/tabs to indent by when rendering nodes.

Set by mode-local grammar file.

TODO: Is there a builtin way to infer this from the buffer?")
;;* User settings
(defgroup tree-edit nil
  "Structural editing library for tree-sitter languages."
  :group 'bindings
  :prefix "tree-edit-")
(defcustom tree-edit-language-alist '((java-mode . java)
                                      (python-mode . python)
                                      (python-ts-mode . python)
                                      (c-mode . c)
                                      (emacs-lisp-mode . elisp))
  "Mapping from TS grammar name to language file."
  :type '(alist :key-type symbol :value-type symbol)
  :group 'tree-edit)

(defcustom tree-edit-storage-dir
  (expand-file-name
   "tree-edit/"
   (let ((cache-dir (when (fboundp 'xdg-cache-home)
                      (xdg-cache-home))))
     (if (or (seq-empty-p cache-dir)
             (not (file-exists-p cache-dir))
             (file-exists-p (expand-file-name
                             "tree-edit"
                             user-emacs-directory)))
         user-emacs-directory
       cache-dir)))
  "Storage location for preprocessed grammar files."
  :type 'string)

;; TODO: Add timeout to queries
;; (defcustom tree-edit-query-timeout 0.1
;;   "How long a query should take before giving up."
;;   :type 'float
;;   :group 'tree-edit)

;;* Data definitions
(define-error 'tree-edit-transformation-error "[Failed transformation]")
(defun tree-edit-transformation-error (format &rest args)
  "Signal transformation error with string FORMAT formatted with ARGS."
  (signal 'tree-edit-transformation-error (list (apply #'format-message format args))))

(cl-defstruct (tree-edit-result)
  "The result of a successful structural edit."
  start-index end-index tokens)

;;* Utilities
(defun tree-edit--node-type (node)
  "Emulates the behavior of `tsc-node-type'."
  (if (treesit-node-check node 'named)
      (intern (treesit-node-type node))
    (treesit-node-type node)))

(defun tree-edit--boring-nodep (node)
  "Check if the NODE is not a named node."
  (and (treesit-node-p node)
       (not (treesit-node-check node 'named))))

(defun tree-edit--all-named-descendants (node)
  "Retrieve all named descendants of NODE."
  (let (result (stack `(,node)))
    (while stack
      (let* ((item (pop stack))
             (children (treesit-node-children item :named)))
        (setq stack (append children stack))
        (setq result (append result children))))
    result))

(defun tree-edit--node-steps (node)
  "Return the sequence of steps from the root node to NODE.

Each step has the form (CHILD-NODE . NTH), where CHILD-NODE is the node to
descend into, and NTH is its 0-based ordinal position within the parent node.

If NODE is the root node, the sequence is empty."
  (let (steps parent (this node))
    (while (setq parent (treesit-node-parent this))
      (cl-block nil
        (let ((i 0))
          (-map
           (lambda (child)
             (when (treesit-node-eq child this)
               (push (cons (tree-edit--node-type this) i) steps)
               (cl-return steps))
             (setq i (1+ i)))
           (treesit-node-children parent :named)))
        (error "Unable to calculate node steps"))
      (setq this parent))
    steps))

(defun tree-edit--node-from-steps (steps)
  "Follow STEPS from TREE's root node; return the final node.
STEPS should be a sequence of steps, as described by `tree-edit--node-steps'.

If a step cannot be followed, return the node followed by all the valid steps."
  (cl-block nil
    (let ((this (treesit-buffer-root-node)))
      (pcase-dolist (`(,_ . ,step) steps)
        (when (zerop (treesit-node-child-count this :named))
          (cl-return this))
        (let ((new-node (treesit-node-child
                         this
                         (min step (1- (treesit-node-child-count this :named)))
                         :named)))
          (setq this new-node)))
      this)))

(defun tree-edit--node-from-steps-strict (steps)
  "Follow STEPS from TREE's root node; return the final node.
STEPS should be a sequence of steps, as described by `tree-edit--node-steps'.

Unlike `tree-edit--node-from-steps', the type of the node must match the type described by the steps.

If a step cannot be followed, return the node followed by all the valid steps."
  (cl-block nil
    (let ((this (treesit-buffer-root-node)))
      (pcase-dolist (`(,old-type . ,i) steps)
        (let ((new-node (treesit-node-child this i :named)))
          (unless new-node
            (cl-return nil))
          (let ((new-type (tree-edit--node-type new-node)))
            (unless (equal old-type new-type)
              (cl-return nil)))
          (setq this new-node)))
      this)))

(defun tree-edit--apply-until-interesting (fun node)
  "Apply FUN to NODE until a named node is hit."
  (let ((parent (funcall fun node)))
    (if (tree-edit--boring-nodep parent)
        (tree-edit--apply-until-interesting fun parent)
      parent)))

(defun tree-edit-all-aliases-for-type (type)
  "Retrieve all aliases that a given node TYPE has. Useful for querying."
  (--mapcat
   (if-let (alias (alist-get type (cdr it))) `(,alias))
   (tree-edit--alias-map)))

(defun tree-edit--format-query-string (node-type)
  "Format a query string for NODE-TYPE.

NODE-TYPE can be a symbol or a list of symbol."
  (interactive)
  (--> (if (listp node-type) node-type `(,node-type))
       (-map (lambda (type) (format "(%s)" type)) it)
       (string-join it " ")
       ;; Query string needs an @name here, or it won't return any results
       (format "[%s] @node" it)))

(cl-defun tree-edit-query (patterns node &key want-text)
  "Execute query PATTERNS against the children of NODE and return captures.

If WANT-TEXT is non-nil, the text will be retrieved for each
node. Only use this if actually operating on the text for
performance."
  (-map #'cdr (treesit-query-capture node patterns)))

(defun tree-edit--relevant-types (type parent-type)
  "Return a list of the TYPE and all relevant types that occur in PARENT-TYPE.

Relevant types are either supertypes of TYPE or alias names referring to TYPE."
  (-intersection
   (cons
    (alist-get type (alist-get parent-type (tree-edit--alias-map)))
    (alist-get type (tree-edit--supertypes) `(,type)))
   (alist-get parent-type (tree-edit--containing-types))))

;; TODO: Do lazily
(defun tree-edit-load-grammar-for-major-mode ()
  "Load the grammar for the major mode, or error if none is registered."
  (let* ((parser (alist-get major-mode tree-edit-language-alist)))
    (unless parser
      (user-error
       "No parser specified for mode `%s' in `tree-edit-language-alist'!"
       (symbol-name major-mode)))

    (thread-last
      (symbol-name parser)
      (format "tree-edit-%s")
      (intern)
      (require))
    
    ;; TODO: error handling
    (treesit-parser-create parser)

    (let ((grammar-file (expand-file-name
                         (format "tree-edit-%s-grammar.el" parser)
                         tree-edit-storage-dir)))
      (if (file-exists-p grammar-file)
          (let (sexp)
            (with-temp-buffer
              (insert-file-contents grammar-file)
              (goto-char (point-min))
              ;; TODO: Error handling?
              (setq sexp (read (current-buffer))))

            (--each sexp
              (-let [(var . val) it]
                ;; FIXME
                (set var (cons (cons parser val) (symbol-value var))))))
        (user-error "No grammar file present for %s! Either install the grammar or remove %s from `tree-edit-language-alist'" major-mode major-mode)))) ())

;;* Locals: node transformations


;; Error recovery seems to be a bit arbitrary:
;; - "foo.readl" in java parses as (program (expression_statement (...) (MISSING \";\")))
;; - "foo.read" in java parses as (program (ERROR (...)))
(defun tree-edit--parse-fragment (fragment)
  "Return the possible nodes of FRAGMENT, or nil if unparseable.

For example, `foo()` in Python parses as an expression_statement
with a call inside. Depending on the context, we may want either:
so we return both.

Fragments should parse as one of the following structures:
- (program (type ...)
- (program (ERROR (type ...))
- (program (... (type ...) (MISSING ...))"
  (cl-flet ((tree-edit--get-only-child
              (lambda (node) (if (equal (treesit-node-child-count node :named) 1)
                                 (treesit-node-child node 0 :named)))))
    (if-let ((first-node
              ;; FIXME document or refactor
              (progn
                (tree-edit--get-only-child
                 fragment))))
        (if-let (node (if (treesit-node-check first-node 'has-error)
                          (-some-> first-node
                            (tree-edit--get-only-child))
                        first-node))
            (let (result)
              (while node
                (let ((tmp (tree-edit--get-only-child node)))
                  (if (or (not result)
                          (and (equal (treesit-node-start node)
                                      (treesit-node-start (car result)))
                               (equal (treesit-node-end node)
                                      (treesit-node-end (car result)))))
                      (progn
                        (push node result)
                        (setq node tmp))
                    (setq node nil))))
              (unless result
                (error "tree-edit internal error: could not parse fragment"))
              (reverse result))))))

(defun tree-edit--get-parent-tokens (node)
  "Return a pair containing the siblings of the NODE and the index of itself."
  (let* ((parent (treesit-node-parent node))
         (children (treesit-node-children parent)))
    (cons (-map #'tree-edit--node-type children)
          (--find-index (treesit-node-eq node it) children))))

(defun tree-edit-simple-delete-override (_ _ start-index end-index)
  "Allow deletion of NODE and surrounding syntax."
  (make-tree-edit-result :start-index start-index :end-index end-index :tokens nil))

(defun tree-edit-simple-insertion-replacement-override (type parent start-index end-index)
  "Allow insertion of TYPE if it appears in NODE's parent's grammar."
  (when-let (types (tree-edit--relevant-types type (tree-edit--node-type parent)))
    (make-tree-edit-result :start-index start-index :end-index (or end-index start-index) :tokens `(,type))))

;;* Globals: Syntax generation
(defun tree-edit--find-raise-ancestor (ancestor child)
  "Find a suitable ANCESTOR to be replaced with CHILD."
  (interactive)
  (let ((child-type (tree-edit--node-type child)))
    (cond
     ((not (and ancestor (treesit-node-parent ancestor)))
      (tree-edit-transformation-error "Can't raise node!"))
     ;; XXX: For cases like (expression_statement (call)), where both represent the same text.
     ;;      This might only apply to Python.
     ;; TODO: Likely more efficient to compare the start & end points
     ((equal (treesit-node-text ancestor) (treesit-node-text child))
      (tree-edit--find-raise-ancestor (treesit-node-parent ancestor) child))
     ((tree-edit--valid-replacement-p child-type ancestor) ancestor)
     (t (tree-edit--find-raise-ancestor (treesit-node-parent ancestor) child)))))

(defun tree-edit--valid-node-including-type (type parent-type)
  "Return a valid sequence of tokens for PARENT-TYPE containing TYPE, or nil."
  (tree-edit-load-grammar-for-major-mode)
  (-let* ((reazon-occurs-check nil)
          (grammar (alist-get parent-type (tree-edit-grammar)))
          (relevant-types (tree-edit--relevant-types type parent-type)))
    (car (tree-edit--run-relation 1 q
           (lambda (tokens) (and (listp tokens) (equal 1 (-count #'symbolp tokens))))
           (tree-edit--max-lengtho q 5)
           (tree-edit--includes-typeo q relevant-types)
           (tree-edit-parseo grammar q '())))))

(defun tree-edit--remove-node-and-surrounding-syntax (tokens idx)
  "Return a pair of indices to remove the node at IDX in TOKENS and all surrounding syntax."
  (let ((end (1+ idx))
        (start (1- idx)))
    (while (stringp (nth end tokens))
      (setq end (1+ end)))
    (while (and (>= start 0) (stringp (nth start tokens)))
      (setq start (1- start)))
    (cons (1+ start) end)))

(defmacro tree-edit--run-relation (tries var pred &rest goals)
  "Run GOALS against VAR, returning the first matching PRED.

Run PRED against TRIES answers, or all if TRIES is nil."
  (declare (indent 2))
  `(let ((,var (reazon--make-variable ',var))
         (reazon--stop-time (and reazon-timeout (+ reazon-timeout (float-time)))))
     (tree-edit--take-first (reazon--run-goal (reazon-conj ,@goals)) ,pred ,tries ,var)))

(cl-defun tree-edit--attempt-structural-edit
    (parent &key (new-type nil) start-index (end-index nil) (overrides-alist nil))
  "Attempt to perform structural edit.

Returns a `tree-edit-result' if successful, otherwise nil.

This is the main interface to the miniKanren syntax relation.

Properties
  :new-type  The type to be inserted, if any
  :start-index  Where the new tokens should start
  :end-index  Where the old tokens should be removed, if set
  :overrides-alist  Alist from type to function for overrides
"
  (tree-edit-load-grammar-for-major-mode)
  (if-let (override (alist-get (tree-edit--node-type parent) overrides-alist))
      (funcall override new-type parent start-index end-index)
    (-let* (reazon-occurs-check
            (parent-type (tree-edit--node-type parent))
            (grammar (alist-get parent-type (tree-edit-grammar)))
            (children (-map #'tree-edit--node-type (treesit-node-children parent)))
            (left (-take start-index children))
            (right (-drop (or end-index start-index) children))
            (relevant-types (tree-edit--relevant-types new-type parent-type)))
      (when-let (result (tree-edit--run-relation 3 q
                          (lambda (tokens)
                            (and
                             (listp tokens)
                             (if new-type
                                 (equal 1 (-count #'symbolp tokens))
                               (-all-p #'stringp tokens))))
                          (reazon-fresh (tokens qr ql)
                            (tree-edit--superpositiono right qr parent-type)
                            (tree-edit--superpositiono left ql parent-type)
                            (tree-edit--max-lengtho q 5)
                            (if new-type (tree-edit--includes-typeo q relevant-types) #'reazon-!S)
                            (tree-edit--prefixpostfixo ql q qr tokens)
                            (tree-edit-parseo grammar tokens '()))))
        (make-tree-edit-result
         :start-index start-index
         :end-index (or end-index start-index)
         :tokens (car result))))))

(defun tree-edit--valid-replacement-p (type node)
  "Return non-nil if NODE can be replaced with a node of TYPE."
  (let* ((parent (treesit-node-parent node))
         (index (--find-index (treesit-node-eq node it)
                              (treesit-node-children parent))))
    (tree-edit--attempt-structural-edit
     parent
     :new-type type
     :start-index index
     :end-index (1+ index)
     :overrides-alist (tree-edit-node-replacement-override))))

(defun tree-edit--valid-insertions (type node &optional before)
  "Return a valid sequence of tokens containing the provided TYPE, or nil.

If BEFORE is non-nil, generate the tokens after NODE, otherwise before."
  (let* ((parent (treesit-node-parent node))
         (index (--find-index (treesit-node-eq node it)
                              (treesit-node-children parent))))
    (tree-edit--attempt-structural-edit
     parent
     :new-type type
     :start-index (+ index (if before 0 1))
     :overrides-alist (tree-edit-node-insertion-override))))

(defun tree-edit--valid-deletions (node)
  "Return a set of edits if NODE can be deleted, else nil.

If successful, the return type will give a range of siblings to
delete, and what syntax needs to be inserted after, if any."
  (-let* (((children . index) (tree-edit--get-parent-tokens node))
          ((left-idx . right-idx) (tree-edit--remove-node-and-surrounding-syntax children index)))
    (tree-edit--attempt-structural-edit
     (treesit-node-parent node)
     :start-index left-idx
     :end-index right-idx
     :overrides-alist (tree-edit-node-deletion-override))))

;; NOTE: The node rendering code is pretty crappy, sorry.
;; I'll probably rewrite this whole thing later on.
;;
;;* Locals: node rendering
(defun tree-edit--generate-node (type &optional parent-type)
  "Given a NODE-TYPE and a set of RULES, generate a node string.

If TOKENS is passed in, that will be used as a basis for node
construction, instead of looking up the rules for node-type."
  (if (or (keywordp type) (stringp type)) `(,type)
    (--mapcat
     (tree-edit--add-whitespace-rules-to-tokens it parent-type (tree-edit--generate-node it type))
     (or (alist-get type (tree-edit-syntax-snippets))
         (user-error "No node definition for %s" type)))))

(defun tree-edit--needs-space-p (to-insert)
  "Check if the token TO-INSERT needs a space before it.

This function uses the point.

This function inserts and deletes into the buffer! This may be a
very bad idea, I'm not sure."
  (and (stringp to-insert)
       (let (result)
         (insert to-insert)
         (goto-char (- (point) (length to-insert)))
         (setq result
               (not (equal (treesit-node-start (treesit-node-at (point)))
                           (point))))
         (delete-char (length to-insert))
         result)))

(defun tree-edit--whitespace-rules-for-type (type parent-type)
  "Retrieve whitespace rules for TYPE.

Will search for the most specific rule first and travel through
the TYPE's supertypes until exhausted."
  (or (--any (alist-get it (alist-get parent-type (tree-edit-whitespace-rules)) nil nil #'equal)
             (alist-get type (tree-edit--supertypes) `(,type)))
      (--any (alist-get it (alist-get nil (tree-edit-whitespace-rules)) nil nil #'equal)
             (alist-get type (tree-edit--supertypes) `(,type)))))

(defun tree-edit--add-whitespace-rules-to-tokens (type parent-type tokens)
  "Wrap TOKENS in the whitespace defined for TYPE, if any."
  (if-let (ws (tree-edit--whitespace-rules-for-type type parent-type))
      (-let [(l r) ws] (append l tokens r))
    tokens))

(defun tree-edit--render-node (left-tokens new-tokens right-token indentation)
  "Insert NEW-TOKENS into the buffer, properly formatting as needed.

LEFT-TOKENS are used for calculating the formatting of
NEW-TOKENS, while RIGHT-TOKEN is used for adding any addition
newlines or spaces before the remaining text.

Pre-existing nodes in the tokens are assumed to be already
formatted correctly and thus are inserted as-is.

New nodes are inserted according `tree-edit-syntax-snippets'.

Text nodes (likely from the `kill-ring') are not assumed to be
formatted correctly and thus decomposed by
`tree-edit--split-node-for-insertion' into chunks where formatting
matters (i.e. expressions are left alone but blocks are split)."
  (-let* ((prev nil)
          (deferred-newline nil))
    (cl-flet ((process-tokens
                (stack do-insert do-whitespace)
                (while stack
                  (-let ((current (pop stack)))
                    ;; TODO: use `pcase'
                    (cond
                     ((equal current :newline)
                      (setq deferred-newline t))
                     ((equal current :indent)
                      (setq indentation (+ indentation (tree-edit-indentation-level))))
                     ((equal current :dedent)
                      (setq indentation (- indentation (tree-edit-indentation-level))))
                     ((stringp current)
                      (when deferred-newline
                        (when (or do-insert do-whitespace)
                          (newline)
                          (indent-line-to indentation))
                        (setq deferred-newline nil))
                      (when (and (or do-insert do-whitespace)
                                 (tree-edit--needs-space-p current))
                        (insert " "))
                      (when do-insert
                        (insert current)))
                     (t (error "bad data: %s" current)))
                    (when (or (equal :newline current) (stringp current))
                      (setq prev current))))))
      (process-tokens left-tokens nil nil)
      (process-tokens new-tokens t nil)
      (process-tokens right-token nil t))))

(defun tree-edit--render-result (result parent)
  "Convert RESULT into a textual representation by editing the children of PARENT."
  (-let* ((start-index (tree-edit-result-start-index result))
          (end-index (tree-edit-result-end-index result))
          (fragment (tree-edit-result-tokens result))
          (children (treesit-node-children parent))
          (left (--mapcat (tree-edit--split-node-for-insertion
                           it
                           (tree-edit--node-type parent))
                          (-slice children 0 start-index)))
          (right (-some--> end-index
                   (nth it children)
                   (tree-edit--add-whitespace-rules-to-tokens
                    (tree-edit--node-type it)
                    (tree-edit--node-type parent)
                    `(,(treesit-node-text it)))))
          ;; FIXME
          (render-fragment
           fragment)
          (indentation
           (save-excursion
             (if children (goto-char (treesit-node-start (car children))))
             (current-indentation)))
          (start-edit
           (if (zerop start-index)
               (treesit-node-start parent)
             (treesit-node-end (nth (1- start-index) children))))
          (end-edit
           (if (nth end-index children)
               (treesit-node-start (nth end-index children))
             (treesit-node-end parent))))
    (combine-change-calls
        start-edit
        end-edit
      (save-excursion
        (goto-char start-edit)
        (delete-region start-edit end-edit)
        (tree-edit--render-node left (if fragment render-fragment) right indentation)))))

(cl-defun tree-edit--split-node-for-insertion (node &optional _parent-type skip-ws)
  "Split NODE into chunks of text as necessary for formatting."
  (let* ((type (tree-edit--node-type node))
         (rules (tree-edit--whitespace-rules-for-type type _parent-type)))
    (let ((result
           (if (or (zerop (treesit-node-child-count node :named)) (not rules))
               (with-current-buffer (or tree-edit--string-parse-buffer
                                        (current-buffer))
                 `(,(treesit-node-text node)))
             (--mapcat (tree-edit--split-node-for-insertion it type)
                       (treesit-node-children node)))))
      (if skip-ws result
        (tree-edit--add-whitespace-rules-to-tokens type _parent-type result)))))

;;* Globals: Structural editing functions
(defun tree-edit-exchange (new-node node)
  "Attempt to exchange NODE for NEW-NODE.

If NEW-NODE is a string, the tree-edit will attempt to infer the type of
the text."
  (when (treesit-node-eq node (treesit-buffer-root-node))
    (tree-edit-transformation-error "Cannot exchange the root node!"))
  (if-let (result (tree-edit--try-transformation
                   new-node
                   (-> node (treesit-node-parent) (tree-edit--node-type))
                   (lambda (type) (tree-edit--valid-replacement-p type node))))
      (tree-edit--render-result result (treesit-node-parent node))
    (tree-edit-transformation-error "Cannot replace %s with %s!" (tree-edit--node-type node) new-node)))

(defun tree-edit-raise (node)
  "Move NODE up the syntax tree until a valid replacement is found."
  (let ((ancestor-to-replace (tree-edit--find-raise-ancestor (treesit-node-parent node) node)))
    (let ((node-text (treesit-node-text node))
          (ancestor-steps (tree-edit--node-steps ancestor-to-replace)))
      (tree-edit-cache-node node)
      ;; FIXME: Relational parser is being run twice
      (tree-edit-exchange node-text ancestor-to-replace)
      (tree-edit--node-from-steps ancestor-steps))))

(defun tree-edit-insert-sibling (new-node node &optional before)
  "Attempt to insert NEW-NODE adjacent to NODE.

If NEW-NODE is a string, the tree-edit will attempt to infer the type of
the text.

if BEFORE is t, the sibling node will be inserted before NODE, else after."
  (when (treesit-node-eq node (treesit-buffer-root-node))
    (tree-edit-transformation-error "Cannot perform insertions on the root node!"))
  (if-let (result (tree-edit--try-transformation
                   new-node
                   (-> node (treesit-node-parent) (tree-edit--node-type))
                   (lambda (type) (tree-edit--valid-insertions type node before))))
      (tree-edit--render-result result (treesit-node-parent node))
    (tree-edit-transformation-error "Cannot insert %s %s %s!" new-node
                                    (if before "before" "after") (tree-edit--node-type node))))

(defun tree-edit--post-process-tokens (parent-type type replace-with result)
  "Replace alternate representations of TYPE in result with TYPE itself.

For example, an identifier type may be represented by it's
supertype (expression), so we need to replace it with the
original."
  (setf (tree-edit-result-tokens result)
        ;; TODO: Check supertype
        (->> result
             (tree-edit-result-tokens)
             (--mapcat (tree-edit--add-whitespace-rules-to-tokens type parent-type (if (symbolp it) replace-with `(,it))))))
  result)

(defun tree-edit--try-transformation (new-node parent-type pred)
  "Run PRED on NEW-NODE and return tokens if valid.

If NEW-NODE is a symbol, the symbol will be passed directly to
the predicate.

If NEW-NODE is a list (presumably of symbols), the types will be
tried in order.

If NEW-NODE is a string, the type cache will be used if an entry
exists. Otherwise, the string will be parsed by the tree-sitter
parser."
  (cond
   ((symbolp new-node)
    (-some->> new-node
      (funcall pred)
      (tree-edit--post-process-tokens parent-type new-node (tree-edit--generate-node new-node))))
   ((listp new-node)
    (cl-dolist (type new-node)
      (-some->> type
        (funcall pred)
        (tree-edit--post-process-tokens parent-type type (tree-edit--generate-node type))
        (cl-return))))
   ((stringp new-node)
    (cl-block nil
      (if-let ((cached-node (gethash new-node tree-edit--type-cache)))
          (-let [(type . split-node) cached-node]
            ;; If the cached type doesn't match, we'll try to reparse it.
            (-some->> type
              (funcall pred)
              (tree-edit--post-process-tokens parent-type type split-node)
              (cl-return))))
      (let* ((parser-lang
              (thread-last
                (treesit-buffer-root-node)
                (treesit-node-parser)
                (treesit-parser-language)))
             (tree-edit--string-parse-buffer (get-buffer-create "*tree-edit string parse*"))
             ;; Parse the node fragment in a temporary buffer.
             (temp-node
              (with-current-buffer tree-edit--string-parse-buffer
                (erase-buffer)
                (insert new-node)
                (treesit-parser-create parser-lang)
                (treesit-buffer-root-node))))

        (dolist (fragment-node (tree-edit--parse-fragment temp-node))
          (if-let (result (funcall pred (tree-edit--node-type fragment-node)))
              (cl-return
               (tree-edit--post-process-tokens
                parent-type
                (tree-edit--node-type fragment-node)
                (tree-edit--split-node-for-insertion fragment-node)
                result)))))))
   (t (user-error "Bad data: %s" new-node))))

(defun tree-edit-insert-child (new-node node &optional idx)
  "Attempt to insert NEW-NODE inside of NODE.

If NODE already has named children, the new node will be inserted
before the first child.

If NEW-NODE is a string, the tree-edit will attempt to infer the
type of the text."
  (if (> (treesit-node-child-count node :named) 0)
      (tree-edit-insert-sibling new-node (treesit-node-child node (or idx 0) :named) t)
    (if-let (result (tree-edit--try-transformation
                     new-node
                     (tree-edit--node-type node)
                     (lambda (type)
                       (when-let ((tokens (tree-edit--valid-node-including-type type (tree-edit--node-type node))))
                         (make-tree-edit-result :tokens tokens
                                                :start-index 0
                                                :end-index (treesit-node-child-count node))))))
        (tree-edit--render-result result node)
      (tree-edit-transformation-error "Cannot insert %s into %s!" new-node (tree-edit--node-type node)))))

(defun tree-edit--get-next-node (node)
  "Get the next node to the left of NODE."
  (when-let ((parent (treesit-node-parent node)))
    (or (treesit-node-next-sibling node :named)
        (tree-edit--get-next-node parent))))

(defun tree-edit-slurp (node)
  "Transform NODE's next sibling into it's leftmost child, if possible."
  (let ((slurp-candidate (tree-edit--get-next-node node)))
    (cond ((not slurp-candidate)
           (tree-edit-transformation-error "Nothing to slurp!"))
          ((zerop (treesit-node-child-count node))
           (tree-edit-transformation-error "Current node has no children, can't slurp!"))
          ;; No named children, use insert child
          ((equal (treesit-node-child-count node :named) 0)
           (let ((slurper (tree-edit--node-steps node))
                 (slurp-text (treesit-node-text slurp-candidate)))
             (atomic-change-group
               (tree-edit-delete slurp-candidate)
               (tree-edit-insert-child slurp-text (tree-edit--node-from-steps-strict slurper)))))
          ;; Named children, use insert sibling
          (t
           (let* ((slurper
                   (treesit-node-child node (1- (treesit-node-child-count node :named)) :named))
                  (slurper-steps (tree-edit--node-steps slurper))
                  (slurp-text (treesit-node-text slurp-candidate)))
             (atomic-change-group
               (tree-edit-delete slurp-candidate)
               ;; TODO: Propogate known type to slurp-text
               (tree-edit-insert-sibling slurp-text
                                         (tree-edit--node-from-steps-strict slurper-steps))))))))

(defun tree-edit-barf (node)
  "Transform NODE's leftmost child into it's next sibling, if possible."
  (unless (> (treesit-node-child-count node :named) 0)
    (tree-edit-transformation-error "Cannot barf a node with no named children!"))
  (let ((barfer (treesit-node-parent node))
        (barfee (treesit-node-child node (1- (treesit-node-child-count node :named)) :named)))
    (unless (tree-edit--valid-deletions barfee)
      (tree-edit-transformation-error "Cannot delete %s!" (treesit-node-text barfee)))
    (cl-block nil
      (while barfer
        (let* ((barfer-steps (tree-edit--node-steps barfer))
               (barfee-text (treesit-node-text barfee)))
          (when (tree-edit--valid-insertions (tree-edit--node-type barfer)
                                             (treesit-node-child node 0 :named))
            (tree-edit-delete barfee)
            (tree-edit-insert-sibling barfee-text (tree-edit--node-from-steps-strict barfer-steps))
            (cl-return))
          (setq barfer (treesit-node-parent barfer))))
      (tree-edit-transformation-error "Cannot barf %s!" (tree-edit--node-type node)))))

(defun tree-edit-delete (node)
  "Delete NODE, and any surrounding syntax that accompanies it."
  (when (treesit-node-eq node (treesit-buffer-root-node))
    (tree-edit-transformation-error "Cannot delete the root node!"))
  (if-let (result (tree-edit--valid-deletions node))
      (tree-edit--render-result result (treesit-node-parent node))
    (tree-edit-transformation-error "Cannot delete the current node")))

(defun tree-edit-cache-node (node)
  "Store a mapping from NODE's text to type."
  (interactive)
  (puthash (treesit-node-text node)
           (cons (tree-edit--node-type node)
                 (tree-edit--split-node-for-insertion node nil t))
           tree-edit--type-cache))

(defun tree-edit-copy (node)
  "Copy NODE and cache it's type."
  (interactive)
  (tree-edit-cache-node node)
  (kill-ring-save (treesit-node-start node) (treesit-node-end node)))

;;* Locals: Relational parser
;; Upstream this to `reazon'?
(defun tree-edit--take-first (stream pred tries var)
  "Return the first item from STREAM for which PRED returns non-nil.

Modified from `reazon--take'."
  (declare (indent 1))
  (cond
   ((funcall pred (funcall (reazon--reify var) (car stream)))
    `(,(funcall (reazon--reify var) (car stream))))
   ((or (equal 1 tries) (functionp stream) (null stream)) nil)
   (t
    (cl-block nil
      (let ((count (if tries (1- tries) -1))
            (stream (reazon--pull (cdr stream))))
        (while (and stream (not (zerop count)) (not (functionp stream)))
          (let ((reified-var (funcall (reazon--reify var) (car stream))))
            (when (funcall pred reified-var)
              (cl-return `(,reified-var))))
          (setq count (1- count))
          (setq stream (reazon--pull (cdr stream)))))))))

(reazon-defrel tree-edit-parseo (grammar tokens out)
  "TOKENS are a valid prefix of a node in GRAMMAR and OUT is unused tokens in TOKENS."
  (reazon-disj
   (if tree-edit-parse-comments
       (reazon-fresh (next)
         ;; FIXME: Should use the `extras` field from grammar
         (tree-edit--takeo 'comment tokens next)
         (tree-edit-parseo grammar next out))
     #'reazon-!U)
   (pcase grammar
     (`((type . "STRING")
        (value . ,value))
      (tree-edit--takeo value tokens out))
     (`((type . "PATTERN")
        (value . ,_))
      (tree-edit--takeo :regex tokens out))
     (`((type . "BLANK"))
      (reazon-== tokens out))
     ((and `((type . ,type)
             (value . ,_)
             (content . ,content))
           (guard (s-starts-with-p "PREC" type)))
      ;; Silence the foolish linter.
      (ignore type)
      (tree-edit-parseo content tokens out))
     (`((type . "TOKEN")
        (content . ,content))
      (tree-edit-parseo content tokens out))
     (`((type . "SEQ")
        (members . ,members))
      (tree-edit--seqo members tokens out))
     (`((type . "ALIAS")
        (content . ,_)
        (named . ,named)
        (value . ,alias-name))
      (tree-edit--takeo (if (eq named :json-false) (symbol-name alias-name) alias-name) tokens out))
     (`((type . "IMMEDIATE_TOKEN")
        (content . ,content))
      (tree-edit-parseo content tokens out))
     (`((type . "REPEAT")
        (content . ,content))
      (tree-edit--repeato content tokens out))
     (`((type . "REPEAT1")
        (content . ,content))
      (tree-edit--repeat1o content tokens out))
     (`((type . "FIELD")
        (name . ,_)
        (content . ,content))
      (tree-edit-parseo content tokens out))
     (`((type . "SYMBOL")
        (name . ,name))
      (tree-edit--takeo name tokens out))
     (`((type . "CHOICE")
        (members . ,members))
      (tree-edit--choiceo members tokens out))
     (_ (error "Bad data: %s" grammar)))))

(reazon-defrel tree-edit--max-lengtho (ls len)
  "LS contains at most LEN elements."
  (cond
   ((> len 0)
    (reazon-disj
     (reazon-nullo ls)
     (reazon-fresh (d)
       (reazon-cdro ls d)
       (tree-edit--max-lengtho d (1- len)))))
   (t (reazon-nullo ls))))

(reazon-defrel tree-edit--seqo (members tokens out)
  "TOKENS parse sequentially for each grammar in MEMBERS, with OUT as leftovers."
  (if members
      (reazon-fresh (next)
        (tree-edit-parseo (car members) tokens next)
        (tree-edit--seqo (cdr members) next out))
    (reazon-== tokens out)))

(reazon-defrel tree-edit--choiceo (members tokens out)
  "TOKENS parse for each grammar in MEMBERS, with OUT as leftovers."
  (if members
      (reazon-disj
       (tree-edit-parseo (car members) tokens out)
       (tree-edit--choiceo (cdr members) tokens out))
    #'reazon-!U))

(reazon-defrel tree-edit--repeato (grammar tokens out)
  "TOKENS parse for GRAMMAR an abritrary amount of times, with OUT as leftovers."
  (reazon-disj
   (reazon-== tokens out)
   (reazon-fresh (next)
     (tree-edit-parseo grammar tokens next)
     (tree-edit--repeato grammar next out))))

(reazon-defrel tree-edit--repeat1o (grammar tokens out)
  "TOKENS parse for GRAMMAR at least once, up to an abritrary amount of times, with OUT as leftovers."
  (reazon-fresh (next)
    (tree-edit-parseo grammar tokens next)
    (tree-edit--repeato grammar next out)))

(reazon-defrel tree-edit--takeo (expected tokens out)
  "TOKENS is a cons, with car as EXPECTED and cdr as OUT."
  (reazon-conso expected out tokens))

(reazon-defrel tree-edit--prefixpostfixo (prefix middle postfix out)
  "OUT is equivalent to (append PREFIX MIDDLE POSTFIX)."
  (reazon-fresh (tmp)
    (reazon-appendo prefix middle tmp)
    (reazon-appendo tmp postfix out)))

(reazon-defrel tree-edit--includes-typeo (tokens relevant-types)
  "One of the types in RELEVANT-TYPES appears in TOKENS."
  (reazon-fresh (a d)
    (reazon-conso a d tokens)
    (reazon-disj
     (reazon-membero a relevant-types)
     (tree-edit--includes-typeo d relevant-types))))

(reazon-defrel tree-edit--superpositiono (tokens out parent-type)
  "OUT is TOKENS where each token is either itself or any relevant type occurring in PARENT-TYPE."
  (cond
   ((not tokens) (reazon-== out '()))
   ((and (not (equal (car tokens) 'comment)) (symbolp (car tokens)))
    (reazon-fresh (a d)
      (reazon-conso a d out)
      (reazon-membero a (tree-edit--relevant-types (car tokens) parent-type))
      (tree-edit--superpositiono (cdr tokens) d parent-type)))
   (t
    (reazon-fresh (a d)
      (reazon-conso a d out)
      (reazon-== a (car tokens))
      (tree-edit--superpositiono (cdr tokens) d parent-type)))))

(provide 'tree-edit)
;;; tree-edit.el ends here