test: for flexible reducibility attributes

- We can set `[irreducible]`, `[semireducible]`, and `[reducible]` for
imported declarations.
- Support for `scoped` and `local` versions

TODO: discuss whether we need all this power after we add the module
system.

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

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