PY-91640 code vision: bound widespread-name usage searches per pass

Counting usages of a name too widespread for a cheap search needs a
project-wide ReferencesSearch. The 100ms budget PY-82336 gave it never held:
the deadline was checked from the result processor, which is not called while
the search walks files with no match, and the budget was per declaration, so it
multiplied by their number - test_driver.py has 1351, enough to turn the first
code analysis of it into a 5-minute timeout.

The deadline is now enforced by cancelling the search, and one budget is shared
by the whole pass. The hint stays: the counters live outside the timed block, so
a cancelled search still reports its partial count as "N+".


(cherry picked from commit 8e5ec90c1b9c47a0458a4cae08da5b2125d38bfa)

IJ-MR-218867

GitOrigin-RevId: a666c00b36b51bca4e52d623b302c1b9ff405f9d

python/src/com/jetbrains/python/codeInsight/codeVision/PyReferencesCodeVisionProvider.kt

@@ -1,9 +1,13 @@
 // Copyright 2000-2026 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.jetbrains.python.codeInsight.codeVision
 
+import com.intellij.codeInsight.codeVision.CodeVisionEntry
 import com.intellij.codeInsight.codeVision.CodeVisionRelativeOrdering
 import com.intellij.codeInsight.hints.codeVision.ReferencesCodeVisionProvider
+import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.options.advanced.AdvancedSettings.Companion.getInt
+import com.intellij.openapi.progress.util.ProgressIndicatorUtils
+import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiNamedElement
@@ -29,7 +33,7 @@ import kotlin.time.Duration
 import kotlin.time.Duration.Companion.milliseconds
 import kotlin.time.TimeSource
 
-/** Test-only override for the widespread-search budget; the timing is otherwise machine-dependent. */
+/** Test-only override for both budgets below, so a test expecting an exact count can let the search finish. */
 @TestOnly
 var pyUsagesSearchBudgetTestOverride: Duration? = null
 
@@ -37,11 +41,33 @@ class PyReferencesCodeVisionProvider : ReferencesCodeVisionProvider() {
   companion object {
     const val ID = "python.references"
 
-    // When a symbol's name is too widespread for a cheap exact count, ReferencesSearch could end up
-    // scanning the whole project (worst case: a common-word name with few real references). For such
-    // names we cap the search by this wall-clock budget and render the partial count as "N+",
-    // instead of suppressing the hint entirely as before (PY-82336).
-    private val WIDESPREAD_SEARCH_BUDGET = 100.milliseconds
+    // One widespread-name search; its partial count is rendered as "N+"
+    private val SINGLE_SEARCH_BUDGET = 100.milliseconds
+
+    /** Shared by all widespread-name searches of one pass, so the per-search budget
+     * is not multiplied by the number of declarations */
+    private val PASS_BUDGET = 300.milliseconds
+  }
+
+  /** `computeForEditor` walks the file on one thread, so no synchronization is needed. */
+  private val passBudget = ThreadLocal<PassBudget>()
+
+  private class PassBudget(private var remaining: Duration) {
+    fun nextSlice(singleSearchBudget: Duration): Duration = minOf(remaining, singleSearchBudget)
+
+    fun charge(spent: Duration) {
+      remaining -= spent
+    }
+  }
+
+  override fun computeForEditor(editor: Editor, file: PsiFile): List<Pair<TextRange, CodeVisionEntry>> {
+    passBudget.set(PassBudget(pyUsagesSearchBudgetTestOverride ?: PASS_BUDGET))
+    try {
+      return super.computeForEditor(editor, file)
+    }
+    finally {
+      passBudget.remove()
+    }
   }
 
   override fun acceptsFile(file: PsiFile): Boolean = file is PyFile
@@ -76,33 +102,33 @@ class PyReferencesCodeVisionProvider : ReferencesCodeVisionProvider() {
     val configuredLimit = getInt("python.code.vision.usages.limit")
     val maxUsagesToCount = if (configuredLimit > 0) configuredLimit else Int.MAX_VALUE
 
-    // If the name is widespread, the search could be expensive, so bound it by time and render "N+".
-    // For cheap names we keep an exact, deterministic count (no time budget, so results don't flicker).
-    val isWidespreadName = PsiSearchHelper.getInstance(project).isCheapEnoughToSearch(elementName, scope, file) ==
-                     SearchCostResult.TOO_MANY_OCCURRENCES
-    val budget = pyUsagesSearchBudgetTestOverride ?: WIDESPREAD_SEARCH_BUDGET
-    val deadline = if (isWidespreadName) TimeSource.Monotonic.markNow() + budget else null
+    val isWidespreadName = PsiSearchHelper.getInstance(project)
+      .isCheapEnoughToSearch(elementName, scope, file) == SearchCostResult.TOO_MANY_OCCURRENCES
 
     val regularUsages = AtomicInteger()
     val dynamicUsages = AtomicInteger()
     val isTruncated = AtomicBoolean()
+    val search = { countUsages(element, scope, maxUsagesToCount, regularUsages, dynamicUsages, isTruncated) }
 
-    ReferencesSearch.search(ReferencesSearch.SearchParameters(element, scope, false))
-      .allowParallelProcessing()
-      .forEach(Processor {
-        if (it == null) true
-        else if (element.reference == it) true
-        else if (PsiTreeUtil.getParentOfType(it.element, PyImportStatementBase::class.java) != null) true // imports are not usages
-        else {
-          if (UsageInfo(it).isDynamicUsage) dynamicUsages.incrementAndGet()
-          val count = regularUsages.incrementAndGet()
-          if (count > maxUsagesToCount || deadline?.hasPassedNow() == true) {
-            isTruncated.set(true)
-            false
-          }
-          else true
-        }
-      })
+    if (!isWidespreadName) {
+      search()
+    }
+    else {
+      val singleSearchBudget = pyUsagesSearchBudgetTestOverride ?: SINGLE_SEARCH_BUDGET
+      val budget = passBudget.get()
+      val slice = budget?.nextSlice(singleSearchBudget) ?: singleSearchBudget
+      if (slice <= Duration.ZERO) return null
+
+      val startedAt = TimeSource.Monotonic.markNow()
+      // The deadline has to be enforced by cancellation rather than checked from the result processor: the
+      // processor is not called while the search walks files that contain no match, which is exactly the
+      // case this budget exists for. `ReferencesSearch` calls `checkCanceled` often enough for the
+      // indicator `withTimeout` cancels to stop it. The counters are updated in place, so a timed-out
+      // search still contributes what it managed to find.
+      val completed = ProgressIndicatorUtils.withTimeout(slice.inWholeMilliseconds) { search(); true } != null
+      budget?.charge(startedAt.elapsedNow())
+      if (!completed) isTruncated.set(true)
+    }
 
     val result = regularUsages.get()
     val dynamicResult = dynamicUsages.get()
@@ -115,6 +141,31 @@ class PyReferencesCodeVisionProvider : ReferencesCodeVisionProvider() {
     return CodeVisionInfo(PyBundle.message("inlay.hints.usages.with.dynamic.text", result, dynamicResult), result)
   }
 
+  private fun countUsages(
+    element: PsiNamedElement,
+    scope: GlobalSearchScope,
+    maxUsagesToCount: Int,
+    regularUsages: AtomicInteger,
+    dynamicUsages: AtomicInteger,
+    isTruncated: AtomicBoolean,
+  ) {
+    ReferencesSearch.search(ReferencesSearch.SearchParameters(element, scope, false))
+      .allowParallelProcessing()
+      .forEach(Processor {
+        if (it == null) true
+        else if (element.reference == it) true
+        else if (PsiTreeUtil.getParentOfType(it.element, PyImportStatementBase::class.java) != null) true // imports are not usages
+        else {
+          if (UsageInfo(it).isDynamicUsage) dynamicUsages.incrementAndGet()
+          if (regularUsages.incrementAndGet() > maxUsagesToCount) {
+            isTruncated.set(true)
+            false
+          }
+          else true
+        }
+      })
+  }
+
   override fun getHint(element: PsiElement, file: PsiFile): String? {
     return getVisionInfo(element, file)?.text
   }