From fc5455819f9cd8e67dd73b6befe4d9815af73bab Mon Sep 17 00:00:00 2001
From: Samuel Giuseppe Schlesinger <sgschlesinger@gmail.com>
Date: Sun, 1 Nov 2026 14:49:31 -0500
Subject: [PATCH] perf(polytime): compute the raw-import classification once
 per file
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Every `polytime` and `ppt` call rebuilt a table saying which imported
modules can see a tidy (`ComputableTidily…`) declaration, by rendering
every imported constant name as a string. With the library's imports
that scan costs about 0.45 s, so it dominated short searches: a trivial
`polytime` goal took about 0.5 s, now about 45 ms.

The table depends only on the import header, which is fixed while a
file is elaborated, so keep the latest one in an `IO.Ref` keyed by the
module-name list. Rule classification for the current module, local
evidence and all search tables are still computed per invocation.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015DLg8kwPs8yi7ePPdtoNvG
---
 Complexity/Tactic/PolyTime.lean | 44 ++++++++++++++++++++++++---------
 1 file changed, 33 insertions(+), 11 deletions(-)

diff --git a/Complexity/Tactic/PolyTime.lean b/Complexity/Tactic/PolyTime.lean
index 1bf9796699..5dd5692d19 100644
--- a/Complexity/Tactic/PolyTime.lean
+++ b/Complexity/Tactic/PolyTime.lean
@@ -436,6 +436,38 @@ private meta def plainBuiltinEnvironment (env : Environment) : MetaM Environment
   unless matched.size == 1 do throwError "expected one builtin local-application rule"
   return searchRuleState ext env fun state => { state with unsafeRules := index }
 
+open Lean in
+/-- For each imported module, whether its import closure declares a raw-family constant. -/
+private meta def rawImportClosure (header : EnvironmentHeader) : Std.HashMap Name Bool := Id.run do
+  -- Imports are dependency-first. A theorem in a module with no raw-family declaration in
+  -- its import closure cannot hide one behind an alias: aliases cannot refer to later modules.
+  let mut rawModules : Std.HashMap Name Bool := {}
+  for i in [:header.moduleNames.size] do
+    let moduleName := header.moduleNames[i]!
+    let data := header.moduleData[i]!
+    let ownsRaw := data.constNames.any fun name =>
+      name.toString.startsWith "Turing.MultiTapeTM.ComputableTidily" ||
+        name.toString.startsWith "Turing.MultiTapeTM.ComputesTidily"
+    let importsRaw := data.imports.any fun imported => rawModules.getD imported.module true
+    rawModules := rawModules.insert moduleName (ownsRaw || importsRaw)
+  return rawModules
+
+open Lean in
+/-- The latest import list with its `rawImportClosure`. Scanning every imported constant name
+takes about half a second, and the import header is fixed while a file is elaborated. -/
+private meta initialize rawImportClosureCache : IO.Ref (Array Name × Std.HashMap Name Bool) ←
+  IO.mkRef (#[], {})
+
+open Lean in
+/-- `rawImportClosure`, computed once per import list rather than once per invocation. -/
+private meta def cachedRawImportClosure (header : EnvironmentHeader) :
+    BaseIO (Std.HashMap Name Bool) := do
+  let (moduleNames, rawModules) ← rawImportClosureCache.get
+  if moduleNames == header.moduleNames then return rawModules
+  let rawModules := rawImportClosure header
+  rawImportClosureCache.set (header.moduleNames, rawModules)
+  return rawModules
+
 open Lean Meta Mathlib.Meta.FunProp in
 /-- Prepare only immutable global declaration classifications, never proofs or local evidence. -/
 private meta def preparePlainRuleMemo : MetaM PlainRuleCache := do
@@ -453,18 +485,8 @@ private meta def preparePlainRuleMemo : MetaM PlainRuleCache := do
     for node in rules.theorems.tries do
       for rule in node.values do names := names.insert rule.thmName
       for (_, rule) in node.pending do names := names.insert rule.thmName
-  -- Imports are dependency-first. A theorem in a module with no raw-family declaration in
-  -- its import closure cannot hide one behind an alias: aliases cannot refer to later modules.
   -- Unknown module metadata remains conservative, and current-module rules are always checked.
-  let mut rawModules : Std.HashMap Name Bool := {}
-  for i in [:env.header.moduleNames.size] do
-    let moduleName := env.header.moduleNames[i]!
-    let data := env.header.moduleData[i]!
-    let ownsRaw := data.constNames.any fun name =>
-      name.toString.startsWith "Turing.MultiTapeTM.ComputableTidily" ||
-        name.toString.startsWith "Turing.MultiTapeTM.ComputesTidily"
-    let importsRaw := data.imports.any fun imported => rawModules.getD imported.module true
-    rawModules := rawModules.insert moduleName (ownsRaw || importsRaw)
+  let rawModules ← cachedRawImportClosure env.header
   let mut cache : Std.HashMap Name Bool := {}
   for name in names do
     let mayUseRaw := match env.getModuleIdxFor? name with