From 141e519009fc74d444cb8a32b55daf3c2d192bbd Mon Sep 17 00:00:00 2001 From: Sebastian Ullrich Date: Thu, 6 Mar 2025 15:27:45 +0100 Subject: [PATCH] feat: add async support to more extensions and constructions (#7363) --- src/Lean/Environment.lean | 18 ++++++++++++++++-- src/Lean/Meta/Match/Match.lean | 3 ++- src/Lean/Meta/Match/MatcherInfo.lean | 5 ++++- src/Lean/Meta/Tactic/AuxLemma.lean | 9 ++++++--- src/Lean/MonadEnv.lean | 5 +++-- src/Lean/ReducibilityAttrs.lean | 6 ++++-- 6 files changed, 35 insertions(+), 11 deletions(-) diff --git a/src/Lean/Environment.lean b/src/Lean/Environment.lean index 430a348dd3..eec87d637f 100644 --- a/src/Lean/Environment.lean +++ b/src/Lean/Environment.lean @@ -556,10 +556,18 @@ private def modifyCheckedAsync (env : Environment) (f : Kernel.Environment → K private def setCheckedSync (env : Environment) (newChecked : Kernel.Environment) : Environment := { env with checked := .pure newChecked, base := newChecked } +/-- The declaration prefix to which the environment is restricted to, if any. -/ +def asyncPrefix? (env : Environment) : Option Name := + env.asyncCtx?.map (·.declPrefix) + /-- True while inside `realizeConst`'s `realize`. -/ def isRealizing (env : Environment) : Bool := env.asyncCtx?.any (·.realizing) +/-- Forgets about the asynchronous context restrictions. Used only for `withoutModifyingEnv`. -/ +def unlockAsync (env : Environment) : Environment := + { env with asyncCtx? := none } + /-- Checks whether the given declaration name may potentially added, or have been added, to the current environment branch, which is the case either if this is the main branch or if the declaration name @@ -1453,6 +1461,7 @@ def mkTagDeclarationExtension (name : Name := by exact decl_name%) : IO TagDecla addImportedFn := fun _ => {}, addEntryFn := fun s n => s.insert n, toArrayFn := fun es => es.toArray.qsort Name.quickLt + asyncMode := .async } namespace TagDeclarationExtension @@ -1463,12 +1472,13 @@ instance : Inhabited TagDeclarationExtension := def tag (ext : TagDeclarationExtension) (env : Environment) (declName : Name) : Environment := have : Inhabited Environment := ⟨env⟩ assert! env.getModuleIdxFor? declName |>.isNone -- See comment at `TagDeclarationExtension` + assert! env.asyncMayContain declName ext.addEntry env declName def isTagged (ext : TagDeclarationExtension) (env : Environment) (declName : Name) : Bool := match env.getModuleIdxFor? declName with | some modIdx => (ext.getModuleEntries env modIdx).binSearchContains declName Name.quickLt - | none => (ext.getState env).contains declName + | none => (ext.findStateAsync env declName).contains declName end TagDeclarationExtension @@ -1861,7 +1871,11 @@ def displayStats (env : Environment) : IO Unit := do IO.println ("number of extensions: " ++ toString env.base.extensions.size); pExtDescrs.forM fun extDescr => do IO.println ("extension '" ++ toString extDescr.name ++ "'") - let s := extDescr.toEnvExtension.getState env + -- get state from `checked` at the end if `async`; it would otherwise panic + let mut asyncMode := extDescr.toEnvExtension.asyncMode + if asyncMode matches .async then + asyncMode := .sync + let s := extDescr.toEnvExtension.getState (asyncMode := asyncMode) env let fmt := extDescr.statsFn s.state unless fmt.isNil do IO.println (" " ++ toString (Format.nest 2 (extDescr.statsFn s.state))) IO.println (" number of imported entries: " ++ toString (s.importedEntries.foldl (fun sum es => sum + es.size) 0)) diff --git a/src/Lean/Meta/Match/Match.lean b/src/Lean/Meta/Match/Match.lean index 4892e14b56..244ce8be75 100644 --- a/src/Lean/Meta/Match/Match.lean +++ b/src/Lean/Meta/Match/Match.lean @@ -762,7 +762,8 @@ register_builtin_option bootstrap.genMatcherCode : Bool := { descr := "disable code generation for auxiliary matcher function" } -builtin_initialize matcherExt : EnvExtension (PHashMap (Expr × Bool) Name) ← registerEnvExtension (pure {}) +builtin_initialize matcherExt : EnvExtension (PHashMap (Expr × Bool) Name) ← + registerEnvExtension (pure {}) (asyncMode := .local) -- mere cache, keep it local /-- Similar to `mkAuxDefinition`, but uses the cache `matcherExt`. It also returns an Boolean that indicates whether a new matcher function was added to the environment or not. -/ diff --git a/src/Lean/Meta/Match/MatcherInfo.lean b/src/Lean/Meta/Match/MatcherInfo.lean index 45ce62ee1d..a59b21553f 100644 --- a/src/Lean/Meta/Match/MatcherInfo.lean +++ b/src/Lean/Meta/Match/MatcherInfo.lean @@ -82,13 +82,16 @@ builtin_initialize extension : SimplePersistentEnvExtension Entry State ← registerSimplePersistentEnvExtension { addEntryFn := State.addEntry addImportedFn := fun es => (mkStateFromImportedEntries State.addEntry {} es).switch + asyncMode := .async } def addMatcherInfo (env : Environment) (matcherName : Name) (info : MatcherInfo) : Environment := + let _ : Inhabited Environment := ⟨env⟩ + assert! env.asyncMayContain matcherName extension.addEntry env { name := matcherName, info := info } def getMatcherInfo? (env : Environment) (declName : Name) : Option MatcherInfo := - (extension.getState env).map.find? declName + (extension.findStateAsync env declName).map.find? declName end Extension diff --git a/src/Lean/Meta/Tactic/AuxLemma.lean b/src/Lean/Meta/Tactic/AuxLemma.lean index 744ff2b907..6cd7bc74b6 100644 --- a/src/Lean/Meta/Tactic/AuxLemma.lean +++ b/src/Lean/Meta/Tactic/AuxLemma.lean @@ -14,13 +14,16 @@ structure AuxLemmas where lemmas : PHashMap Expr (Name × List Name) := {} deriving Inhabited -builtin_initialize auxLemmasExt : EnvExtension AuxLemmas ← registerEnvExtension (pure {}) +builtin_initialize auxLemmasExt : EnvExtension AuxLemmas ← + registerEnvExtension (pure {}) (asyncMode := .local) -- a mere cache, keep local /-- Helper method for creating auxiliary lemmas in the environment. It uses a cache that maps `type` to declaration name. The cache is not stored in `.olean` files. - It is useful to make sure the same auxiliary lemma is not created over and over again in the same file. + It is useful to make sure the same auxiliary lemma is not created over and over again in the same + environment branch. For expensive auxiliary lemmas that should be deduplicated even across + different environment branches, consider using `realizeConst` instead. This method is useful for tactics (e.g., `simp`) that may perform preprocessing steps to lemmas provided by users. For example, `simp` preprocessor may convert a lemma into multiple ones. @@ -29,7 +32,7 @@ def mkAuxLemma (levelParams : List Name) (type : Expr) (value : Expr) : MetaM Na let env ← getEnv let s := auxLemmasExt.getState env let mkNewAuxLemma := do - let auxName := Name.mkNum (env.mainModule ++ `_auxLemma) s.idx + let auxName := Name.mkNum (env.asyncPrefix?.getD env.mainModule ++ `_auxLemma) s.idx addDecl <| Declaration.thmDecl { name := auxName levelParams, type, value diff --git a/src/Lean/MonadEnv.lean b/src/Lean/MonadEnv.lean index da32b0b9ef..1f4072644d 100644 --- a/src/Lean/MonadEnv.lean +++ b/src/Lean/MonadEnv.lean @@ -39,8 +39,9 @@ def isRec [Monad m] [MonadEnv m] (declName : Name) : m Bool := return isRecCore (← getEnv) declName @[inline] def withoutModifyingEnv [Monad m] [MonadEnv m] [MonadFinally m] {α : Type} (x : m α) : m α := do - let env ← getEnv - try x finally setEnv env + -- Allow `x` to define new declarations even outside the asynchronous prefix (if any) as all + -- results will be discarded anway. + withEnv (← getEnv).unlockAsync x /-- Similar to `withoutModifyingEnv`, but also returns the updated environment -/ @[inline] def withoutModifyingEnv' [Monad m] [MonadEnv m] [MonadFinally m] {α : Type} (x : m α) : m (α × Environment) := do diff --git a/src/Lean/ReducibilityAttrs.lean b/src/Lean/ReducibilityAttrs.lean index ab6bf712c3..f096fa2737 100644 --- a/src/Lean/ReducibilityAttrs.lean +++ b/src/Lean/ReducibilityAttrs.lean @@ -31,6 +31,7 @@ builtin_initialize reducibilityCoreExt : PersistentEnvExtension (Name × Reducib let r : Array (Name × ReducibilityStatus) := m.fold (fun a n p => a.push (n, p)) #[] r.qsort (fun a b => Name.quickLt a.1 b.1) statsFn := fun s => "reducibility attribute core extension" ++ Format.line ++ "number of local entries: " ++ format s.size + asyncMode := .async } builtin_initialize reducibilityExtraExt : SimpleScopedEnvExtension (Name × ReducibilityStatus) (SMap Name ReducibilityStatus) ← @@ -51,7 +52,7 @@ def getReducibilityStatusCore (env : Environment) (declName : Name) : Reducibili match (reducibilityCoreExt.getModuleEntries env modIdx).binSearch (declName, .semireducible) (fun a b => Name.quickLt a.1 b.1) with | some (_, status) => status | none => .semireducible - | none => (reducibilityCoreExt.getState env).find? declName |>.getD .semireducible + | none => (reducibilityCoreExt.findStateAsync env declName).find? declName |>.getD .semireducible private def setReducibilityStatusCore (env : Environment) (declName : Name) (status : ReducibilityStatus) (attrKind : AttributeKind) (currNamespace : Name) : Environment := if attrKind matches .global then @@ -60,7 +61,8 @@ private def setReducibilityStatusCore (env : Environment) (declName : Name) (sta -- Trying to set the attribute of a declaration defined in an imported module. reducibilityExtraExt.addEntry env (declName, status) | none => - -- + let _ : Inhabited Environment := ⟨env⟩ + assert! env.asyncMayContain declName reducibilityCoreExt.addEntry env (declName, status) else -- `scoped` and `local` must be handled by `reducibilityExtraExt`