From 4f4f6c4c6d035da35de570546a7785bd389906c8 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Fri, 17 Jan 2025 16:14:33 +0100
Subject: [PATCH] perf: optimize `setImportedEntries`

---
 src/Lean/Environment.lean | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/Lean/Environment.lean b/src/Lean/Environment.lean
index 039973a374c6..f2fa65079adc 100644
--- a/src/Lean/Environment.lean
+++ b/src/Lean/Environment.lean
@@ -1208,19 +1208,23 @@ def mkExtNameMap (startingAt : Nat) : IO (Std.HashMap Name Nat) := do
   return result
 
 private def setImportedEntries (env : Environment) (mods : Array ModuleData) (startingAt : Nat := 0) : IO Environment := do
-  let mut env := env
+  -- We work directly on the states array instead of `env` as `Environment.modifyState` introduces
+  -- significant overhead on such frequent calls
+  let mut states := env.checkedWithoutAsync.extensions
   let extDescrs ← persistentEnvExtensionsRef.get
   /- For extensions starting at `startingAt`, ensure their `importedEntries` array have size `mods.size`. -/
   for extDescr in extDescrs[startingAt:] do
-    env := extDescr.toEnvExtension.modifyState env fun s => { s with importedEntries := mkArray mods.size #[] }
+    states := EnvExtensionInterfaceImp.modifyState extDescr.toEnvExtension states fun s =>
+      { s with importedEntries := mkArray mods.size #[] }
   /- For each module `mod`, and `mod.entries`, if the extension name is one of the extensions after `startingAt`, set `entries` -/
   let extNameIdx ← mkExtNameMap startingAt
   for h : modIdx in [:mods.size] do
     let mod := mods[modIdx]
     for (extName, entries) in mod.entries do
       if let some entryIdx := extNameIdx[extName]? then
-        env := extDescrs[entryIdx]!.toEnvExtension.modifyState env fun s => { s with importedEntries := s.importedEntries.set! modIdx entries }
-  return env
+        states := EnvExtensionInterfaceImp.modifyState extDescrs[entryIdx]!.toEnvExtension states fun s =>
+          { s with importedEntries := s.importedEntries.set! modIdx entries }
+  return env.setCheckedSync { env.checkedWithoutAsync with extensions := states }
 
 /--
   "Forward declaration" needed for updating the attribute table with user-defined attributes.