refactor(frontends/lean/local_context): simplify local_context representation

src/frontends/lean/local_context.cpp

@@ -5,9 +5,36 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Author: Leonardo de Moura
 */
 #include "kernel/abstract.h"
+#include "kernel/replace_fn.h"
 #include "frontends/lean/local_context.h"
 
 namespace lean {
+/** \brief Given a list of local constants \c locals
+              (x_n : A_n) ... (x_0 : A_0)
+    and a term \c e
+              t[x_0, ..., x_n]
+    return
+              t[#n, ..., #0]
+*/
+expr abstract_locals(expr const & e, list<expr> const & locals) {
+    lean_assert(std::all_of(locals.begin(), locals.end(), [](expr const & e) { return closed(e) && is_local(e); }));
+    if (!has_local(e))
+        return e;
+    return replace(e, [=](expr const & m, unsigned offset) -> optional<expr> {
+            if (!has_local(m))
+                return some_expr(m); // expression m does not contain local constants
+            if (is_local(m)) {
+                unsigned i = 0;
+                for (expr const & l : locals) {
+                    if (mlocal_name(l) == mlocal_name(m))
+                        return some_expr(copy_tag(m, mk_var(offset + i)));
+                    i++;
+                }
+            }
+            return none_expr();
+        });
+}
+
 local_context::local_context(name const & prefix):m_ngen(prefix) {}
 local_context::local_context(name const & prefix, list<expr> const & ctx):
     m_ngen(prefix) {
@@ -16,31 +43,32 @@ local_context::local_context(name const & prefix, list<expr> const & ctx):
 
 void local_context::set_ctx(list<expr> const & ctx) {
     m_ctx = ctx;
-    m_ctx_buffer.clear();
-    m_ctx_domain_buffer.clear();
-    to_buffer(ctx, m_ctx_buffer);
-    std::reverse(m_ctx_buffer.begin(), m_ctx_buffer.end());
-    for (unsigned i = 0; i < m_ctx_buffer.size(); i++) {
-        m_ctx_domain_buffer.push_back(abstract_locals(m_ctx_buffer[i], i, m_ctx_buffer.data()));
+    buffer<expr> tmp;
+    list<expr> it = ctx;
+    while (it) {
+        tmp.push_back(abstract_locals(head(it), tail(it)));
+        it = tail(it);
     }
+    m_ctx_abstracted = to_list(tmp.begin(), tmp.end());
+    lean_assert(std::all_of(m_ctx_abstracted.begin(), m_ctx_abstracted.end(), [](expr const & e) { return is_local(e); }));
 }
 
 expr local_context::pi_abstract_context(expr e, tag g) const {
-    e = abstract_locals(e, m_ctx_buffer.size(), m_ctx_buffer.data());
-    unsigned i = m_ctx_domain_buffer.size();
-    while (i > 0) {
-        --i;
-        expr const & l = m_ctx_domain_buffer[i];
+    e = abstract_locals(e, m_ctx);
+    for (expr const & l : m_ctx_abstracted)
         e = mk_pi(local_pp_name(l), mlocal_type(l), e, local_info(l)).set_tag(g);
-    }
     return e;
 }
 
+static expr apply_context_core(expr const & f, list<expr> const & ctx, tag g) {
+    if (ctx)
+        return mk_app(apply_context_core(f, tail(ctx), g), head(ctx)).set_tag(g);
+    else
+        return f;
+}
+
 expr local_context::apply_context(expr const & f, tag g) const {
-    expr r = f;
-    for (unsigned i = 0; i < m_ctx_buffer.size(); i++)
-        r = mk_app(r, m_ctx_buffer[i]).set_tag(g);
-    return r;
+    return apply_context_core(f, m_ctx, g);
 }
 
 expr local_context::mk_type_metavar(tag g) {
@@ -69,9 +97,11 @@ expr local_context::mk_meta(optional<expr> const & type, tag g) {
 }
 
 void local_context::add_local(expr const & l) {
-    m_ctx = cons(l, m_ctx);
-    m_ctx_domain_buffer.push_back(abstract_locals(l, m_ctx_buffer.size(), m_ctx_buffer.data()));
-    m_ctx_buffer.push_back(l);
+    lean_assert(is_local(l));
+    m_ctx_abstracted = cons(abstract_locals(l, m_ctx), m_ctx_abstracted);
+    m_ctx            = cons(l, m_ctx);
+    lean_assert(length(m_ctx) == length(m_ctx_abstracted));
+    lean_assert(is_local(head(m_ctx_abstracted)));
 }
 
 optional<expr> local_context::find_meta(name const & n) const {
@@ -86,10 +116,9 @@ list<expr> const & local_context::get_data() const {
 }
 
 local_context::scope::scope(local_context & main):
-    m_main(main), m_old_ctx(main.m_ctx), m_old_sz(main.m_ctx_buffer.size()) {}
+    m_main(main), m_old_ctx(main.m_ctx), m_old_ctx_abstracted(main.m_ctx_abstracted) {}
 local_context::scope::~scope() {
-    m_main.m_ctx = m_old_ctx;
-    m_main.m_ctx_buffer.shrink(m_old_sz);
-    m_main.m_ctx_domain_buffer.shrink(m_old_sz);
+    m_main.m_ctx            = m_old_ctx;
+    m_main.m_ctx_abstracted = m_old_ctx_abstracted;
 }
 }

src/frontends/lean/local_context.h

@@ -19,8 +19,7 @@ class local_context {
     name_generator  m_ngen;
     mvar2meta       m_mvar2meta;
     list<expr>      m_ctx; // current local context: a list of local constants
-    buffer<expr>    m_ctx_buffer; // m_ctx as a buffer
-    buffer<expr>    m_ctx_domain_buffer; // m_ctx_domain_buffer[i] == abstract_locals(m_ctx_buffer[i], i, m_ctx_buffer.beg
+    list<expr>      m_ctx_abstracted; // m_ctx where elements have been abstracted
 public:
     local_context(name const & prefix);
     local_context(name const & prefix, list<expr> const & ctx);
@@ -94,7 +93,7 @@ public:
     class scope {
         local_context & m_main;
         list<expr>      m_old_ctx;
-        unsigned        m_old_sz;
+        list<expr>      m_old_ctx_abstracted;
     public:
         scope(local_context & main);
         ~scope();