Skip to content

Commit 21312f9

Browse files
KhaKha
committed
fix: escalate .mainEnv to .asyncEnv on realized consts
1 parent 6cca822 commit 21312f9

1 file changed

Lines changed: 23 additions & 17 deletions

File tree

‎src/Lean/Environment.lean‎

Lines changed: 23 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -434,17 +434,18 @@ private def AsyncContext.mayContain (ctx : AsyncContext) (n : Name) : Bool :=
434434
Constant info and environment extension states eventually resulting from async elaboration.
435435
-/
436436
private structure AsyncConst where
437-
constInfo : AsyncConstantInfo
437+
constInfo : AsyncConstantInfo
438438
/--
439439
Reported extension state eventually fulfilled by promise; may be missing for tasks (e.g. kernel
440440
checking) that can eagerly guarantee they will not report any state.
441441
-/
442-
exts? : Option (Task (Array EnvExtensionState))
442+
exts? : Option (Task (Array EnvExtensionState))
443443
/--
444444
`Task AsyncConsts` except for problematic recursion. The set of nested constants created while
445445
elaborating this constant.
446446
-/
447-
consts : Task Dynamic
447+
aconstsImpl : Task Dynamic
448+
isRealized : Bool := false
448449

449450
/-- Data structure holding a sequence of `AsyncConst`s optimized for efficient access. -/
450451
private structure AsyncConsts where
@@ -456,6 +457,10 @@ private structure AsyncConsts where
456457
normalizedTrie : NameTrie AsyncConst
457458
deriving Inhabited, TypeName
458459

460+
private def AsyncConst.aconsts (c : AsyncConst) : Task AsyncConsts :=
461+
c.aconstsImpl.map (sync := true) fun dyn =>
462+
dyn.get? AsyncConsts |>.getD default
463+
459464
private def AsyncConsts.add (aconsts : AsyncConsts) (aconst : AsyncConst) : AsyncConsts :=
460465
let normalizedName := privateToUserName aconst.constInfo.name
461466
if let some aconst' := aconsts.normalizedTrie.find? normalizedName then
@@ -488,16 +493,15 @@ private partial def AsyncConsts.findRec? (aconsts : AsyncConsts) (declName : Nam
488493
-- If privacy is the only difference between `declName` and `findPrefix?` result, we can assume
489494
-- `declName` does not exist according to the `add` invariant
490495
guard <| privateToUserName c.constInfo.name != privateToUserName declName
491-
let aconsts ← c.consts.get.get? AsyncConsts
496+
let aconsts ← c.aconsts.get
492497
AsyncConsts.findRec? aconsts declName
493498

494499
/-- Like `findRec?`; allocating tasks is (currently?) too costly to do always. -/
495500
private partial def AsyncConsts.findRecTask (aconsts : AsyncConsts) (declName : Name) : Task (Option AsyncConst) := Id.run do
496501
let some c := aconsts.findPrefix? declName | .pure none
497502
if c.constInfo.name == declName then
498503
return .pure c
499-
c.consts.bind (sync := true) fun aconsts => Id.run do
500-
let some aconsts := aconsts.get? AsyncConsts | .pure none
504+
c.aconsts.bind (sync := true) fun aconsts => Id.run do
501505
AsyncConsts.findRecTask aconsts declName
502506

503507
/-- Like `findRec?` but returns the constant that has `declName` in its `consts`, if any. -/
@@ -510,8 +514,7 @@ where go parent? aconsts := do
510514
-- If privacy is the only difference between `declName` and `findPrefix?` result, we can assume
511515
-- `declName` does not exist according to the `add` invariant
512516
guard <| privateToUserName c.constInfo.name != privateToUserName declName
513-
let aconsts ← c.consts.get.get? AsyncConsts
514-
go (some c) aconsts
517+
go (some c) c.aconsts.get
515518

516519
/-- Accessibility levels of declarations in `Lean.Environment`. -/
517520
private inductive Visibility where
@@ -733,14 +736,14 @@ def addDeclCore (env : Environment) (maxHeartbeats : USize) (decl : @& Declarati
733736
env := { env with asyncConstsMap.private := env.asyncConstsMap.private.add {
734737
constInfo := .ofConstantInfo info
735738
exts? := none
736-
consts := .pure <| .mk (α := AsyncConsts) default
739+
aconstsImpl := .pure <| .mk (α := AsyncConsts) default
737740
} }
738741
-- TODO
739742
if true /- !isPrivateName n-/ then
740743
env := { env with asyncConstsMap.public := env.asyncConstsMap.public.add {
741744
constInfo := .ofConstantInfo info
742745
exts? := none
743-
consts := .pure <| .mk (α := AsyncConsts) default
746+
aconstsImpl := .pure <| .mk (α := AsyncConsts) default
744747
} }
745748

746749
return env
@@ -762,7 +765,7 @@ private def lakeAdd (env : Environment) (cinfo : ConstantInfo) : Environment :=
762765
asyncConstsMap := env.asyncConstsMap.map (·.add {
763766
constInfo := .ofConstantInfo cinfo
764767
exts? := none
765-
consts := .pure <| .mk (α := AsyncConsts) default
768+
aconstsImpl := .pure <| .mk (α := AsyncConsts) default
766769
})
767770
}
768771

@@ -1045,7 +1048,7 @@ def addConstAsync (env : Environment) (constName : Name) (kind : ConstantKind)
10451048
| some v => v.exts
10461049
-- any value should work here, `base` does not block
10471050
| none => env.base.private.extensions)
1048-
consts := constPromise.result?.map (sync := true) fun
1051+
aconstsImpl := constPromise.result?.map (sync := true) fun
10491052
| some v => .mk v.nestedConsts.private
10501053
| none => .mk (α := AsyncConsts) default
10511054
}
@@ -1056,7 +1059,7 @@ def addConstAsync (env : Environment) (constName : Name) (kind : ConstantKind)
10561059
| some c => c.exportedConstInfo
10571060
| none => mkFallbackConstInfo constName exportedKind
10581061
}
1059-
consts := constPromise.result?.map (sync := true) fun
1062+
aconstsImpl := constPromise.result?.map (sync := true) fun
10601063
| some v => .mk v.nestedConsts.public
10611064
| none => .mk (α := AsyncConsts) default
10621065
}
@@ -1393,12 +1396,13 @@ private unsafe def getStateUnsafe {σ : Type} [Inhabited σ] (ext : EnvExtension
13931396
-- safety: `ext`'s constructor is private, so we can assume the entry at `ext.idx` is of type `σ`
13941397
if env.base.get env |>.constants.contains asyncDecl then
13951398
return ext.getStateImpl env.base.private.extensions
1396-
if let some exts? := match branch with
1399+
if let some exts? := (match branch with
13971400
| .asyncEnv => env.asyncConsts.findRec? asyncDecl |>.map AsyncConst.exts?
13981401
| .mainEnv =>
13991402
match env.asyncConsts.find? asyncDecl with
1400-
| some _ => some <| some <| .pure env.base.private.extensions
1401-
| _ => env.asyncConsts.findRecParent? asyncDecl |>.map AsyncConst.exts? then
1403+
| some c => if c.isRealized then c.exts? else some <| some <| .pure env.base.private.extensions
1404+
| _ => env.asyncConsts.findRecParent? asyncDecl |>.bind fun c =>
1405+
if c.isRealized then c.aconsts.get.find? asyncDecl |>.map AsyncConst.exts? else c.exts?) then
14021406
if let some exts := exts? then
14031407
return ext.getStateImpl exts.get
14041408
-- NOTE: if `exts?` is `none`, we should *not* try the following, more expensive branches that
@@ -2244,7 +2248,7 @@ private def updateBaseAfterKernelAdd (env : Environment) (kenv : Kernel.Environm
22442248
asyncConsts.add {
22452249
constInfo := .ofConstantInfo (kenv.find? n |>.get!)
22462250
exts? := none
2247-
consts := .pure <| .mk (α := AsyncConsts) default
2251+
aconstsImpl := .pure <| .mk (α := AsyncConsts) default
22482252
}
22492253
else asyncConsts
22502254
}
@@ -2431,12 +2435,14 @@ def realizeConst (env : Environment) (forConst : Name) (constName : Name)
24312435
let numNewPrivateConsts := realizeEnv'.asyncConstsMap.private.size - realizeEnv.asyncConstsMap.private.size
24322436
let newPrivateConsts := realizeEnv'.asyncConstsMap.private.revList.take numNewPrivateConsts |>.reverse
24332437
let newPrivateConsts := newPrivateConsts.map fun c =>
2438+
let c := { c with isRealized := true }
24342439
if c.exts?.isNone then
24352440
{ c with exts? := some <| .pure realizeEnv'.base.private.extensions }
24362441
else c
24372442
let numNewPublicConsts := realizeEnv'.asyncConstsMap.public.size - realizeEnv.asyncConstsMap.public.size
24382443
let newPublicConsts := realizeEnv'.asyncConstsMap.public.revList.take numNewPublicConsts |>.reverse
24392444
let newPublicConsts := newPublicConsts.map fun c =>
2445+
let c := { c with isRealized := true }
24402446
if c.exts?.isNone then
24412447
{ c with exts? := some <| .pure realizeEnv'.base.private.extensions }
24422448
else c

0 commit comments

Comments
 (0)