feat(tactic): make docstrings of imported modules accessible

Author: Vtec234

library/init/meta/tactic.lean

@@ -543,7 +543,7 @@ It updates the environment in the tactic_state, and returns an expression of the
 where l_i's and a_j's are the collected dependencies.
 -/
 meta constant add_aux_decl (c : name) (type : expr) (val : expr) (is_lemma : bool) : tactic expr
-meta constant module_doc_strings : tactic (list (option name × string))
+meta constant module_doc_strings : tactic (list (option string × string))
 /-- Set attribute `attr_name` for constant `c_name` with the given priority.
    If the priority is none, then use default -/
 meta constant set_basic_attribute (attr_name : name) (c_name : name) (persistent := ff) (prio : option nat := none) : tactic unit

src/library/documentation.cpp

@@ -15,7 +15,7 @@ Author: Leonardo de Moura
 namespace lean {
 struct documentation_ext : public environment_extension {
     /** Doc string for the current module being processed. It does not include imported doc strings. */
-    list<doc_entry>       m_module_doc;
+    optional<std::string> m_module_doc;
     /** Doc strings for declarations (including imported ones). We store doc_strings for declarations in the .olean files. */
     name_map<std::string> m_doc_string_map;
 };
@@ -36,25 +36,16 @@ static environment update(environment const & env, documentation_ext const & ext
 struct doc_modification : public modification {
     LEAN_MODIFICATION("doc")
 
-    /** If non-empty, this is a doc on a declaration. Otherwise,
-      * it's a doc on the whole module. */
     name m_decl;
     std::string m_doc;
 
     doc_modification() {}
-    /** A docstring for the entire module. */
-    doc_modification(std::string const & doc) : m_decl(""), m_doc(doc) {}
     /** A docstring for a declaration in the module. */
     doc_modification(name const & decl, std::string const & doc) : m_decl(decl), m_doc(doc) {}
 
     void perform(environment & env) const override {
         auto ext = get_extension(env);
-        if (m_decl != "") {
-            ext.m_doc_string_map.insert(m_decl, m_doc);
-        }
-        // Note that we do NOT add anything to `m_module_doc` here, because `perform` is called
-        // when applying modifications from imported .olean modules whose doc strings are NOT
-        // part of the current module's documentation.
+        ext.m_doc_string_map.insert(m_decl, m_doc);
         env = update(env, ext);
     }
 
@@ -179,9 +170,13 @@ static std::string process_doc(std::string s) {
 environment add_module_doc_string(environment const & env, std::string doc) {
     doc = process_doc(doc);
     auto ext = get_extension(env);
-    ext.m_module_doc = cons(doc_entry(doc), ext.m_module_doc);
+    if (ext.m_module_doc) {
+        *ext.m_module_doc += "\n" + doc;
+    } else {
+        ext.m_module_doc = doc;
+    }
     auto new_env = update(env, ext);
-    return module::add(new_env, std::make_shared<doc_modification>(doc));
+    return module::add_doc_string(new_env, doc);
 }
 
 environment add_doc_string(environment const & env, name const & n, std::string doc) {
@@ -191,7 +186,6 @@ environment add_doc_string(environment const & env, name const & n, std::string
         throw exception(sstream() << "environment already contains a doc string for '" << n << "'");
     }
     ext.m_doc_string_map.insert(n, doc);
-    ext.m_module_doc = cons(doc_entry(n, doc), ext.m_module_doc);
     auto new_env = update(env, ext);
     return module::add(new_env, std::make_shared<doc_modification>(n, doc));
 }
@@ -204,10 +198,15 @@ optional<std::string> get_doc_string(environment const & env, name const & n) {
         return optional<std::string>();
 }
 
-void get_module_doc_strings(environment const & env, buffer<doc_entry> & result) {
+void get_module_doc_strings(environment const & env, buffer<mod_doc_entry> & result) {
     auto ext = get_extension(env);
-    to_buffer(ext.m_module_doc, result);
-    std::reverse(result.begin(), result.end());
+    auto const & mod_docs = module::get_doc_strings(env);
+    for (auto const & pr : mod_docs) {
+        result.push_back({ optional<std::string>{ pr.first }, pr.second });
+    }
+    if (ext.m_module_doc) {
+        result.push_back({ {}, *ext.m_module_doc });
+    }
 }
 
 void initialize_documentation() {

src/library/documentation.h

@@ -8,19 +8,14 @@ Author: Leonardo de Moura
 #include <string>
 #include "kernel/environment.h"
 namespace lean {
-class doc_entry {
-    optional<name> m_decl_name;
-    std::string    m_doc;
-public:
-    doc_entry(std::string const & doc):m_doc(doc) {}
-    doc_entry(name const & decl_name, std::string const & doc):m_decl_name(decl_name), m_doc(doc) {}
-    optional<name> const & get_decl_name() const { return m_decl_name; }
-    std::string const & get_doc() const { return m_doc; }
+struct mod_doc_entry {
+    optional<std::string> m_mod_name;
+    std::string           m_doc;
 };
 environment add_module_doc_string(environment const & env, std::string doc);
 environment add_doc_string(environment const & env, name const & n, std::string);
 optional<std::string> get_doc_string(environment const & env, name const & n);
-void get_module_doc_strings(environment const & env, buffer<doc_entry> & result);
+void get_module_doc_strings(environment const & env, buffer<mod_doc_entry> & result);
 void initialize_documentation();
 void finalize_documentation();
 }

src/library/module.cpp

@@ -86,6 +86,11 @@ struct module_ext : public environment_extension {
     list<name>        m_module_decls;
     name_set          m_module_defs;
     name_set          m_imported;
+    /** Top-level doc strings for modules which have them. Lean doesn't have a notion
+     * of module different from that of a source file, so we use file names to index
+     * the docstring map. */
+    // TODO(Vtec234): lean::rb_map misbehaves here for some reason
+    std::map<std::string, std::string> m_module_docs;
     // Map from declaration name to olean file where it was defined
     name_map<std::string>     m_decl2olean;
     name_map<pos_info>        m_decl2pos_info;
@@ -416,6 +421,30 @@ struct quot_modification : public modification {
     }
 };
 
+struct mod_doc_modification : public modification {
+    LEAN_MODIFICATION("mod_doc")
+
+    std::string m_doc;
+
+    mod_doc_modification() {}
+    /** A docstring for a declaration in the module. */
+    mod_doc_modification(std::string const & doc) : m_doc(doc) {}
+
+    void perform(environment & env) const override {
+        // No-op. See `import_module` for actual action
+    }
+
+    void serialize(serializer & s) const override {
+        s << m_doc;
+    }
+
+    static std::shared_ptr<modification const> deserialize(deserializer & d) {
+        std::string doc;
+        d >> doc;
+        return std::make_shared<mod_doc_modification>(doc);
+    }
+};
+
 namespace module {
 environment add(environment const & env, std::shared_ptr<modification const> const & modf) {
     module_ext ext = get_extension(env);
@@ -506,6 +535,14 @@ environment add(environment const & env, certified_declaration const & d) {
     return add_decl_pos_info(new_env, _d.get_name());
 }
 
+environment add_doc_string(environment const & env, std::string const & doc) {
+    return add(env, std::make_shared<mod_doc_modification>(doc));
+}
+
+std::map<std::string, std::string> const & get_doc_strings(environment const & env) {
+    return get_extension(env).m_module_docs;
+}
+
 bool is_definition(environment const & env, name const & n) {
     module_ext const & ext = get_extension(env);
     return ext.m_module_defs.contains(n);
@@ -693,6 +730,14 @@ void import_module(modification_list const & modifications, std::string const &
             env = add_decl_olean(env, dm->m_decl.get_name(), file_name);
         } else if (auto im = dynamic_cast<inductive_modification const *>(m.get())) {
             env = add_decl_olean(env, im->m_decl.get_decl().m_name, file_name);
+        } else if (auto mdm = dynamic_cast<mod_doc_modification const *>(m.get())) {
+            auto ext = get_extension(env);
+            if (ext.m_module_docs.count(file_name) != 0) {
+                ext.m_module_docs[file_name] += "\n" + mdm->m_doc;
+            } else {
+                ext.m_module_docs[file_name] = mdm->m_doc;
+            }
+            env = update(env, ext);
         }
     }
 }
@@ -724,13 +769,15 @@ void initialize_module() {
     inductive_modification::init();
     quot_modification::init();
     pos_info_mod::init();
+    mod_doc_modification::init();
 }
 
 void finalize_module() {
     quot_modification::finalize();
     pos_info_mod::finalize();
     inductive_modification::finalize();
     decl_modification::finalize();
+    mod_doc_modification::finalize();
     delete g_object_readers;
     delete g_ext;
 }

src/library/module.h

@@ -9,6 +9,7 @@ Author: Leonardo de Moura
 #include <iostream>
 #include <utility>
 #include <vector>
+#include <map>
 #include "util/serializer.h"
 #include "util/optional.h"
 #include "kernel/pos_info_provider.h"
@@ -143,6 +144,12 @@ environment add_and_perform(environment const & env, std::shared_ptr<modificatio
 /** \brief Add the given declaration to the environment, and mark it to be exported. */
 environment add(environment const & env, certified_declaration const & d);
 
+/** \brief Adds a module-level doc to the current module. */
+environment add_doc_string(environment const & env, std::string const & doc);
+
+/** \brief Returns the map of module-level docs indexed by source file name. */
+std::map<std::string, std::string> const & get_doc_strings(environment const & env);
+
 /** \brief Return true iff \c n is a definition added to the current module using #module::add */
 bool is_definition(environment const & env, name const & n);
 

src/library/tactic/tactic_state.cpp

@@ -758,21 +758,21 @@ vm_obj tactic_add_doc_string(vm_obj const & n, vm_obj const & doc, vm_obj const
     }
 }
 
-/* meta constant module_doc_strings : tactic (list (option name × string)) */
+/* meta constant module_doc_strings : tactic (list (option string × string)) */
 vm_obj tactic_module_doc_strings(vm_obj const & _s) {
     tactic_state const & s  = tactic::to_state(_s);
-    buffer<doc_entry> entries;
+    buffer<mod_doc_entry> entries;
     get_module_doc_strings(s.env(), entries);
     unsigned i = entries.size();
     vm_obj r   = mk_vm_simple(0);
     while (i > 0) {
         --i;
         vm_obj decl_name;
-        if (auto d = entries[i].get_decl_name())
+        if (auto d = entries[i].m_mod_name)
             decl_name = mk_vm_some(to_obj(*d));
         else
             decl_name = mk_vm_none();
-        vm_obj doc = to_obj(entries[i].get_doc());
+        vm_obj doc = to_obj(entries[i].m_doc);
         vm_obj e   = mk_vm_pair(decl_name, doc);
         r          = mk_vm_constructor(1, e, r);
     }

tests/lean/doc_strings.lean.expected.out

@@ -1,8 +1,5 @@
 Makes foo.
 Makes foo of foo.
 Doesn't make foo.
-[(none, This module has a doc..),
- (none, ..or two.),
- ((some foo.a), Makes foo.),
- ((some foo.b), Makes foo of foo.),
- ((some two), Doesn't make foo.)]
+[(none, This module has a doc..
+..or two.)]