feat(library/init/lean/attributes): add pushScope/popScope, missing APIs, and @[export]

library/init/lean/attributes.lean

@@ -40,19 +40,39 @@ constant scopeManagerExt : SimplePersistentEnvExtension Name ScopeManagerState :
 
 namespace Environment
 
+@[export lean.get_namespaces_core]
 def getNamespaces (env : Environment) : List Name :=
 (scopeManagerExt.getState env).namespaces
 
+def getNamespaceSet (env : Environment) : NameSet :=
+(scopeManagerExt.getState env).allNamespaces
+
+@[export lean.is_namespace_core]
+def isNamespace (env : Environment) (n : Name) : Bool :=
+env.getNamespaceSet.contains n
+
+@[export lean.in_section_core]
+def inSection (env : Environment) : Bool :=
+match (scopeManagerExt.getState env).isNamespace with
+| (b::_) := !b
+| _      := false
+
+def hasOpenScopes (env : Environment) : Bool :=
+!env.getNamespaces.isEmpty
+
+@[export lean.get_namespace_core]
 def getNamespace (env : Environment) : Name :=
 match env.getNamespaces with
 | (n::_) := n
 | _      := Name.anonymous
 
+@[export lean.get_scope_header_core]
 def getScopeHeader (env : Environment) : Name :=
 match (scopeManagerExt.getState env).headers with
 | (n::_) := n
 | _      := Name.anonymous
 
+@[export lean.to_valid_namespace_core]
 def toValidNamespace (env : Environment) (n : Name) : Option Name :=
 let s := scopeManagerExt.getState env in
 if s.allNamespaces.contains n then some n
@@ -64,6 +84,33 @@ else s.namespaces.foldl
       if s.allNamespaces.contains c then some c else none)
   none
 
+def registerNamespaceAux (env : Environment) (n : Name) : Environment :=
+if env.getNamespaceSet.contains n then env else scopeManagerExt.addEntry env n
+
+@[export lean.register_namespace_core]
+def registerNamespace : Environment → Name → Environment
+| env n@(Name.mkString p _) := registerNamespace (registerNamespaceAux env n) p
+| env _ := env
+
+def pushScopeCore (env : Environment) (header : Name) (isNamespace : Bool) : Environment :=
+let ns    := env.getNamespace in
+let newNs := if isNamespace then ns ++ header else ns in
+let env   := env.registerNamespaceAux newNs in
+let env   := scopeManagerExt.modifyState env $ λ s,
+  { headers     := header :: s.headers,
+    namespaces  := newNs :: s.namespaces,
+    isNamespace := isNamespace :: s.isNamespace,
+    .. s } in
+env
+
+def popScopeCore (env : Environment) : Environment :=
+if env.getNamespaces.isEmpty then env
+else scopeManagerExt.modifyState env $ λ s,
+  { headers     := s.headers.tail,
+    namespaces  := s.namespaces.tail,
+    isNamespace := s.isNamespace.tail,
+    .. s }
+
 end Environment
 
 structure AttributeImpl :=
@@ -84,13 +131,20 @@ IO.mkRef {}
 @[init mkAttributeMapRef]
 constant attributeMapRef : IO.Ref (HashMap Name AttributeImpl) := default _
 
+def mkAttributeArrayRef : IO (IO.Ref (Array AttributeImpl)) :=
+IO.mkRef Array.empty
+
+@[init mkAttributeArrayRef]
+constant attributeArrayRef : IO.Ref (Array AttributeImpl) := default _
+
 /- Low level attribute registration function. -/
 def registerAttribute (attr : AttributeImpl) : IO Unit :=
 do m ← attributeMapRef.get,
    when (m.contains attr.name) $ throw (IO.userError ("invalid attribute declaration, '" ++ toString attr.name ++ "' has already been used")),
    initializing ← IO.initializing,
    unless initializing $ throw (IO.userError ("failed to register attribute, attributes can only be registered during initialization")),
-   attributeMapRef.modify (λ m, m.insert attr.name attr)
+   attributeMapRef.modify (λ m, m.insert attr.name attr),
+   attributeArrayRef.modify (λ attrs, attrs.push attr)
 
 namespace Environment
 
@@ -146,18 +200,24 @@ do attr ← env.getAttributeImpl attrName,
 
 /- Activate all scoped attributes at `scope` -/
 def activateScopedAttributes (env : Environment) (scope : Name) : IO Environment :=
-do m ← attributeMapRef.get,
-   m.mfold (λ env _ attr, attr.activateScoped env scope) env
+do attrs ← attributeArrayRef.get,
+   attrs.mfoldl (λ env attr, attr.activateScoped env scope) env
 
 /- We use this function to implement commands `namespace foo` and `section foo`.
    It activates scoped attributes in the new resulting namespace. -/
+@[export lean.push_scope_core]
 def pushScope (env : Environment) (header : Name) (isNamespace : Bool) : IO Environment :=
-pure env -- TODO
+do let env := env.pushScopeCore header isNamespace,
+   let ns  := env.getNamespace,
+   attrs ← attributeArrayRef.get,
+   attrs.mfoldl (λ env attr, do env ← attr.pushScope env, if isNamespace then attr.activateScoped env ns else pure env) env
 
 /- We use this function to implement commands `end foo` for closing namespaces and sections. -/
+@[export lean.pop_scope_core]
 def popScope (env : Environment) : IO Environment :=
-pure env -- TODO
+do let env := env.popScopeCore,
+   attrs ← attributeArrayRef.get,
+   attrs.mfoldl (λ env attr, attr.popScope env) env
 
 end Environment
-
 end Lean

library/init/lean/environment.lean

@@ -202,6 +202,12 @@ ext.toEnvExtension.modifyState env $ λ s,
 def getState {α σ : Type} (ext : PersistentEnvExtension α σ) (env : Environment) : σ :=
 (ext.toEnvExtension.getState env).state
 
+def setState {α σ : Type} (ext : PersistentEnvExtension α σ) (env : Environment) (s : σ) : Environment :=
+ext.toEnvExtension.modifyState env $ λ ps, { state := s, .. ps }
+
+def modifyState {α σ : Type} (ext : PersistentEnvExtension α σ) (env : Environment) (f : σ → σ) : Environment :=
+ext.toEnvExtension.modifyState env $ λ ps, { state := f (ps.state), .. ps }
+
 end PersistentEnvExtension
 
 private def mkPersistentEnvExtensionsRef : IO (IO.Ref (Array (PersistentEnvExtension EnvExtensionEntry EnvExtensionState))) :=
@@ -273,6 +279,12 @@ def getEntries {α σ : Type} (ext : SimplePersistentEnvExtension α σ) (env :
 def getState {α σ : Type} (ext : SimplePersistentEnvExtension α σ) (env : Environment) : σ :=
 (PersistentEnvExtension.getState ext env).2
 
+def setState {α σ : Type} (ext : SimplePersistentEnvExtension α σ) (env : Environment) (s : σ) : Environment :=
+PersistentEnvExtension.modifyState ext env (λ ⟨entries, _⟩, (entries, s))
+
+def modifyState {α σ : Type} (ext : SimplePersistentEnvExtension α σ) (env : Environment) (f : σ → σ) : Environment :=
+PersistentEnvExtension.modifyState ext env (λ ⟨entries, s⟩, (entries, f s))
+
 end SimplePersistentEnvExtension
 
 /- API for creating extensions in C++.