feat(library/simplifier): add API for extracting simplification rules defined in a given namespace

src/frontends/lean/builtin_cmds.cpp

@@ -381,11 +381,22 @@ static void print_reducible_info(parser & p, reducible_status s1) {
 
 static void print_simp_sets(parser & p) {
     io_state_stream out = p.regular_stream();
-    simp_rule_sets s = get_simp_rule_sets(p.env());
+    simp_rule_sets s;
+    name ns;
+    if (p.curr_is_identifier()) {
+        ns = p.get_name_val();
+        p.next();
+        s = get_simp_rule_sets(p.env(), ns);
+    } else {
+        s = get_simp_rule_sets(p.env());
+    }
     name prev_eqv;
     s.for_each([&](name const & eqv, simp_rule const & rw) {
             if (prev_eqv != eqv) {
-                out << "simplification rules for " << eqv << "\n";
+                out << "simplification rules for " << eqv;
+                if (!ns.is_anonymous())
+                    out << " at namespace '" << ns << "'";
+                out << "\n";
                 prev_eqv = eqv;
             }
             out << rw.pp(out.get_formatter()) << "\n";

src/library/simplifier/simp_rule_set.cpp

@@ -100,7 +100,7 @@ void simp_rule_sets::for_each(std::function<void(name const &, simp_rule const &
 static name * g_prefix = nullptr;
 
 simp_rule_sets add_core(type_checker & tc, simp_rule_sets const & s,
-                           name const & id, levels const & univ_metas, expr const & e, expr const & h) {
+                        name const & id, levels const & univ_metas, expr const & e, expr const & h) {
     list<expr_pair> ceqvs   = to_ceqvs(tc, e, h);
     environment const & env = tc.env();
     simp_rule_sets new_s = s;
@@ -139,25 +139,27 @@ simp_rule_sets join(simp_rule_sets const & s1, simp_rule_sets const & s2) {
 static name * g_class_name = nullptr;
 static std::string * g_key = nullptr;
 
+static simp_rule_sets add_core(type_checker & tc, simp_rule_sets const & s, name const & cname) {
+    declaration const & d = tc.env().get(cname);
+    buffer<level> us;
+    unsigned num_univs = d.get_num_univ_params();
+    for (unsigned i = 0; i < num_univs; i++) {
+        us.push_back(mk_meta_univ(name(*g_prefix, i)));
+    }
+    levels ls = to_list(us);
+    expr e    = instantiate_type_univ_params(d, ls);
+    expr h    = mk_constant(cname, ls);
+    return add_core(tc, s, cname, ls, e, h);
+}
+
 struct rrs_state {
     simp_rule_sets           m_sets;
-    name_set                    m_rnames;
-    name_map<simp_rule_sets> m_namespace_cache;
+    name_set                 m_snames;
 
     void add(environment const & env, name const & cname) {
         type_checker tc(env);
-        declaration const & d = env.get(cname);
-        buffer<level> us;
-        unsigned num_univs = d.get_num_univ_params();
-        for (unsigned i = 0; i < num_univs; i++) {
-            us.push_back(mk_meta_univ(name(*g_prefix, i)));
-        }
-        levels ls = to_list(us);
-        expr e    = instantiate_type_univ_params(d, ls);
-        expr h    = mk_constant(cname, ls);
-        m_sets = add_core(tc, m_sets, cname, ls, e, h);
-        m_rnames.insert(cname);
-        m_namespace_cache.erase(get_namespace(env));
+        m_sets = add_core(tc, m_sets, cname);
+        m_snames.insert(cname);
     }
 };
 
@@ -197,13 +199,25 @@ environment add_congr_rule(environment const & env, name const & n, bool persist
 }
 
 bool is_simp_rule(environment const & env, name const & n) {
-    return rrs_ext::get_state(env).m_rnames.contains(n);
+    return rrs_ext::get_state(env).m_snames.contains(n);
 }
 
 simp_rule_sets get_simp_rule_sets(environment const & env) {
     return rrs_ext::get_state(env).m_sets;
 }
 
+simp_rule_sets get_simp_rule_sets(environment const & env, name const & ns) {
+    simp_rule_sets set;
+    list<name> const * cnames = rrs_ext::get_entries(env, ns);
+    if (cnames) {
+        type_checker tc(env);
+        for (name const & cname : *cnames) {
+            set = add_core(tc, set, cname);
+        }
+    }
+    return set;
+}
+
 void initialize_simp_rule_set() {
     g_prefix     = new name(name::mk_internal_unique_name());
     g_class_name = new name("rrs");

src/library/simplifier/simp_rule_set.h

@@ -76,7 +76,7 @@ bool is_simp_rule(environment const & env, name const & n);
 /** \brief Get current rewrite rule sets */
 simp_rule_sets get_simp_rule_sets(environment const & env);
 /** \brief Get rewrite rule sets in the given namespace. */
-simp_rule_sets get_simp_rule_set(environment const & env, name const & ns);
+simp_rule_sets get_simp_rule_sets(environment const & env, name const & ns);
 void initialize_simp_rule_set();
 void finalize_simp_rule_set();
 }

tests/lean/run/rw_set2.lean

@@ -0,0 +1,4 @@
+import data.nat
+
+
+print [simp] nat