Skip to content

Commit

Permalink
fix: Thread unique name generator through library_search init
Browse files Browse the repository at this point in the history
  • Loading branch information
joehendrix authored and kim-em committed Mar 6, 2024
1 parent 50eca90 commit 6fce8f7
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 10 deletions.
20 changes: 11 additions & 9 deletions src/Lean/Meta/LazyDiscrTree.lean
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
5 changes: 4 additions & 1 deletion src/Lean/Meta/Tactic/LibrarySearch.lean
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 6fce8f7

Please sign in to comment.