refactor(kernel/free_vars): use get_free_var_range to improve lift_free_vars and lower_free_vars performance

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

src/kernel/free_vars.cpp

@@ -99,12 +99,17 @@ bool has_free_var(expr const & e, unsigned low, unsigned high) {
 bool has_free_var(expr const & e, unsigned i) { return has_free_var(e, i, i+1); }
 
 expr lower_free_vars(expr const & e, unsigned s, unsigned d) {
-    if (d == 0 || closed(e))
+    if (d == 0 || s >= get_free_var_range(e))
         return e;
     lean_assert(s >= d);
     lean_assert(!has_free_var(e, s-d, s));
     return replace(e, [=](expr const & e, unsigned offset) -> optional<expr> {
-            if (is_var(e) && var_idx(e) >= s + offset) {
+            unsigned s1 = s + offset;
+            if (s1 < s)
+                return some_expr(e); // overflow, vidx can't be >= max unsigned
+            if (s1 >= get_free_var_range(e))
+                return some_expr(e); // expression e does not contain free variables with idx >= s1
+            if (is_var(e) && var_idx(e) >= s1) {
                 lean_assert(var_idx(e) >= offset + d);
                 return some_expr(mk_var(var_idx(e) - d));
             } else {
@@ -115,11 +120,19 @@ expr lower_free_vars(expr const & e, unsigned s, unsigned d) {
 expr lower_free_vars(expr const & e, unsigned d) { return lower_free_vars(e, d, d); }
 
 expr lift_free_vars(expr const & e, unsigned s, unsigned d) {
-    if (d == 0 || closed(e))
+    if (d == 0 || s >= get_free_var_range(e))
         return e;
     return replace(e, [=](expr const & e, unsigned offset) -> optional<expr> {
+            unsigned s1 = s + offset;
+            if (s1 < s)
+                return some_expr(e); // overflow, vidx can't be >= max unsigned
+            if (s1 >= get_free_var_range(e))
+                return some_expr(e); // expression e does not contain free variables with idx >= s1
             if (is_var(e) && var_idx(e) >= s + offset) {
-                return some_expr(mk_var(var_idx(e) + d));
+                unsigned new_idx = var_idx(e) + d;
+                if (new_idx < var_idx(e))
+                    throw exception("invalid lift_free_vars operation, index overflow");
+                return some_expr(mk_var(new_idx));
             } else {
                 return none_expr();
             }

src/kernel/replace_fn.h

@@ -25,11 +25,12 @@ public:
    \brief Functional for applying <tt>F</tt> to the subexpressions of a given expression.
 
    The signature of \c F is
-   expr const &, unsigned -> expr
+   expr const &, unsigned -> optional(expr)
 
    F is invoked for each subexpression \c s of the input expression e.
    In a call <tt>F(s, n)</tt>, n is the scope level, i.e., the number of
-   bindings operators that enclosing \c s.
+   bindings operators that enclosing \c s. The replaces only visits children of \c e
+   if F return none_expr
 
    P is a "post-processing" functional object that is applied to each
    pair (old, new)

src/tests/kernel/free_vars.cpp

@@ -57,10 +57,24 @@ static void tst3() {
     }
 }
 
+static void tst4() {
+    expr f = Const("f");
+    expr x = Const("x");
+    expr y = Const("y");
+    expr B = Bool;
+    expr t = f(Fun({x, B}, Fun({y, B}, f(x, y)))(f(Var(1), Var(2))), x);
+    lean_assert_eq(lift_free_vars(t, 1, 2),
+                   f(Fun({x, B}, Fun({y, B}, f(x, y)))(f(Var(3), Var(4))), x));
+    lean_assert_eq(lift_free_vars(t, 0, 3),
+                   f(Fun({x, B}, Fun({y, B}, f(x, y)))(f(Var(4), Var(5))), x));
+    lean_assert_eq(lift_free_vars(t, 3, 2), t);
+}
+
 int main() {
     save_stack_info();
     tst1();
     tst2();
     tst3();
+    tst4();
     return has_violations() ? 1 : 0;
 }