From dc648c3ba0a10fd80653039398dd944ddeac34bf Mon Sep 17 00:00:00 2001
From: Samuel Giuseppe Schlesinger <sgschlesinger@gmail.com>
Date: Sun, 1 Nov 2026 16:40:56 -0500
Subject: [PATCH] refactor(polytime): keep the raw-import table in an
 environment extension

#2191 cached the table of imported modules whose import closure declares a
raw-family constant in a process-global `IO.Ref` keyed by the import list.
Store it instead as the state of a persistent environment extension whose
`addImportedFn` creates a `Thunk` over the import header. The scan runs once
per file on first use, never in files that do not call the tactics, and
proofs elaborated in parallel share one computation instead of each
repeating it.

Replace the two `Name.toString` prefix tests with one structural predicate,
`isRawFamilyName`, used by both the import scan and `rawTransitionPremise`.
It classifies every imported constant of the library identically and makes
the one-time scan about 40% cheaper.

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

diff --git a/Complexity/Tactic/PolyTime.lean b/Complexity/Tactic/PolyTime.lean
index 5dd5692d19..b3d535b8d8 100644
--- a/Complexity/Tactic/PolyTime.lean
+++ b/Complexity/Tactic/PolyTime.lean
@@ -277,16 +277,22 @@ private meta def polytimeKnownFacts : TacticM Unit := withMainContext do
     else (``ComputableInPolyTime, ``ComputableInPolyTime.tidy)
   restateKnownFacts source conversion
 
+open Lean in
+/-- Whether a constant name is in the raw family: its component after `Turing.MultiTapeTM` starts
+with `ComputableTidily` or `ComputesTidily`. -/
+private meta def isRawFamilyName : Name → Bool
+  | .str (.str (.str .anonymous "Turing") "MultiTapeTM") head =>
+    head.startsWith "ComputableTidily" || head.startsWith "ComputesTidily"
+  | .str parent _ | .num parent _ => isRawFamilyName parent
+  | .anonymous => false
+
 open Lean Meta in
 /-- Inspect proposition aliases with the same reducible transparency used by `fun_prop`.
 `Meta.transform` maintains binder context, including dependent and quantified premises; data and
 proof implementations are not unfolded merely to classify a transition's proposition. -/
 private meta def rawTransitionPremise (type : Expr) : MetaM Bool := do
-  let rawName := fun n : Name =>
-    n.toString.startsWith "Turing.MultiTapeTM.ComputableTidily" ||
-      n.toString.startsWith "Turing.MultiTapeTM.ComputesTidily"
   if (type.find? fun e => match e with
-      | .const n _ => rawName n
+      | .const n _ => isRawFamilyName n
       | _ => false).isSome then
     return true
   let inspect : StateRefT Bool MetaM Unit := do
@@ -294,7 +300,7 @@ private meta def rawTransitionPremise (type : Expr) : MetaM Bool := do
       if ← get then return .done e
       let reduced ← if ← isProp e then whnfR e else pure e
       if let .const name _ := reduced.getAppFn then
-        if rawName name then
+        if isRawFamilyName name then
           set true
           return .done reduced
       return .continue (some reduced))
@@ -445,28 +451,24 @@ private meta def rawImportClosure (header : EnvironmentHeader) : Std.HashMap Nam
   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 ownsRaw := data.constNames.any isRawFamilyName
     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
+/-- `rawImportClosure` of the imports, which are fixed for the whole file. Importing only creates
+the thunk: the first tactic invocation scans the imported constants, and every later one,
+including proofs elaborated in parallel, reuses the result. The extension stores no entries. -/
+private meta initialize rawImportClosureExt :
+    PersistentEnvExtension Empty Empty (Thunk (Std.HashMap Name Bool)) ←
+  registerPersistentEnvExtension {
+    mkInitial := pure (.pure {})
+    addImportedFn := fun _ => do
+      let header := (← read).env.header
+      return .mk fun _ => rawImportClosure header
+    addEntryFn := fun rawModules _ => rawModules
+    exportEntriesFn := fun _ => #[] }
 
 open Lean Meta Mathlib.Meta.FunProp in
 /-- Prepare only immutable global declaration classifications, never proofs or local evidence. -/
@@ -486,7 +488,7 @@ private meta def preparePlainRuleMemo : MetaM PlainRuleCache := do
       for rule in node.values do names := names.insert rule.thmName
       for (_, rule) in node.pending do names := names.insert rule.thmName
   -- Unknown module metadata remains conservative, and current-module rules are always checked.
-  let rawModules ← cachedRawImportClosure env.header
+  let rawModules := (rawImportClosureExt.getState env).get
   let mut cache : Std.HashMap Name Bool := {}
   for name in names do
     let mayUseRaw := match env.getModuleIdxFor? name with