feat(library): add function free_var_range for computing the range [0, R) of free variables occurring in an expression

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

src/kernel/free_vars.cpp

@@ -106,6 +106,125 @@ bool has_free_vars(expr const & e) {
     return has_free_vars_fn()(e);
 }
 
+/**
+   \brief Functional object for computing the range [0, R) of free variables occurring
+   in an expression.
+*/
+class free_var_range_fn {
+    expr_map<unsigned>  m_cached;
+    metavar_env const & m_menv;
+
+    static unsigned dec(unsigned s) { return (s == 0) ? 0 : s - 1; }
+
+    /*
+      \brief If a metavariable \c m was defined in a context \c ctx and <tt>ctx.size() == R</tt>,
+      then \c m can only contain free variables in the range <tt>[0, R)</tt>
+
+      So, if \c m does not have an associated local context, the answer is just \c R.
+      If \c m has an associated local context, we process it using the following rules
+
+           [inst:s v] R  ===>  if s >= R then R else max(R-1, range_of(v))
+           [lift:s:n] R  ===>  if s >= R then R else R + n
+    */
+    unsigned process_metavar(expr const & m) {
+        lean_assert(is_metavar(m));
+        context ctx = m_menv.get_context(metavar_name(m));
+        unsigned R  = ctx.size();
+        if (has_local_context(m)) {
+            local_context lctx = metavar_lctx(m);
+            buffer<local_entry> lentries;
+            to_buffer(lctx, lentries);
+            unsigned i = lentries.size();
+            while (i > 0) {
+                --i;
+                local_entry const & entry = lentries[i];
+                if (entry.is_inst()) {
+                    if (entry.s() < R) {
+                        R = std::max(dec(R), apply(entry.v()));
+                    }
+                } else {
+                    if (entry.s() < R)
+                        R += entry.n();
+                }
+            }
+        }
+        return R;
+    }
+
+    unsigned apply(optional<expr> const & e) {
+        return e ? apply(*e) : 0;
+    }
+
+    unsigned apply(expr const & e) {
+        // handle easy cases
+        switch (e.kind()) {
+        case expr_kind::Constant:
+            if (!const_type(e))
+                return 0;
+            break;
+        case expr_kind::Type: case expr_kind::Value:
+            return 0;
+        case expr_kind::Var:
+            return var_idx(e) + 1;
+        case expr_kind::MetaVar: case expr_kind::App: case expr_kind::Eq:
+        case expr_kind::Lambda: case expr_kind::Pi: case expr_kind::Let:
+            break;
+        }
+
+        if (e.raw()->is_closed())
+            return 0;
+
+        bool shared = false;
+        if (is_shared(e)) {
+            shared = true;
+            auto it = m_cached.find(e);
+            if (it != m_cached.end())
+                return it->second;
+        }
+
+        unsigned result = 0;
+
+        switch (e.kind()) {
+        case expr_kind::Constant:
+            lean_assert(const_type(e));
+            result = apply(const_type(e));
+            break;
+        case expr_kind::Type: case expr_kind::Value: case expr_kind::Var:
+            // easy cases were already handled
+            lean_unreachable(); // LCOV_EXCL_LINE
+        case expr_kind::MetaVar:
+            result = process_metavar(e);
+            break;
+        case expr_kind::App:
+            for (auto const & c : args(e))
+                result = std::max(result, apply(c));
+            break;
+        case expr_kind::Eq:
+            result = std::max(apply(eq_lhs(e)), apply(eq_rhs(e)));
+            break;
+        case expr_kind::Lambda:
+        case expr_kind::Pi:
+            result = std::max(apply(abst_domain(e)), dec(apply(abst_body(e))));
+            break;
+        case expr_kind::Let:
+            result = std::max({apply(let_type(e)), apply(let_value(e)), dec(apply(let_body(e)))});
+            break;
+        }
+
+        if (shared)
+            m_cached.insert(mk_pair(e, result));
+
+        return result;
+    }
+public:
+    free_var_range_fn(metavar_env const & menv):m_menv(menv) {}
+    unsigned operator()(expr const & e) { return apply(e); }
+};
+
+unsigned free_var_range(expr const & e, metavar_env const & menv) {
+    return free_var_range_fn(menv)(e);
+}
+
 /**
     \brief Functional object for checking whether a kernel expression has a free variable in the range <tt>[low, high)</tt> or not.
 

src/kernel/free_vars.h

@@ -16,6 +16,26 @@ bool has_free_vars(expr const & a);
 */
 inline bool closed(expr const & a) { return !has_free_vars(a); }
 
+class metavar_env;
+/**
+   \brief Return \c R s.t. the de Bruijn index of all free variables
+   occurring in \c e is in the interval <tt>[0, R)</tt>.
+
+   \pre All metavariables occurring in \c e must have been created
+   at \c menv.
+
+   \remark Regarding metavariables, if a metavariable \c m was defined
+   in a context \c ctx and <tt>ctx.size() == R</tt>, then \c m can
+   only contain free variables in the range <tt>[0, R)</tt>
+
+   So, if \c m does not have an associated local context, the answer is just \c R.
+   If \c m has an associated local context, we process it using the following rules
+
+   [inst:s v] R  ===>  if s >= R then R else max(R-1, range_of(v))
+   [lift:s:n] R  ===>  if s >= R then R else R + n
+*/
+unsigned free_var_range(expr const & e, metavar_env const & menv);
+
 /**
     \brief Return true iff \c e contains the free variable <tt>(var i)</tt>.
 */

src/tests/kernel/free_vars.cpp

@@ -7,6 +7,8 @@ Author: Leonardo de Moura
 #include "util/test.h"
 #include "kernel/free_vars.h"
 #include "kernel/abstract.h"
+#include "kernel/metavar.h"
+#include "kernel/builtin.h"
 using namespace lean;
 
 static void tst1() {
@@ -57,10 +59,58 @@ static void tst3() {
     }
 }
 
+static void tst4() {
+    metavar_env menv;
+    auto fn = [&](expr const & e) { return free_var_range(e, menv); };
+    expr f = Const("f");
+    expr a = Const("a");
+    expr b = Const("b");
+    expr x = Const("x");
+    expr m1 = menv.mk_metavar();
+    lean_assert(fn(m1) == 0);
+    lean_assert(fn(Var(0)) == 1);
+    lean_assert(fn(Var(0)(Var(2), Var(1))) == 3);
+    lean_assert(fn(Type()) == 0);
+    lean_assert(fn(Bool) == 0);
+    lean_assert(fn(Fun({x, Type()}, Var(0))) == 0);
+    lean_assert(fn(Fun({x, Var(0)}, Var(0))) == 1);
+    lean_assert(fn(Fun({x, Var(0)}, Var(2))) == 2);
+    lean_assert(fn(Fun({x, Var(0)}, Eq(Var(2), Var(1)))) == 2);
+    lean_assert(fn(Pi({x, Type()}, Var(0))) == 0);
+    lean_assert(fn(Pi({x, Var(0)}, Var(0))) == 1);
+    lean_assert(fn(Pi({x, Var(0)}, Var(2))) == 2);
+    lean_assert(fn(Pi({x, Var(0)}, Eq(Var(2), Var(1)))) == 2);
+    context ctx;
+    ctx = extend(ctx, name("x"), Bool);
+    ctx = extend(ctx, name("y"), Bool);
+    expr m2 = menv.mk_metavar(ctx);
+    lean_assert_eq(fn(m2), 2);
+    lean_assert_eq(fn(add_lift(m2, 3, 5)), 2);
+    lean_assert_eq(fn(add_lift(m2, 2, 5)), 2);
+    lean_assert_eq(fn(add_lift(m2, 1, 5)), 7);
+    lean_assert_eq(fn(add_inst(m2, 3, Var(10))), 2);
+    lean_assert_eq(fn(add_inst(m2, 1, Var(10))), 11);
+    // Here is the explanation for the following assertion.
+    // If m2 is assigned to #1, that m2[1:f(#2)] becomes f(#2),
+    // and then lift:2:2 transforms it to f(#4)
+    lean_assert_eq(fn(add_lift(add_inst(m2, 1, f(Var(2))), 2, 2)), 5);
+    ctx = extend(ctx, name("w"), Bool);
+    ctx = extend(ctx, name("z"), Bool);
+    expr m3 = menv.mk_metavar(ctx);
+    lean_assert_eq(fn(m3), 4);
+    lean_assert_eq(fn(add_lift(add_inst(m3, 1, f(Var(0))), 1, 1)), 4);
+    lean_assert_eq(fn(add_lift(add_inst(m3, 1, f(Var(3))), 1, 1)), 5);
+    lean_assert_eq(fn(mk_let("x", Var(0), Var(1))), 1);
+    lean_assert_eq(fn(mk_let("x", Var(1), Var(1))), 2);
+    lean_assert_eq(fn(mk_let("x", Var(2), Var(1), Var(1))), 3);
+    lean_assert_eq(fn(mk_let("x", Var(2), Var(1), Var(4))), 4);
+}
+
 int main() {
     save_stack_info();
     tst1();
     tst2();
     tst3();
+    tst4();
     return has_violations() ? 1 : 0;
 }