From 76eb88df0243dd15492e71f068fd476be8efc800 Mon Sep 17 00:00:00 2001
From: Yuan Fu <casouri@gmail.com>
Date: Sun, 20 Nov 2022 17:15:42 -0800
Subject: [PATCH] Add another heuristic to speed up tree-sitter fontification

* lisp/treesit.el (treesit-font-lock-fontify-region): Use
treesit--children-covering-range to reduce the node size.
---
 lisp/treesit.el | 68 ++++++++++++++++++++++++++++---------------------
 1 file changed, 39 insertions(+), 29 deletions(-)

diff --git a/lisp/treesit.el b/lisp/treesit.el
index d65ad31b473..f9f807947c2 100644
--- a/lisp/treesit.el
+++ b/lisp/treesit.el
@@ -816,41 +816,51 @@ If LOUDLY is non-nil, display some debugging information."
                  ;; Only activate if ENABLE flag is t.
                  (activate (eq t enable)))
         (ignore activate)
+
         ;; The node seems small, enlarge it.
         (while (and (< (- (treesit-node-end node-on)
                           (treesit-node-start node-on))
                        40)
                     (treesit-node-parent node-on))
           (setq node-on (treesit-node-parent node-on)))
-        (let ((captures (treesit-query-capture
-                         node-on query start end))
-              (inhibit-point-motion-hooks t))
-          (with-silent-modifications
-            (dolist (capture captures)
-              (let* ((face (car capture))
-                     (node (cdr capture))
-                     (node-start (treesit-node-start node))
-                     (node-end (treesit-node-end node)))
-                ;; Turns out it is possible to capture a node that's
-                ;; completely outside the region between START and
-                ;; END.  If the node is outside of that region, (max
-                ;; node-start start) and friends return bad values.
-                (when (and (< start node-end)
-                           (< node-start end))
-                  (cond
-                   ((facep face)
-                    (treesit-fontify-with-override
-                     (max node-start start) (min node-end end)
-                     face override))
-                   ((functionp face)
-                    (funcall face node override start end)))
-                  ;; Don't raise an error if FACE is neither a face nor
-                  ;; a function.  This is to allow intermediate capture
-                  ;; names used for #match and #eq.
-                  (when (or loudly treesit--font-lock-verbose)
-                    (message "Fontifying text from %d to %d, Face: %s, Node: %s"
-                             (max node-start start) (min node-end end)
-                             face (treesit-node-type node)))))))))))
+
+        ;; Maybe the node returned by `treesit-node-on' is the root
+        ;; node, because the region between START and END contains
+        ;; several top-level constructs (e.g., variable declarations
+        ;; in C).
+        (setq node-on (treesit--children-covering-range
+                       node-on start end))
+
+        (dolist (sub-node node-on)
+          (let ((captures (treesit-query-capture
+                           sub-node query start end))
+                (inhibit-point-motion-hooks t))
+            (with-silent-modifications
+              (dolist (capture captures)
+                (let* ((face (car capture))
+                       (node (cdr capture))
+                       (node-start (treesit-node-start node))
+                       (node-end (treesit-node-end node)))
+                  ;; Turns out it is possible to capture a node that's
+                  ;; completely outside the region between START and
+                  ;; END.  If the node is outside of that region, (max
+                  ;; node-start start) and friends return bad values.
+                  (when (and (< start node-end)
+                             (< node-start end))
+                    (cond
+                     ((facep face)
+                      (treesit-fontify-with-override
+                       (max node-start start) (min node-end end)
+                       face override))
+                     ((functionp face)
+                      (funcall face node override start end)))
+                    ;; Don't raise an error if FACE is neither a face nor
+                    ;; a function.  This is to allow intermediate capture
+                    ;; names used for #match and #eq.
+                    (when (or loudly treesit--font-lock-verbose)
+                      (message "Fontifying text from %d to %d, Face: %s, Node: %s"
+                               (max node-start start) (min node-end end)
+                               face (treesit-node-type node))))))))))))
   `(jit-lock-bounds ,start . ,end))
 
 (defun treesit--font-lock-notifier (ranges parser)
-- 
2.39.5