From c658648ee8923a66ce3ff8bb79bbd4550dbb10b4 Mon Sep 17 00:00:00 2001 From: Sebastian Ullrich Date: Wed, 2 Apr 2025 18:19:12 +0200 Subject: [PATCH] refactor: split Lean.EnvironmentExtension from Lean.Environment (#7794) --- src/Lean.lean | 1 + src/Lean/AuxRecursor.lean | 2 +- src/Lean/Compiler/NoncomputableAttr.lean | 2 +- src/Lean/EnvExtension.lean | 161 +++++++++++++++++++++++ src/Lean/Environment.lean | 152 +-------------------- src/Lean/Modifiers.lean | 2 +- src/Lean/Namespace.lean | 2 +- src/Lean/ProjFns.lean | 2 +- 8 files changed, 168 insertions(+), 156 deletions(-) create mode 100644 src/Lean/EnvExtension.lean diff --git a/src/Lean.lean b/src/Lean.lean index e6a1cf9f9c..50aef24aff 100644 --- a/src/Lean.lean +++ b/src/Lean.lean @@ -40,3 +40,4 @@ import Lean.Replay import Lean.PrivateName import Lean.PremiseSelection import Lean.Namespace +import Lean.EnvExtension diff --git a/src/Lean/AuxRecursor.lean b/src/Lean/AuxRecursor.lean index c49a369828..4671af8bfc 100644 --- a/src/Lean/AuxRecursor.lean +++ b/src/Lean/AuxRecursor.lean @@ -4,7 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE. Authors: Leonardo de Moura -/ prelude -import Lean.Environment +import Lean.EnvExtension namespace Lean diff --git a/src/Lean/Compiler/NoncomputableAttr.lean b/src/Lean/Compiler/NoncomputableAttr.lean index a572182a4b..8d834da59e 100644 --- a/src/Lean/Compiler/NoncomputableAttr.lean +++ b/src/Lean/Compiler/NoncomputableAttr.lean @@ -4,7 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE. Authors: Leonardo de Moura -/ prelude -import Lean.Environment +import Lean.EnvExtension namespace Lean diff --git a/src/Lean/EnvExtension.lean b/src/Lean/EnvExtension.lean new file mode 100644 index 0000000000..80c945743e --- /dev/null +++ b/src/Lean/EnvExtension.lean @@ -0,0 +1,161 @@ +/- +Copyright (c) 2025 Lean FRO. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura, Sebastian Ullrich +-/ +prelude +import Lean.Environment + +/-! Further environment extension API; the primitives live in `Lean.Environment`. -/ + +namespace Lean + +/-- Simple `PersistentEnvExtension` that implements `exportEntriesFn` using a list of entries. -/ +def SimplePersistentEnvExtension (α σ : Type) := PersistentEnvExtension α α (List α × σ) + +@[specialize] def mkStateFromImportedEntries {α σ : Type} (addEntryFn : σ → α → σ) (initState : σ) (as : Array (Array α)) : σ := + as.foldl (fun r es => es.foldl (fun r e => addEntryFn r e) r) initState + +structure SimplePersistentEnvExtensionDescr (α σ : Type) where + name : Name := by exact decl_name% + addEntryFn : σ → α → σ + addImportedFn : Array (Array α) → σ + toArrayFn : List α → Array α := fun es => es.toArray + asyncMode : EnvExtension.AsyncMode := .mainOnly + replay? : Option ((newEntries : List α) → (newState : σ) → σ → List α × σ) := none + +/-- +Returns a function suitable for `SimplePersistentEnvExtensionDescr.replay?` that replays all new +entries onto the state using `addEntryFn`. `p` should filter out entries that have already been +added to the state by a prior replay of the same realizable constant. +-/ +def SimplePersistentEnvExtension.replayOfFilter (p : σ → α → Bool) + (addEntryFn : σ → α → σ) : List α → σ → σ → List α × σ := + fun newEntries _ s => + let newEntries := newEntries.filter (p s) + (newEntries, newEntries.foldl (init := s) addEntryFn) + +def registerSimplePersistentEnvExtension {α σ : Type} [Inhabited σ] (descr : SimplePersistentEnvExtensionDescr α σ) : IO (SimplePersistentEnvExtension α σ) := + registerPersistentEnvExtension { + name := descr.name, + mkInitial := pure ([], descr.addImportedFn #[]), + addImportedFn := fun as => pure ([], descr.addImportedFn as), + addEntryFn := fun s e => match s with + | (entries, s) => (e::entries, descr.addEntryFn s e), + exportEntriesFn := fun s => descr.toArrayFn s.1.reverse, + statsFn := fun s => format "number of local entries: " ++ format s.1.length + asyncMode := descr.asyncMode + replay? := descr.replay?.map fun replay oldState newState _ (entries, s) => + let newEntries := newState.1.take (newState.1.length - oldState.1.length) + let (newEntries, s) := replay newEntries newState.2 s + (newEntries ++ entries, s) + } + +namespace SimplePersistentEnvExtension + +instance {α σ : Type} [Inhabited σ] : Inhabited (SimplePersistentEnvExtension α σ) := + inferInstanceAs (Inhabited (PersistentEnvExtension α α (List α × σ))) + +/-- Get the list of values used to update the state of the given +`SimplePersistentEnvExtension` in the current file. -/ +def getEntries {α σ : Type} [Inhabited σ] (ext : SimplePersistentEnvExtension α σ) (env : Environment) : List α := + (PersistentEnvExtension.getState ext env).1 + +/-- Get the current state of the given `SimplePersistentEnvExtension`. -/ +def getState {α σ : Type} [Inhabited σ] (ext : SimplePersistentEnvExtension α σ) (env : Environment) + (asyncMode := ext.toEnvExtension.asyncMode) : σ := + (PersistentEnvExtension.getState (asyncMode := asyncMode) ext env).2 + +/-- Set the current state of the given `SimplePersistentEnvExtension`. This change is *not* persisted across files. -/ +def setState {α σ : Type} (ext : SimplePersistentEnvExtension α σ) (env : Environment) (s : σ) : Environment := + PersistentEnvExtension.modifyState ext env (fun ⟨entries, _⟩ => (entries, s)) + +/-- Modify the state of the given extension in the given environment by applying the given function. This change is *not* persisted across files. -/ +def modifyState {α σ : Type} (ext : SimplePersistentEnvExtension α σ) (env : Environment) (f : σ → σ) : Environment := + PersistentEnvExtension.modifyState ext env (fun ⟨entries, s⟩ => (entries, f s)) + +@[inherit_doc PersistentEnvExtension.findStateAsync] +def findStateAsync {α σ : Type} [Inhabited σ] (ext : SimplePersistentEnvExtension α σ) + (env : Environment) (declPrefix : Name) : σ := + PersistentEnvExtension.findStateAsync ext env declPrefix |>.2 + +end SimplePersistentEnvExtension + +/-- Environment extension for tagging declarations. + Declarations must only be tagged in the module where they were declared. -/ +def TagDeclarationExtension := SimplePersistentEnvExtension Name NameSet + +def mkTagDeclarationExtension (name : Name := by exact decl_name%) + (asyncMode : EnvExtension.AsyncMode := .mainOnly) : IO TagDeclarationExtension := + registerSimplePersistentEnvExtension { + name := name, + addImportedFn := fun _ => {}, + addEntryFn := fun s n => s.insert n, + toArrayFn := fun es => es.toArray.qsort Name.quickLt + asyncMode + } + +namespace TagDeclarationExtension + +instance : Inhabited TagDeclarationExtension := + inferInstanceAs (Inhabited (SimplePersistentEnvExtension Name NameSet)) + +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 => if ext.toEnvExtension.asyncMode matches .async then + (ext.findStateAsync env declName).contains declName + else + (ext.getState env).contains declName + +end TagDeclarationExtension + +/-- Environment extension for mapping declarations to values. + Declarations must only be inserted into the mapping in the module where they were declared. -/ + +structure MapDeclarationExtension (α : Type) extends PersistentEnvExtension (Name × α) (Name × α) (NameMap α) +deriving Inhabited + +def mkMapDeclarationExtension (name : Name := by exact decl_name%) : IO (MapDeclarationExtension α) := + .mk <$> registerPersistentEnvExtension { + name := name, + mkInitial := pure {} + addImportedFn := fun _ => pure {} + addEntryFn := fun s (n, v) => s.insert n v + exportEntriesFn := fun s => s.toArray + asyncMode := .async + replay? := some fun _ newState newConsts s => + newConsts.foldl (init := s) fun s c => + if let some a := newState.find? c then + s.insert c a + else s + } + +namespace MapDeclarationExtension + +def insert (ext : MapDeclarationExtension α) (env : Environment) (declName : Name) (val : α) : Environment := + have : Inhabited Environment := ⟨env⟩ + assert! env.getModuleIdxFor? declName |>.isNone -- See comment at `MapDeclarationExtension` + assert! env.asyncMayContain declName + ext.addEntry env (declName, val) + +def find? [Inhabited α] (ext : MapDeclarationExtension α) (env : Environment) (declName : Name) : Option α := + match env.getModuleIdxFor? declName with + | some modIdx => + match (ext.getModuleEntries env modIdx).binSearch (declName, default) (fun a b => Name.quickLt a.1 b.1) with + | some e => some e.2 + | none => none + | none => (ext.findStateAsync env declName).find? declName + +def contains [Inhabited α] (ext : MapDeclarationExtension α) (env : Environment) (declName : Name) : Bool := + match env.getModuleIdxFor? declName with + | some modIdx => (ext.getModuleEntries env modIdx).binSearchContains (declName, default) (fun a b => Name.quickLt a.1 b.1) + | none => (ext.findStateAsync env declName).contains declName + +end MapDeclarationExtension diff --git a/src/Lean/Environment.lean b/src/Lean/Environment.lean index 8040594ee9..4671c50277 100644 --- a/src/Lean/Environment.lean +++ b/src/Lean/Environment.lean @@ -1044,7 +1044,7 @@ def instantiateValueLevelParams! (c : ConstantInfo) (ls : List Level) : Expr := end ConstantInfo /-- -Async access mode for environment extensions used in `EnvironmentExtension.get/set/modifyState`. +Async access mode for environment extensions used in `EnvExtension.get/set/modifyState`. When modified in concurrent contexts, extensions may need to switch to a different mode than the default `mainOnly`, which will panic in such cases. The access mode is set at environment extension registration time but can be overriden when calling the mentioned functions in order to weaken it @@ -1465,156 +1465,6 @@ unsafe def registerPersistentEnvExtensionUnsafe {α β σ : Type} [Inhabited σ] @[implemented_by registerPersistentEnvExtensionUnsafe] opaque registerPersistentEnvExtension {α β σ : Type} [Inhabited σ] (descr : PersistentEnvExtensionDescr α β σ) : IO (PersistentEnvExtension α β σ) -/-- Simple `PersistentEnvExtension` that implements `exportEntriesFn` using a list of entries. -/ -def SimplePersistentEnvExtension (α σ : Type) := PersistentEnvExtension α α (List α × σ) - -@[specialize] def mkStateFromImportedEntries {α σ : Type} (addEntryFn : σ → α → σ) (initState : σ) (as : Array (Array α)) : σ := - as.foldl (fun r es => es.foldl (fun r e => addEntryFn r e) r) initState - -structure SimplePersistentEnvExtensionDescr (α σ : Type) where - name : Name := by exact decl_name% - addEntryFn : σ → α → σ - addImportedFn : Array (Array α) → σ - toArrayFn : List α → Array α := fun es => es.toArray - asyncMode : EnvExtension.AsyncMode := .mainOnly - replay? : Option ((newEntries : List α) → (newState : σ) → σ → List α × σ) := none - -/-- -Returns a function suitable for `SimplePersistentEnvExtensionDescr.replay?` that replays all new -entries onto the state using `addEntryFn`. `p` should filter out entries that have already been -added to the state by a prior replay of the same realizable constant. --/ -def SimplePersistentEnvExtension.replayOfFilter (p : σ → α → Bool) - (addEntryFn : σ → α → σ) : List α → σ → σ → List α × σ := - fun newEntries _ s => - let newEntries := newEntries.filter (p s) - (newEntries, newEntries.foldl (init := s) addEntryFn) - -def registerSimplePersistentEnvExtension {α σ : Type} [Inhabited σ] (descr : SimplePersistentEnvExtensionDescr α σ) : IO (SimplePersistentEnvExtension α σ) := - registerPersistentEnvExtension { - name := descr.name, - mkInitial := pure ([], descr.addImportedFn #[]), - addImportedFn := fun as => pure ([], descr.addImportedFn as), - addEntryFn := fun s e => match s with - | (entries, s) => (e::entries, descr.addEntryFn s e), - exportEntriesFn := fun s => descr.toArrayFn s.1.reverse, - statsFn := fun s => format "number of local entries: " ++ format s.1.length - asyncMode := descr.asyncMode - replay? := descr.replay?.map fun replay oldState newState _ (entries, s) => - let newEntries := newState.1.take (newState.1.length - oldState.1.length) - let (newEntries, s) := replay newEntries newState.2 s - (newEntries ++ entries, s) - } - -namespace SimplePersistentEnvExtension - -instance {α σ : Type} [Inhabited σ] : Inhabited (SimplePersistentEnvExtension α σ) := - inferInstanceAs (Inhabited (PersistentEnvExtension α α (List α × σ))) - -/-- Get the list of values used to update the state of the given -`SimplePersistentEnvExtension` in the current file. -/ -def getEntries {α σ : Type} [Inhabited σ] (ext : SimplePersistentEnvExtension α σ) (env : Environment) : List α := - (PersistentEnvExtension.getState ext env).1 - -/-- Get the current state of the given `SimplePersistentEnvExtension`. -/ -def getState {α σ : Type} [Inhabited σ] (ext : SimplePersistentEnvExtension α σ) (env : Environment) - (asyncMode := ext.toEnvExtension.asyncMode) : σ := - (PersistentEnvExtension.getState (asyncMode := asyncMode) ext env).2 - -/-- Set the current state of the given `SimplePersistentEnvExtension`. This change is *not* persisted across files. -/ -def setState {α σ : Type} (ext : SimplePersistentEnvExtension α σ) (env : Environment) (s : σ) : Environment := - PersistentEnvExtension.modifyState ext env (fun ⟨entries, _⟩ => (entries, s)) - -/-- Modify the state of the given extension in the given environment by applying the given function. This change is *not* persisted across files. -/ -def modifyState {α σ : Type} (ext : SimplePersistentEnvExtension α σ) (env : Environment) (f : σ → σ) : Environment := - PersistentEnvExtension.modifyState ext env (fun ⟨entries, s⟩ => (entries, f s)) - -@[inherit_doc PersistentEnvExtension.findStateAsync] -def findStateAsync {α σ : Type} [Inhabited σ] (ext : SimplePersistentEnvExtension α σ) - (env : Environment) (declPrefix : Name) : σ := - PersistentEnvExtension.findStateAsync ext env declPrefix |>.2 - -end SimplePersistentEnvExtension - -/-- Environment extension for tagging declarations. - Declarations must only be tagged in the module where they were declared. -/ -def TagDeclarationExtension := SimplePersistentEnvExtension Name NameSet - -def mkTagDeclarationExtension (name : Name := by exact decl_name%) - (asyncMode : EnvExtension.AsyncMode := .mainOnly) : IO TagDeclarationExtension := - registerSimplePersistentEnvExtension { - name := name, - addImportedFn := fun _ => {}, - addEntryFn := fun s n => s.insert n, - toArrayFn := fun es => es.toArray.qsort Name.quickLt - asyncMode - } - -namespace TagDeclarationExtension - -instance : Inhabited TagDeclarationExtension := - inferInstanceAs (Inhabited (SimplePersistentEnvExtension Name NameSet)) - -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 => if ext.toEnvExtension.asyncMode matches .async then - (ext.findStateAsync env declName).contains declName - else - (ext.getState env).contains declName - -end TagDeclarationExtension - -/-- Environment extension for mapping declarations to values. - Declarations must only be inserted into the mapping in the module where they were declared. -/ - -structure MapDeclarationExtension (α : Type) extends PersistentEnvExtension (Name × α) (Name × α) (NameMap α) -deriving Inhabited - -def mkMapDeclarationExtension (name : Name := by exact decl_name%) : IO (MapDeclarationExtension α) := - .mk <$> registerPersistentEnvExtension { - name := name, - mkInitial := pure {} - addImportedFn := fun _ => pure {} - addEntryFn := fun s (n, v) => s.insert n v - exportEntriesFn := fun s => s.toArray - asyncMode := .async - replay? := some fun _ newState newConsts s => - newConsts.foldl (init := s) fun s c => - if let some a := newState.find? c then - s.insert c a - else s - } - -namespace MapDeclarationExtension - -def insert (ext : MapDeclarationExtension α) (env : Environment) (declName : Name) (val : α) : Environment := - have : Inhabited Environment := ⟨env⟩ - assert! env.getModuleIdxFor? declName |>.isNone -- See comment at `MapDeclarationExtension` - assert! env.asyncMayContain declName - ext.addEntry env (declName, val) - -def find? [Inhabited α] (ext : MapDeclarationExtension α) (env : Environment) (declName : Name) : Option α := - match env.getModuleIdxFor? declName with - | some modIdx => - match (ext.getModuleEntries env modIdx).binSearch (declName, default) (fun a b => Name.quickLt a.1 b.1) with - | some e => some e.2 - | none => none - | none => (ext.findStateAsync env declName).find? declName - -def contains [Inhabited α] (ext : MapDeclarationExtension α) (env : Environment) (declName : Name) : Bool := - match env.getModuleIdxFor? declName with - | some modIdx => (ext.getModuleEntries env modIdx).binSearchContains (declName, default) (fun a b => Name.quickLt a.1 b.1) - | none => (ext.findStateAsync env declName).contains declName - -end MapDeclarationExtension - @[extern "lean_save_module_data"] opaque saveModuleData (fname : @& System.FilePath) (mod : @& Name) (data : @& ModuleData) : IO Unit @[extern "lean_read_module_data"] diff --git a/src/Lean/Modifiers.lean b/src/Lean/Modifiers.lean index a95d82cf8b..7ee15352b7 100644 --- a/src/Lean/Modifiers.lean +++ b/src/Lean/Modifiers.lean @@ -4,7 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE. Authors: Leonardo de Moura -/ prelude -import Lean.Environment +import Lean.EnvExtension import Lean.PrivateName namespace Lean diff --git a/src/Lean/Namespace.lean b/src/Lean/Namespace.lean index fbb984d4a9..30d3432027 100644 --- a/src/Lean/Namespace.lean +++ b/src/Lean/Namespace.lean @@ -4,7 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE. Authors: Leonardo de Moura -/ prelude -import Lean.Environment +import Lean.EnvExtension namespace Lean diff --git a/src/Lean/ProjFns.lean b/src/Lean/ProjFns.lean index 9873c4f4a4..2cd82d5a15 100644 --- a/src/Lean/ProjFns.lean +++ b/src/Lean/ProjFns.lean @@ -4,7 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE. Authors: Leonardo de Moura -/ prelude -import Lean.Environment +import Lean.EnvExtension namespace Lean