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

@@ -139,14 +139,8 @@ 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;
 
-struct rrs_state {
-    simp_rule_sets           m_sets;
-    name_set                    m_rnames;
-    name_map<simp_rule_sets> m_namespace_cache;
-
-    void add(environment const & env, name const & cname) {
-        type_checker tc(env);
-        declaration const & d = env.get(cname);
+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++) {
@@ -155,9 +149,17 @@ struct rrs_state {
     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));
+    return add_core(tc, s, cname, ls, e, h);
+}
+
+struct rrs_state {
+    simp_rule_sets           m_sets;
+    name_set                 m_snames;
+
+    void add(environment const & env, name const & cname) {
+        type_checker tc(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