feat: add moduleDocExt

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

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