diff --git a/src/Lean/ReducibilityAttrs.lean b/src/Lean/ReducibilityAttrs.lean index a058740bd9..5fcf8194d1 100644 --- a/src/Lean/ReducibilityAttrs.lean +++ b/src/Lean/ReducibilityAttrs.lean @@ -5,6 +5,7 @@ Authors: Leonardo de Moura -/ prelude import Lean.Attributes +import Lean.ScopedEnvExtension namespace Lean @@ -15,34 +16,98 @@ inductive ReducibilityStatus where | reducible | semireducible | irreducible deriving Inhabited, Repr -/-- -Environment extension for storing the reducibility attribute for definitions. --/ -builtin_initialize reducibilityAttrs : EnumAttributes ReducibilityStatus ← - registerEnumAttributes - [(`reducible, "reducible", ReducibilityStatus.reducible), - (`semireducible, "semireducible", ReducibilityStatus.semireducible), - (`irreducible, "irreducible", ReducibilityStatus.irreducible)] +builtin_initialize reducibilityCoreExt : PersistentEnvExtension (Name × ReducibilityStatus) (Name × ReducibilityStatus) (NameMap ReducibilityStatus) ← + registerPersistentEnvExtension { + name := `reducibilityCore + mkInitial := pure {} + addImportedFn := fun _ _ => pure {} + addEntryFn := fun (s : NameMap ReducibilityStatus) (p : Name × ReducibilityStatus) => s.insert p.1 p.2 + exportEntriesFn := fun m => + 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 + } + +builtin_initialize reducibilityExtraExt : SimpleScopedEnvExtension (Name × ReducibilityStatus) (SMap Name ReducibilityStatus) ← + registerSimpleScopedEnvExtension { + name := `reducibilityExtra + initial := {} + addEntry := fun d (declName, status) => d.insert declName status + finalizeImport := fun d => d.switch + } @[export lean_get_reducibility_status] -private def getReducibilityStatusImp (env : Environment) (declName : Name) : ReducibilityStatus := - match reducibilityAttrs.getValue env declName with - | some s => s - | none => ReducibilityStatus.semireducible +def getReducibilityStatusCore (env : Environment) (declName : Name) : ReducibilityStatus := + let m := reducibilityExtraExt.getState env + if let some status := m.find? declName then + status + else match env.getModuleIdxFor? declName with + | some modIdx => + 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 + +def setReducibilityStatusCore (env : Environment) (declName : Name) (status : ReducibilityStatus) (attrKind : AttributeKind) (currNamespace : Name) : Environment := + if attrKind matches .global then + match env.getModuleIdxFor? declName with + | some _ => + -- Trying to set the attribute of a declaration defined in an imported module. + reducibilityExtraExt.addEntry env (declName, status) + | none => + -- + reducibilityCoreExt.addEntry env (declName, status) + else + -- `scoped` and `local` must be handled by `reducibilityExtraExt` + reducibilityExtraExt.addCore env (declName, status) attrKind currNamespace @[export lean_set_reducibility_status] -private def setReducibilityStatusImp (env : Environment) (declName : Name) (s : ReducibilityStatus) : Environment := - match reducibilityAttrs.setValue env declName s with - | Except.ok env => env - | _ => env -- TODO(Leo): we should extend EnumAttributes.setValue in the future and ensure it never fails +private def setReducibilityStatusImp (env : Environment) (declName : Name) (status : ReducibilityStatus) : Environment := + setReducibilityStatusCore env declName status .global .anonymous + +builtin_initialize + registerBuiltinAttribute { + ref := by exact decl_name% + name := `irreducible + descr := "irreducible declaration" + add := fun declName stx attrKind => do + Attribute.Builtin.ensureNoArgs stx + let ns ← getCurrNamespace + modifyEnv fun env => setReducibilityStatusCore env declName .irreducible attrKind ns + applicationTime := .afterTypeChecking + } + +builtin_initialize + registerBuiltinAttribute { + ref := by exact decl_name% + name := `reducible + descr := "reducible declaration" + add := fun declName stx attrKind => do + Attribute.Builtin.ensureNoArgs stx + let ns ← getCurrNamespace + modifyEnv fun env => setReducibilityStatusCore env declName .reducible attrKind ns + applicationTime := .afterTypeChecking + } + +builtin_initialize + registerBuiltinAttribute { + ref := by exact decl_name% + name := `semireducible + descr := "semireducible declaration" + add := fun declName stx attrKind => do + Attribute.Builtin.ensureNoArgs stx + let ns ← getCurrNamespace + modifyEnv fun env => setReducibilityStatusCore env declName .reducible attrKind ns + applicationTime := .afterTypeChecking + } /-- Return the reducibility attribute for the given declaration. -/ def getReducibilityStatus [Monad m] [MonadEnv m] (declName : Name) : m ReducibilityStatus := do - return getReducibilityStatusImp (← getEnv) declName + return getReducibilityStatusCore (← getEnv) declName /-- Set the reducibility attribute for the given declaration. -/ def setReducibilityStatus [Monad m] [MonadEnv m] (declName : Name) (s : ReducibilityStatus) : m Unit := do - modifyEnv fun env => setReducibilityStatusImp env declName s + modifyEnv fun env => setReducibilityStatusCore env declName s .global .anonymous /-- Set the given declaration as `[reducible]` -/ def setReducibleAttribute [Monad m] [MonadEnv m] (declName : Name) : m Unit := do @@ -51,13 +116,13 @@ def setReducibleAttribute [Monad m] [MonadEnv m] (declName : Name) : m Unit := d /-- Return `true` if the given declaration has been marked as `[reducible]`. -/ def isReducible [Monad m] [MonadEnv m] (declName : Name) : m Bool := do match (← getReducibilityStatus declName) with - | ReducibilityStatus.reducible => return true + | .reducible => return true | _ => return false /-- Return `true` if the given declaration has been marked as `[irreducible]` -/ def isIrreducible [Monad m] [MonadEnv m] (declName : Name) : m Bool := do match (← getReducibilityStatus declName) with - | ReducibilityStatus.irreducible => return true + | .irreducible => return true | _ => return false end Lean diff --git a/src/Lean/ScopedEnvExtension.lean b/src/Lean/ScopedEnvExtension.lean index 917e77d497..68e2e7ec3d 100644 --- a/src/Lean/ScopedEnvExtension.lean +++ b/src/Lean/ScopedEnvExtension.lean @@ -146,11 +146,15 @@ def ScopedEnvExtension.addLocalEntry (ext : ScopedEnvExtension α β σ) (env : let top := { top with state := ext.descr.addEntry top.state b } ext.ext.setState env { s with stateStack := top :: states } -def ScopedEnvExtension.add [Monad m] [MonadResolveName m] [MonadEnv m] (ext : ScopedEnvExtension α β σ) (b : β) (kind := AttributeKind.global) : m Unit := do +def ScopedEnvExtension.addCore (env : Environment) (ext : ScopedEnvExtension α β σ) (b : β) (kind : AttributeKind) (namespaceName : Name) : Environment := match kind with - | AttributeKind.global => modifyEnv (ext.addEntry · b) - | AttributeKind.local => modifyEnv (ext.addLocalEntry · b) - | AttributeKind.scoped => modifyEnv (ext.addScopedEntry · (← getCurrNamespace) b) + | AttributeKind.global => ext.addEntry env b + | AttributeKind.local => ext.addLocalEntry env b + | AttributeKind.scoped => ext.addScopedEntry env namespaceName b + +def ScopedEnvExtension.add [Monad m] [MonadResolveName m] [MonadEnv m] (ext : ScopedEnvExtension α β σ) (b : β) (kind := AttributeKind.global) : m Unit := do + let ns ← getCurrNamespace + modifyEnv (ext.addCore · b kind ns) def ScopedEnvExtension.getState [Inhabited σ] (ext : ScopedEnvExtension α β σ) (env : Environment) : σ := match ext.ext.getState env |>.stateStack with