Lean4
/-- Create a `RefinedDiscrTree` consisting of all entries generated by `act`
from imported constants. (it gets called by `addConstToPreDiscrTree`).
This uses parallel computation.
-/
def createImportedDiscrTree (ngen : NameGenerator) (env : Environment)
(act : Name → ConstantInfo → MetaM (List (α × List (Key × LazyEntry)))) (constantsPerTask capacityPerTask : Nat) :
CoreM (RefinedDiscrTree α) := do
let numModules := env.header.moduleData.size
let cctx ← read
let rec /-- Allocate constants to tasks according to `constantsPerTask`. -/
go (ngen : NameGenerator) (tasks : Array (Task (InitResults α))) (start cnt idx : Nat) := do
if h : idx < numModules then
let mdata := env.header.moduleData[idx]
let cnt := cnt + mdata.constants.size
if cnt > constantsPerTask then
let (childNGen, ngen) := ngen.mkChild
let t ← (createImportInitResults cctx childNGen env act capacityPerTask start (idx + 1)).asTask
go ngen (tasks.push t) (idx + 1) 0 (idx + 1)
else
go ngen tasks start cnt (idx + 1)
else
if start < numModules then
let (childNGen, _) := ngen.mkChild
let t ← (createImportInitResults cctx childNGen env act capacityPerTask start numModules).asTask
pure (tasks.push t)
else
pure tasks
termination_by env.header.moduleData.size - idx
let tasks ← go ngen #[] 0 0 0
let r : InitResults α := tasks.foldl (init := { }) (· ++ ·.get)
r.errors.forM logImportFailure
return r.tree.toRefinedDiscrTree
