From efa282c2aceefc3d28c530081b9802590861b509 Mon Sep 17 00:00:00 2001
From: Yuan Fu <casouri@gmail.com>
Date: Sat, 20 Apr 2024 00:42:25 -0700
Subject: [PATCH] Add treesit-add-font-lock-rules

This function should help users customize font-lock rules.

* lisp/treesit.el (treesit-add-font-lock-rules): New function.

(cherry picked from commit bf5755aa91905a376e4ac9fedde6a645486e1589)
---
 etc/NEWS        | 11 ++++++++++-
 lisp/treesit.el | 48 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/etc/NEWS b/etc/NEWS
index 549bbe01926..be9e7dfbf9f 100644
--- a/etc/NEWS
+++ b/etc/NEWS
@@ -2689,7 +2689,16 @@ parsers with that tag.  Note that passing nil as tag doesn't mean return
 all parsers, but rather "all parsers with no tags".
 
 +++
-*** New function 'treesit-parser-changed-ranges' which returns buffer regions that are affected by the last buffer edits
+*** New function 'treesit-parser-changed-ranges'.
+
+This function returns buffer regions that are affected by the last
+buffer edits.
+
+*** New function 'treesit-add-font-lock-rules'.
+
+This function helps user to add custom font-lock rules to a tree-sitter
+major mode.
+
 
 
 * Changes in Emacs 30.1 on Non-Free Operating Systems
diff --git a/lisp/treesit.el b/lisp/treesit.el
index 4c9cf5a36e7..e55e04e53b3 100644
--- a/lisp/treesit.el
+++ b/lisp/treesit.el
@@ -817,6 +817,17 @@ OVERRIDE is the override flag for this query.  Its value can be
 t, nil, append, prepend, keep.  See more in
 `treesit-font-lock-rules'.")
 
+(defsubst treesit--font-lock-setting-feature (setting)
+  "Reutrn the feature of SETTING.
+SETTING should be a setting in `treesit-font-lock-settings'."
+  (nth 2 setting))
+
+(defsubst treesit--font-lock-setting-enable (setting)
+  "Return enabled SETTING."
+  (let ((new-setting (copy-tree setting)))
+    (setf (nth 1 new-setting) t)
+    new-setting))
+
 (defun treesit--font-lock-level-setter (sym val)
   "Custom setter for `treesit-font-lock-level'.
 Set the default value of SYM to VAL, recompute fontification
@@ -1095,6 +1106,43 @@ and leave settings for other languages unchanged."
                        ((memq feature remove-list) nil)
                        (t current-value))))))
 
+(defun treesit-add-font-lock-rules (rules &optional how feature)
+  "Add font-lock RULES to the current buffer
+
+RULES should be the return value of `treesit-font-lock-rules'.  RULES
+will be enabled and added to `treesit-font-lock-settings'.
+
+HOW can be either :before or :after.  If HOW is :before, prepend RULES
+before all other existing font-lock rules in
+`treesit-font-lock-settings'; if :after or omitted, append RULES after
+all existing rules.
+
+If FEATURE is non-nil, add RULES before/after rules for FEATURE.  See
+docstring of `treesit-font-lock-rules' for what is a feature."
+  (let ((rules (seq-map #'treesit--font-lock-setting-enable rules))
+        (feature-idx
+         (when feature
+           (cl-position-if
+            (lambda (setting)
+              (eq (treesit--font-lock-setting-feature setting) feature))
+            treesit-font-lock-settings))))
+    (pcase (cons how feature)
+      ((or '(:after . nil) '(nil . nil))
+       (setq treesit-font-lock-settings
+             (append treesit-font-lock-settings rules)))
+      ('(:before . nil)
+       (setq treesit-font-lock-settings
+             (append rules treesit-font-lock-settings)))
+      (`(:after . ,_feature)
+       (setf (nthcdr (1+ feature-idx) treesit-font-lock-settings)
+             (append rules
+                     (nthcdr (1+ feature-idx)
+                             treesit-font-lock-settings))))
+      (`(:before . ,_feature)
+       (setf (nthcdr feature-idx treesit-font-lock-settings)
+             (append rules
+                     (nthcdr feature-idx treesit-font-lock-settings)))))))
+
 (defun treesit-fontify-with-override
     (start end face override &optional bound-start bound-end)
   "Apply FACE to the region between START and END.
-- 
2.39.5