refactor(frontends/lean): use expr_struct_set when collecting locals

Signed-off-by: Leonardo de Moura <leonardo@microsoft.com>

src/frontends/lean/builtin_cmds.cpp

@@ -83,8 +83,9 @@ environment end_scoped_cmd(parser & p) {
 environment check_cmd(parser & p) {
     expr e   = p.parse_expr();
     buffer<expr> section_ps;
-    name_set locals = collect_locals(e);
-    mk_section_params(collect_locals(e), p, section_ps);
+    expr_struct_set locals;
+    collect_locals(e, locals);
+    sort_section_params(locals, p, section_ps);
     e = p.lambda_abstract(section_ps, e);
     level_param_names ls = to_level_param_names(collect_univ_params(e));
     level_param_names new_ls;
@@ -95,7 +96,7 @@ environment check_cmd(parser & p) {
     for (unsigned i = 0; i < section_ps.size(); i++) {
         lean_assert(is_lambda(e));
         lean_assert(is_pi(type));
-        expr local = mk_local(binding_name(e), binding_domain(e), binding_info(e));
+        expr local = section_ps[i];
         e    = instantiate(binding_body(e), local);
         type = instantiate(binding_body(type), local);
     }

src/frontends/lean/inductive_cmd.cpp

@@ -328,16 +328,13 @@ struct inductive_cmd_fn {
     }
 
     /** \brief Collect section local parameters used in the inductive decls */
-    name_set collect_section_locals(buffer<inductive_decl> const & decls) {
-        name_set section_locals;
+    void collect_section_locals(buffer<inductive_decl> const & decls, expr_struct_set & ls) {
         for (auto const & d : decls) {
-            section_locals = collect_locals(inductive_decl_type(d), section_locals);
-            for (auto const & ir : inductive_decl_intros(d)) {
-                section_locals = collect_locals(intro_rule_type(ir), section_locals);
+            collect_locals(inductive_decl_type(d), ls);
+            for (auto const & ir : inductive_decl_intros(d))
+                collect_locals(intro_rule_type(ir), ls);
         }
     }
-        return section_locals;
-    }
 
     /** \brief Make sure that every occurrence of an inductive datatype (in decls) in \c type has
         section parameters \c section_params as arguments.
@@ -361,10 +358,11 @@ struct inductive_cmd_fn {
     void abstract_section_locals(buffer<inductive_decl> & decls, buffer<expr> & section_params) {
         if (!in_section(m_env))
             return;
-        name_set section_locals = collect_section_locals(decls);
+        expr_struct_set section_locals;
+        collect_section_locals(decls, section_locals);
         if (section_locals.empty())
             return;
-        mk_section_params(section_locals, m_p, section_params);
+        sort_section_params(section_locals, m_p, section_params);
         // First, add section_params to inductive types type.
         for (inductive_decl & d : decls) {
             d = update_inductive_decl(d, m_p.pi_abstract(section_params, inductive_decl_type(d)));

src/frontends/lean/util.cpp

@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 
 Author: Leonardo de Moura
 */
+#include <algorithm>
 #include "util/sstream.h"
 #include "library/scoped_ext.h"
 #include "library/locals.h"
@@ -28,10 +29,9 @@ name remove_root_prefix(name const & n) {
 }
 
 // Sort local_names by order of occurrence in the section, and copy the associated parameters to section_ps
-void mk_section_params(name_set const & local_names, parser const & p, buffer<expr> & section_ps) {
-    local_names.for_each([&](name const & n) {
-            section_ps.push_back(*p.get_local(n));
-        });
+void sort_section_params(expr_struct_set const & locals, parser const & p, buffer<expr> & section_ps) {
+    for (expr const & l : locals)
+        section_ps.push_back(l);
     std::sort(section_ps.begin(), section_ps.end(), [&](expr const & p1, expr const & p2) {
             return p.get_local_index(mlocal_name(p1)) < p.get_local_index(mlocal_name(p2));
         });
@@ -51,7 +51,9 @@ levels collect_section_levels(level_param_names const & ls, parser & p) {
 
 // Collect local (section) constants occurring in type and value, sort them, and store in section_ps
 void collect_section_locals(expr const & type, expr const & value, parser const & p, buffer<expr> & section_ps) {
-    name_set ls = collect_locals(type, collect_locals(value));
-    return mk_section_params(ls, p, section_ps);
+    expr_struct_set ls;
+    collect_locals(type, ls);
+    collect_locals(value, ls);
+    sort_section_params(ls, p, section_ps);
 }
 }

src/frontends/lean/util.h

@@ -6,18 +6,17 @@ Author: Leonardo de Moura
 */
 #pragma once
 #include "kernel/expr.h"
+#include "kernel/expr_sets.h"
 namespace lean {
 class parser;
 void check_atomic(name const & n);
 void check_in_section(parser const & p);
 bool is_root_namespace(name const & n);
 name remove_root_prefix(name const & n);
-/** \brief Copy the parameters associated with the local names in \c local_names to \c section_ps.
-    Then sort \c section_ps (using the order in which they were declared).
-*/
-void mk_section_params(name_set const & local_names, parser const & p, buffer<expr> & section_ps);
 /** \brief Return the levels in \c ls that are defined in the section. */
 levels collect_section_levels(level_param_names const & ls, parser & p);
 /** \brief Collect local (section) constants occurring in type and value, sort them, and store in section_ps */
 void collect_section_locals(expr const & type, expr const & value, parser const & p, buffer<expr> & section_ps);
+/** \brief Copy the local parameters to \c section_ps, then sort \c section_ps (using the order in which they were declared). */
+void sort_section_params(expr_struct_set const & locals, parser const & p, buffer<expr> & section_ps);
 }

src/library/locals.cpp

@@ -47,19 +47,15 @@ level_param_names to_level_param_names(name_set const & ls) {
     return r;
 }
 
-name_set collect_locals(expr const & e, name_set const & ls) {
-    if (!has_local(e)) {
-        return ls;
-    } else {
-        name_set r = ls;
+void collect_locals(expr const & e, expr_struct_set & ls) {
+    if (!has_local(e))
+        return;
     for_each(e, [&](expr const & e, unsigned) {
             if (!has_local(e))
                 return false;
             if (is_local(e))
-                    r.insert(mlocal_name(e));
+                ls.insert(e);
             return true;
         });
-        return r;
-    }
 }
 }

src/library/locals.h

@@ -7,8 +7,9 @@ Author: Leonardo de Moura
 #pragma once
 #include "util/name_set.h"
 #include "kernel/expr.h"
+#include "kernel/expr_sets.h"
 namespace lean {
 name_set collect_univ_params(expr const & e, name_set const & ls = name_set());
-name_set collect_locals(expr const & e, name_set const & ls = name_set());
+void collect_locals(expr const & e, expr_struct_set & ls);
 level_param_names to_level_param_names(name_set const & ls);
 }