From 5a93a5c4232a6aed6423cb789e006d50e37478bb Mon Sep 17 00:00:00 2001
From: Yuan Fu <casouri@gmail.com>
Date: Sat, 7 Dec 2024 23:18:53 -0800
Subject: [PATCH] Allow treesit--font-lock-level-setter to be an alist

* lisp/treesit.el (treesit-font-lock-level): Allow the
value to be an alist mapping major modes to font lock levels.
(treesit--font-lock-level-setter): Rewrite.
(treesit--compute-font-lock-level): New function.
(treesit-font-lock-recompute-features): Use new function.

* doc/lispref/modes.texi (Parser-based Font Lock): Minor fix.

(cherry picked from commit b953bcb17047998c9e41cede7c5e5ffec22209b2)
---
 doc/lispref/modes.texi |  2 +-
 lisp/treesit.el        | 56 ++++++++++++++++++++++++++++--------------
 2 files changed, 39 insertions(+), 19 deletions(-)

diff --git a/doc/lispref/modes.texi b/doc/lispref/modes.texi
index 34dcf4cb10e..bede5bf2478 100644
--- a/doc/lispref/modes.texi
+++ b/doc/lispref/modes.texi
@@ -4379,7 +4379,7 @@ ignored.
 @defvar treesit-font-lock-feature-list
 This is a list of lists of feature symbols.  Each element of the list
 is a list that represents a decoration level.
-@code{treesit-font-lock-level} controls which levels are
+@var{treesit-font-lock-level} controls which levels are
 activated.
 
 Each element of the list is a list of the form @w{@code{(@var{feature}
diff --git a/lisp/treesit.el b/lisp/treesit.el
index 221c539f86a..10e0e9f885d 100644
--- a/lisp/treesit.el
+++ b/lisp/treesit.el
@@ -897,29 +897,22 @@ t, nil, append, prepend, keep.  See more in
     (setf (nth 1 new-setting) t)
     new-setting))
 
-;; FIXME: Rewrite this in more readable fashion.
 (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
 features and refontify for every buffer where tree-sitter-based
 fontification is enabled."
   (set-default sym val)
-  (and (treesit-available-p)
-       (named-let loop ((res nil)
-                        (buffers (buffer-list)))
-         (if (null buffers)
-             (mapc (lambda (b)
-                     (with-current-buffer b
-                       (setq-local treesit-font-lock-level val)
-                       (treesit-font-lock-recompute-features)
-                       (treesit-font-lock-fontify-region (point-min)
-                                                         (point-max))))
-                   res)
-           (let ((buffer (car buffers)))
-             (with-current-buffer buffer
-               (if treesit-font-lock-settings
-                   (loop (append res (list buffer)) (cdr buffers))
-                 (loop res (cdr buffers)))))))))
+  (when (treesit-available-p)
+    (dolist (buffer (buffer-list))
+      (with-current-buffer buffer
+        ;; FIXME: This doesn't re-run major mode hooks, meaning any
+        ;; customization done in major mode hooks (e.g., with
+        ;; `treesit-font-lock-recompute-features') is lost.
+        (when treesit-font-lock-settings
+          (treesit-font-lock-recompute-features)
+          (treesit-font-lock-fontify-region
+           (point-min) (point-max)))))))
 
 (defcustom treesit-font-lock-level 3
   "Decoration level to be used by tree-sitter fontifications.
@@ -936,6 +929,15 @@ Level 4 adds everything else that can be fontified: delimiters,
 operators, brackets, punctuation, all functions, properties,
 variables, etc.
 
+The value of this variable can be either a number representing a level,
+or an alist of (MAJOR-MODE . LEVEL), where MAJOR-MODE is major mode
+symbols, or t (meaning the default), and LEVEL is the font-lock level
+for that mode.  For example,
+
+    ((c-ts-mode . 3) (c++-ts-mode . 4) (t . 3))
+
+Major mode is checked with `derived-mode-p'.
+
 In addition to the decoration level, individual features can be
 turned on/off by calling `treesit-font-lock-recompute-features'.
 Changing the decoration level requires calling
@@ -1119,6 +1121,23 @@ name, it is ignored."
 (defvar treesit--font-lock-verbose nil
   "If non-nil, print debug messages when fontifying.")
 
+(defun treesit--compute-font-lock-level (level)
+  "Compute the font-lock level for the current major mode.
+
+LEVEL should be the value of `treesit-font-lock-level'.  Return a number
+representing the font-lock level for the current major mode.  If there's
+no match, return 3."
+  (if (numberp level)
+      level
+    (catch 'found
+      (dolist (config level)
+        (let ((mode (car config))
+              (num (cdr config)))
+          (when (derived-mode-p mode)
+            (throw 'found num))))
+      (or (alist-get t level)
+          3))))
+
 (defun treesit-font-lock-recompute-features
     (&optional add-list remove-list language)
   "Enable/disable font-lock features.
@@ -1143,7 +1162,8 @@ and leave settings for other languages unchanged."
     (signal 'treesit-font-lock-error
             (list "ADD-LIST and REMOVE-LIST contain the same feature"
                   intersection)))
-  (let* ((level treesit-font-lock-level)
+  (let* ((level (treesit--compute-font-lock-level
+                 treesit-font-lock-level))
          (base-features (cl-loop
                          for idx = 0 then (1+ idx)
                          for features in treesit-font-lock-feature-list
-- 
2.39.5