From 1970726e26a979243925fabe32686ba2ee757c6b Mon Sep 17 00:00:00 2001
From: Yuan Fu <casouri@gmail.com>
Date: Sun, 29 Jan 2023 00:07:46 -0800
Subject: [PATCH] Use treesit-subtree-stat to determine
 treesit--font-lock-fast-mode

* lisp/treesit.el:
(treesit--children-covering-range-recurse): Return nil if LIMIT is
exceeded.
(treesit--font-lock-fast-mode): Change to a ternary value.
(treesit-font-lock-fontify-region): Enable fast mode based on the
result of treesit-subtree-stat.
---
 lisp/treesit.el | 50 +++++++++++++++++++++++++------------------------
 1 file changed, 26 insertions(+), 24 deletions(-)

diff --git a/lisp/treesit.el b/lisp/treesit.el
index b3029707376..f4976502fc1 100644
--- a/lisp/treesit.el
+++ b/lisp/treesit.el
@@ -913,14 +913,15 @@ LIMIT is the recursion limit, which defaults to 100."
             (push r result))
         (push child result))
       (setq child (treesit-node-next-sibling child)))
-    ;; If NODE has no child, keep NODE.
-    (or result (list node))))
+    ;; If NODE has no child, keep NODE.  If LIMIT is exceeded, return
+    ;; nil.
+    (or result (and (> limit 0) (list node)))))
 
 (defsubst treesit--node-length (node)
   "Return the length of the text of NODE."
   (- (treesit-node-end node) (treesit-node-start node)))
 
-(defvar-local treesit--font-lock-fast-mode nil
+(defvar-local treesit--font-lock-fast-mode 'unspecified
   "If this variable is t, change the way we query so it's faster.
 This is not a general optimization and should be RARELY needed!
 See comments in `treesit-font-lock-fontify-region' for more
@@ -985,36 +986,34 @@ If LOUDLY is non-nil, display some debugging information."
            (enable (nth 1 setting))
            (override (nth 3 setting))
            (language (treesit-query-language query)))
-      (when-let ((nodes (list (treesit-buffer-root-node language)))
-                 ;; Only activate if ENABLE flag is t.
-                 (activate (eq t enable)))
-        (ignore activate)
 
-        ;; If we run into problematic files, use the "fast mode" to
-        ;; try to recover.  See comment #2 above for more explanation.
-        (when treesit--font-lock-fast-mode
-          (setq nodes (treesit--children-covering-range-recurse
-                       (car nodes) start end (* 4 jit-lock-chunk-size))))
+      ;; Use deterministic way to decide whether to turn on "fast
+      ;; mode". (See bug#60691, bug#60223.)
+      (when (eq treesit--font-lock-fast-mode 'unspecified)
+        (pcase-let ((`(,max-depth ,max-width)
+                     (treesit-subtree-stat
+                      (treesit-buffer-root-node language))))
+          (if (or (> max-depth 100) (> max-width 4000))
+              (setq treesit--font-lock-fast-mode t)
+            (setq treesit--font-lock-fast-mode nil))))
+
+      (when-let* ((root (treesit-buffer-root-node language))
+                  (nodes (if (eq t treesit--font-lock-fast-mode)
+                             (treesit--children-covering-range-recurse
+                              root start end (* 4 jit-lock-chunk-size))
+                           (list (treesit-buffer-root-node language))))
+                  ;; Only activate if ENABLE flag is t.
+                  (activate (eq t enable)))
+        (ignore activate)
 
         ;; Query each node.
         (dolist (sub-node nodes)
           (let* ((delta-start (car treesit--font-lock-query-expand-range))
                  (delta-end (cdr treesit--font-lock-query-expand-range))
-                 (start-time (current-time))
                  (captures (treesit-query-capture
                             sub-node query
                             (max (- start delta-start) (point-min))
-                            (min (+ end delta-end) (point-max))))
-                 (end-time (current-time)))
-            ;; If for any query the query time is strangely long,
-            ;; switch to fast mode (see comments above).
-            (when (and (null treesit--font-lock-fast-mode)
-                       (> (time-to-seconds
-                           (time-subtract end-time start-time))
-                          0.01))
-              (if (> treesit--font-lock-fast-mode-grace-count 0)
-                  (cl-decf treesit--font-lock-fast-mode-grace-count)
-                (setq-local treesit--font-lock-fast-mode t)))
+                            (min (+ end delta-end) (point-max)))))
 
             ;; For each captured node, fontify that node.
             (with-silent-modifications
@@ -1023,12 +1022,14 @@ If LOUDLY is non-nil, display some debugging information."
                        (node (cdr capture))
                        (node-start (treesit-node-start node))
                        (node-end (treesit-node-end node)))
+
                   ;; If node is not in the region, take them out.  See
                   ;; comment #3 above for more detail.
                   (if (and (facep face)
                            (or (>= start node-end) (>= node-start end)))
                       (when (or loudly treesit--font-lock-verbose)
                         (message "Captured node %s(%s-%s) but it is outside of fontifing region" node node-start node-end))
+
                     (cond
                      ((facep face)
                       (treesit-fontify-with-override
@@ -1036,6 +1037,7 @@ If LOUDLY is non-nil, display some debugging information."
                        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.
-- 
2.39.2