From 3d485bf3756fa610326705b1d4036d05dcdb326f Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Fri, 6 Aug 2021 13:40:54 -0700
Subject: [PATCH] feat: add `moduleDocExt`

---
 src/Lean/DocString.lean   | 18 +++++++++++++++++-
 src/Lean/Environment.lean | 11 +++++++----
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/src/Lean/DocString.lean b/src/Lean/DocString.lean
index 4528926da6..2e7555f81f 100644
--- a/src/Lean/DocString.lean
+++ b/src/Lean/DocString.lean
@@ -7,7 +7,7 @@ import Lean.MonadEnv
 
 namespace Lean
 
-builtin_initialize docStringExt : MapDeclarationExtension String ← mkMapDeclarationExtension `docstring
+private builtin_initialize docStringExt : MapDeclarationExtension String ← mkMapDeclarationExtension `docstring
 
 def addDocString [MonadEnv m] (declName : Name) (docString : String) : m Unit :=
   modifyEnv fun env => docStringExt.insert env declName docString
@@ -20,4 +20,20 @@ def addDocString' [Monad m] [MonadEnv m] (declName : Name) (docString? : Option
 def findDocString? [Monad m] [MonadEnv m] (declName : Name) : m (Option String) :=
   return docStringExt.find? (← getEnv) declName
 
+private builtin_initialize moduleDocExt : SimplePersistentEnvExtension String (Std.PersistentArray String) ← registerSimplePersistentEnvExtension {
+  name          := `moduleDocExt
+  addImportedFn := fun _ => {}
+  addEntryFn    := fun s e => s.push e
+  toArrayFn     := fun es => es.toArray
+}
+
+def addMainModuleDoc (env : Environment) (doc : String) : Environment :=
+  moduleDocExt.addEntry env doc
+
+def getMainModuleDoc (env : Environment)  : Std.PersistentArray String :=
+  moduleDocExt.getState env
+
+def getModuleDoc? (env : Environment) (moduleName : Name) : Option (Array String) :=
+  env.getModuleIdx? moduleName |>.map fun modIdx => moduleDocExt.getModuleEntries env modIdx
+
 end Lean
diff --git a/src/Lean/Environment.lean b/src/Lean/Environment.lean
index bd5fbee994..6653f9a483 100644
--- a/src/Lean/Environment.lean
+++ b/src/Lean/Environment.lean
@@ -102,14 +102,17 @@ private def isQuotInit (env : Environment) : Bool :=
 private def getTrustLevel (env : Environment) : UInt32 :=
   env.header.trustLevel
 
-def getModuleIdxFor? (env : Environment) (c : Name) : Option ModuleIdx :=
-  env.const2ModIdx.find? c
+def getModuleIdxFor? (env : Environment) (declName : Name) : Option ModuleIdx :=
+  env.const2ModIdx.find? declName
 
-def isConstructor (env : Environment) (c : Name) : Bool :=
-  match env.find? c with
+def isConstructor (env : Environment) (declName : Name) : Bool :=
+  match env.find? declName with
   | ConstantInfo.ctorInfo _ => true
   | _                       => false
 
+def getModuleIdx? (env : Environment) (moduleName : Name) : Option ModuleIdx :=
+  env.header.moduleNames.findIdx? (. == moduleName)
+
 end Environment
 
 inductive KernelException where