From 6fce8f7d5cd18a4419bca7fd51780c71c9b1cc5a Mon Sep 17 00:00:00 2001
From: Joe Hendrix <joe@lean-fro.org>
Date: Tue, 5 Mar 2024 16:33:14 -0800
Subject: [PATCH] fix: Thread unique name generator through library_search init

---
 src/Lean/Meta/LazyDiscrTree.lean        | 20 +++++++++++---------
 src/Lean/Meta/Tactic/LibrarySearch.lean |  5 ++++-
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/src/Lean/Meta/LazyDiscrTree.lean b/src/Lean/Meta/LazyDiscrTree.lean
index f314cc06a401..b472f0b82c88 100644
--- a/src/Lean/Meta/LazyDiscrTree.lean
+++ b/src/Lean/Meta/LazyDiscrTree.lean
@@ -791,10 +791,10 @@ private partial def loadImportedModule (env : Environment)
   else
     pure tree
 
-private def createImportedEnvironmentSeq (env : Environment)
+private def createImportedEnvironmentSeq (ngen : NameGenerator) (env : Environment)
     (act : Name → ConstantInfo → MetaM (Array (InitEntry α)))
     (start stop : Nat) : BaseIO (InitResults α) := do
-      let cacheRef ← IO.mkRef (Cache.empty {})
+      let cacheRef ← IO.mkRef (Cache.empty ngen)
       go (← ImportData.new) cacheRef {} start stop
     where go d cacheRef (tree : PreDiscrTree α) (start stop : Nat) : BaseIO _ := do
             if start < stop then
@@ -811,29 +811,31 @@ private def combineGet [Append α] (z : α) (tasks : Array (Task α)) : α :=
   tasks.foldl (fun x t => x ++ t.get) (init := z)
 
 /-- Create an imported environment for tree. -/
-def createImportedEnvironment (env : Environment)
+def createImportedEnvironment (ngen : NameGenerator) (env : Environment)
     (act : Name → ConstantInfo → MetaM (Array (InitEntry α)))
     (constantsPerTask : Nat := 1000) :
     EIO Exception (LazyDiscrTree α) := do
   let n := env.header.moduleData.size
   let rec
     /-- Allocate constants to tasks according to `constantsPerTask`. -/
-    go tasks start cnt idx := do
+    go ngen tasks start cnt idx := do
       if h : idx < env.header.moduleData.size then
         let mdata := env.header.moduleData[idx]
         let cnt := cnt + mdata.constants.size
         if cnt > constantsPerTask then
-          let t ← createImportedEnvironmentSeq env act start (idx+1) |>.asTask
-          go (tasks.push t) (idx+1) 0 (idx+1)
+          let (childNGen, ngen) := ngen.mkChild
+          let t ← createImportedEnvironmentSeq childNGen env act start (idx+1) |>.asTask
+          go ngen (tasks.push t) (idx+1) 0 (idx+1)
         else
-          go tasks start cnt (idx+1)
+          go ngen tasks start cnt (idx+1)
       else
         if start < n then
-          tasks.push <$> (createImportedEnvironmentSeq env act start n).asTask
+          let (childNGen, _) := ngen.mkChild
+          tasks.push <$> (createImportedEnvironmentSeq childNGen env act start n).asTask
         else
           pure tasks
     termination_by env.header.moduleData.size - idx
-  let tasks ← go #[] 0 0 0
+  let tasks ← go ngen #[] 0 0 0
   let r := combineGet default tasks
   if p : r.errors.size > 0 then
     throw r.errors[0].exception
diff --git a/src/Lean/Meta/Tactic/LibrarySearch.lean b/src/Lean/Meta/Tactic/LibrarySearch.lean
index a21eb50a0d1a..8b0822f35ff3 100644
--- a/src/Lean/Meta/Tactic/LibrarySearch.lean
+++ b/src/Lean/Meta/Tactic/LibrarySearch.lean
@@ -150,9 +150,12 @@ Candidate-finding function that uses a strict discrimination tree for resolution
 def mkImportFinder : IO CandidateFinder := do
   let ref ← IO.mkRef none
   pure fun ty => do
+    let ngen ← getNGen
+    let (childNGen, ngen) := ngen.mkChild
+    setNGen ngen
     let importTree ← (←ref.get).getDM $ do
       profileitM Exception  "librarySearch launch" (←getOptions) $
-        createImportedEnvironment (←getEnv) (constantsPerTask := constantsPerTask) addImport
+        createImportedEnvironment childNGen (←getEnv) (constantsPerTask := constantsPerTask) addImport
     let (imports, importTree) ← importTree.getMatch ty
     ref.set importTree
     pure imports