From 3145994e360c0db40d0fd80e0ecfe60652d6e221 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Wed, 24 Jul 2019 13:24:02 -0700
Subject: [PATCH] feat(library/init/lean/elaborator/basic): register
 declaration name prefixes as namespaces at `Environment.add`

---
 library/init/lean/elaborator/basic.lean | 16 -----------
 library/init/lean/environment.lean      | 36 +++++++++++++++++++++++--
 2 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/library/init/lean/elaborator/basic.lean b/library/init/lean/elaborator/basic.lean
index e7a62b76ee..eb5cff7574 100644
--- a/library/init/lean/elaborator/basic.lean
+++ b/library/init/lean/elaborator/basic.lean
@@ -367,22 +367,6 @@ def getOpenDecls : Elab (List OpenDecl) :=
 do s ← get;
    pure s.scopes.head.openDecls
 
-def regNamespacesExtension : IO (SimplePersistentEnvExtension Name NameSet) :=
-registerSimplePersistentEnvExtension {
-  name            := `namespaces,
-  addImportedFn   := fun as => mkStateFromImportedEntries NameSet.insert {} as,
-  addEntryFn      := fun s n => s.insert n
-}
-
-@[init regNamespacesExtension]
-constant namespacesExt : SimplePersistentEnvExtension Name NameSet := default _
-
-def registerNamespace (env : Environment) (n : Name) : Environment :=
-if (namespacesExt.getState env).contains n then env else namespacesExt.addEntry env n
-
-def isNamespace (env : Environment) (n : Name) : Bool :=
-(namespacesExt.getState env).contains n
-
 def getNamespace : Elab Name :=
 do s ← get;
    match s.scopes with
diff --git a/library/init/lean/environment.lean b/library/init/lean/environment.lean
index 09ead92d8d..24bda16cd0 100644
--- a/library/init/lean/environment.lean
+++ b/library/init/lean/environment.lean
@@ -43,8 +43,7 @@ namespace Environment
 instance : Inhabited Environment :=
 ⟨{ const2ModIdx := {}, constants := {}, extensions := Array.empty }⟩
 
-@[export lean.environment_add_core]
-def add (env : Environment) (cinfo : ConstantInfo) : Environment :=
+def addAux (env : Environment) (cinfo : ConstantInfo) : Environment :=
 { constants := env.constants.insert cinfo.name cinfo, .. env }
 
 @[export lean.environment_find_core]
@@ -481,8 +480,41 @@ env ← finalizePersistentExtensions env;
 env ← mods.miterate env $ fun _ mod env => performModifications env mod.serialized;
 pure env
 
+def regNamespacesExtension : IO (SimplePersistentEnvExtension Name NameSet) :=
+registerSimplePersistentEnvExtension {
+  name            := `namespaces,
+  addImportedFn   := fun as => mkStateFromImportedEntries NameSet.insert {} as,
+  addEntryFn      := fun s n => s.insert n
+}
+
+@[init regNamespacesExtension]
+constant namespacesExt : SimplePersistentEnvExtension Name NameSet := default _
+
+def registerNamespace (env : Environment) (n : Name) : Environment :=
+if (namespacesExt.getState env).contains n then env else namespacesExt.addEntry env n
+
+def isNamespace (env : Environment) (n : Name) : Bool :=
+(namespacesExt.getState env).contains n
+
+def getNamespaceSet (env : Environment) : NameSet :=
+namespacesExt.getState env
+
 namespace Environment
 
+private def isNamespaceName : Name → Bool
+| (Name.mkString Name.anonymous _) := true
+| (Name.mkString p _)              := isNamespaceName p
+| _                                := false
+
+private def registerNamePrefixes : Environment → Name → Environment
+| env (Name.mkString p _) := if isNamespaceName p then registerNamePrefixes (registerNamespace env p) p else env
+| env _                   := env
+
+@[export lean.environment_add_core]
+def add (env : Environment) (cinfo : ConstantInfo) : Environment :=
+let env := registerNamePrefixes env cinfo.name;
+env.addAux cinfo
+
 @[export lean.display_stats_core]
 def displayStats (env : Environment) : IO Unit :=
 do