From ef28af35bb4c43d71fe4c10d02fe93f30e830c5e Mon Sep 17 00:00:00 2001
From: Yuan Fu <casouri@gmail.com>
Date: Thu, 30 Jan 2025 16:54:20 -0800
Subject: [PATCH] Add treesit-add-simple-indent-rules

* lisp/treesit.el:
(treesit-add-simple-indent-rules): New function.
* etc/NEWS: Update NEWS.
*
test/src/treesit-tests.el:
(treesit-test-add-simple-indent-rules): Add a test for the new
function.
---
 etc/NEWS                  |  4 ++++
 lisp/treesit.el           | 29 +++++++++++++++++++++++++++++
 test/src/treesit-tests.el | 21 +++++++++++++++++++++
 3 files changed, 54 insertions(+)

diff --git a/etc/NEWS b/etc/NEWS
index 1e7365258a3..9d0eb7b7d58 100644
--- a/etc/NEWS
+++ b/etc/NEWS
@@ -1159,6 +1159,10 @@ at point to explore.
 *** New variable 'treesit-aggregated-simple-imenu-settings'
 This variable allows major modes to setup Imenu for multiple languages.
 
+*** New function 'treesit-add-simple-indent-rules'
+This new function makes it easier to customize indent rules for
+tree-sitter modes.
+
 +++
 ** New optional BUFFER argument for 'string-pixel-width'.
 If supplied, 'string-pixel-width' will use any face remappings from
diff --git a/lisp/treesit.el b/lisp/treesit.el
index 76b6bdc9cb2..1940dd37795 100644
--- a/lisp/treesit.el
+++ b/lisp/treesit.el
@@ -2401,6 +2401,35 @@ RULES."
                             offset)))))
              (cons lang (mapcar #'optimize-rule indent-rules)))))
 
+(defun treesit-add-simple-indent-rules (language rules &optional where anchor)
+  "Add simple indent RULES for LANGUAGE.
+
+WHERE can be either :before or :after, which means adding RULES before
+or after the existing rules in `treesit-simple-indent-rules'.  If
+ommited, default to adding the rules before (so it overrides existing
+rules).
+
+If ANCHOR is non-nil, add RULES before/after the rules in
+`treesit-simple-indent-rules' that's `equal' to ANCHOR.  If ANCHOR is
+omitted or no existing rules matches it, add RULES at the beginning or
+end of existing rules."
+  (when (not (memq where '(nil :before :after)))
+    (error "WHERE must be either :before, :after, or nil"))
+  (let* ((existing-rules (alist-get language treesit-simple-indent-rules))
+         (anchor-idx (and anchor (seq-position existing-rules anchor)))
+         (new-rules
+          (if anchor-idx
+              (let* ((pivot (if (eq where :after)
+                                (1+ anchor-idx)
+                              anchor-idx))
+                     (first-half (seq-subseq existing-rules 0 pivot))
+                     (second-half (seq-subseq existing-rules pivot)))
+                (append first-half rules second-half))
+            (if (eq where :after)
+                (append existing-rules rules)
+              (append rules existing-rules)))))
+    (setf (alist-get language treesit-simple-indent-rules) new-rules)))
+
 ;;; Search
 
 (defun treesit-search-forward-goto
diff --git a/test/src/treesit-tests.el b/test/src/treesit-tests.el
index f9c64f792d1..22f53243274 100644
--- a/test/src/treesit-tests.el
+++ b/test/src/treesit-tests.el
@@ -411,6 +411,27 @@ BODY is the test body."
    (let ((missing-bracket (treesit-node-child array -1)))
      (treesit-search-forward missing-bracket "" t))))
 
+;;; Indent
+
+(ert-deftest treesit-test-add-simple-indent-rules ()
+  "Test `treesit-add-simple-indent-rules'."
+  (let ((treesit-simple-indent-rules
+         (copy-tree '((c (a a a) (b b b) (c c c))))))
+    (treesit-add-simple-indent-rules 'c '((d d d)))
+    (should (equal treesit-simple-indent-rules
+                   '((c (d d d) (a a a) (b b b) (c c c)))))
+    (treesit-add-simple-indent-rules 'c '((e e e)) :after)
+    (should (equal treesit-simple-indent-rules
+                   '((c (d d d) (a a a) (b b b) (c c c) (e e e)))))
+    (treesit-add-simple-indent-rules 'c '((f f f)) :after '(b b b))
+    (should (equal treesit-simple-indent-rules
+                   '((c (d d d) (a a a) (b b b) (f f f)
+                        (c c c) (e e e)))))
+    (treesit-add-simple-indent-rules 'c '((g g g)) :before '(b b b))
+    (should (equal treesit-simple-indent-rules
+                   '((c (d d d) (a a a) (g g g)
+                        (b b b) (f f f) (c c c) (e e e)))))))
+
 ;;; Query
 
 (defun treesit--ert-pred-last-sibling (node)