feat(library/fun_info_manager): safer "replace" function

src/library/fun_info_manager.cpp

@@ -7,6 +7,8 @@ Author: Leonardo de Moura
 #include <algorithm>
 #include "kernel/for_each_fn.h"
 #include "kernel/instantiate.h"
+#include "kernel/abstract.h"
+#include "library/replace_visitor.h"
 #include "library/fun_info_manager.h"
 
 namespace lean {
@@ -18,7 +20,7 @@ auto fun_info_manager::get(expr const & e) -> fun_info {
     if (auto r = m_fun_info.find(e))
         return *r;
     expr type = m_ctx.infer(e);
-    buffer<arg_info> info;
+    buffer<param_info> info;
     buffer<expr> locals;
     while (is_pi(type)) {
         expr local      = m_ctx.mk_tmp_local_from_binding(type);
@@ -51,4 +53,103 @@ auto fun_info_manager::get(expr const & e) -> fun_info {
     m_fun_info.insert(e, r);
     return r;
 }
+
+struct replace_fn : public replace_visitor {
+    fun_info_manager & m_infom;
+    type_context &     m_ctx;
+    unsigned           m_num;
+    expr const *       m_from;
+    expr const *       m_to;
+
+    struct failed {};
+
+    replace_fn(fun_info_manager & infom, unsigned n, expr const * ts, expr const * rs):
+        m_infom(infom), m_ctx(infom.ctx()), m_num(n), m_from(ts), m_to(rs) {}
+
+    virtual expr visit(expr const & e) {
+        for (unsigned i = 0; i < m_num; i++) {
+            if (e == m_from[i])
+                return m_to[i];
+        }
+        return replace_visitor::visit(e);
+    }
+
+    virtual expr visit_mlocal(expr const & e) {
+        return e;
+    }
+
+    virtual expr visit_binding(expr const & b) {
+        expr new_domain = visit(binding_domain(b));
+        expr l          = m_ctx.mk_tmp_local(binding_name(b), new_domain, binding_info(b));
+        expr new_body   = abstract(visit(instantiate(binding_body(b), l)), l);
+        return update_binding(b, new_domain, new_body);
+    }
+
+    virtual expr visit_app(expr const & e) {
+        buffer<expr> args;
+        expr f        = get_app_args(e, args);
+        expr new_f    = visit(f);
+        fun_info info = m_infom.get(f);
+        if (info.get_arity() < args.size())
+            throw failed();
+        auto ps_info  = info.get_params_info();
+        bool modified = f != new_f;
+        buffer<expr> new_args;
+        buffer<bool> to_check;
+        bool has_to_check = false;
+        for (expr const & arg : args) {
+            expr new_arg = visit(arg);
+            auto pinfo   = head(ps_info);
+            bool c       = false;
+            if (modified) {
+                // if not argument has been modified, then there is nothing to be checked
+                for (unsigned idx : pinfo.get_dependencies()) {
+                    lean_assert(idx < new_args.size());
+                    if (args[idx] != new_args[idx]) {
+                        c = true;
+                        has_to_check = true;
+                        break;
+                    }
+                }
+            }
+            if (new_arg != arg) {
+                modified = true;
+            }
+            to_check.push_back(c);
+            new_args.push_back(new_arg);
+            ps_info      = tail(ps_info);
+        }
+        lean_assert(args.size() == new_args.size());
+        if (!modified)
+            return e;
+        expr new_e = mk_app(new_f, new_args);
+        if (has_to_check) {
+            lean_assert(to_check.size() == new_args.size());
+            expr it    = new_e;
+            unsigned i = to_check.size();
+            while (i > 0) {
+                --i;
+                expr const & fn = app_fn(it);
+                if (to_check[i]) {
+                    expr fn_type  = m_ctx.whnf(m_ctx.infer(fn));
+                    if (!is_pi(fn_type))
+                        throw failed();
+                    expr arg_type = m_ctx.infer(app_arg(it));
+                    if (!m_ctx.is_def_eq(binding_domain(fn_type), arg_type))
+                        throw failed();
+                }
+                it = fn;
+            }
+        }
+        return new_e;
+    }
+};
+
+optional<expr> replace(fun_info_manager & infom, expr const & e, unsigned n, expr const * ts, expr const * rs) {
+    try {
+        return some_expr(replace_fn(infom, n, ts, rs)(e));
+    } catch (replace_fn::failed &) {
+        return none_expr();
+    }
+}
 }

src/library/fun_info_manager.h

@@ -9,16 +9,14 @@ Author: Leonardo de Moura
 #include "library/type_context.h"
 
 namespace lean {
-class fun_info_manager {
-public:
-    class arg_info {
+class param_info {
     unsigned       m_implicit:1;
     unsigned       m_inst_implicit:1;
     unsigned       m_prop:1;
     unsigned       m_subsingleton:1;
     list<unsigned> m_deps; // previous arguments it depends on
 public:
-        arg_info(bool imp, bool inst_imp, bool prop, bool sub, list<unsigned> const & deps):
+    param_info(bool imp, bool inst_imp, bool prop, bool sub, list<unsigned> const & deps):
         m_implicit(imp), m_inst_implicit(inst_imp), m_prop(prop), m_subsingleton(sub), m_deps(deps) {}
     list<unsigned> const & get_dependencies() const { return m_deps; }
     bool is_prop() const { return m_prop; }
@@ -29,19 +27,46 @@ public:
 
 class fun_info {
     unsigned         m_arity;
-        list<arg_info> m_args_info;
+    list<param_info> m_params_info;
 public:
     fun_info():m_arity(0) {}
-        fun_info(unsigned arity, list<arg_info> const & info):m_arity(arity), m_args_info(info) {}
+    fun_info(unsigned arity, list<param_info> const & info):m_arity(arity), m_params_info(info) {}
     unsigned get_arity() const { return m_arity; }
-        list<arg_info> const & get_args_info() const { return m_args_info; }
+    list<param_info> const & get_params_info() const { return m_params_info; }
 };
-private:
+
+class fun_info_manager {
     type_context &                         m_ctx;
     rb_map<expr, fun_info, expr_quick_cmp> m_fun_info;
 
 public:
     fun_info_manager(type_context & ctx);
+    type_context & ctx() { return m_ctx; }
     fun_info get(expr const & e);
 };
+
+
+/** \brief Given a term \c e of the form F[ts[0], ..., ts[n-1]],
+    return F[rs[0], ..., rs[n-1]] only if the replacement does not produce type errors.
+
+    Note that even if each ts[i] and rs[i] have the same type, the result may be none.
+
+    Here is a very simple example:
+    Given (f : Pi (n : nat), bv n -> bv n) (v : bv 10), then (F[v] := f 10 v) is type
+    correct, but (f 11 v) is not.
+
+    If
+      a) F[ts[0], ..., ts[n-1] is type correct, and
+      b) If there is a telescope T s.t. (ts[0], ..., ts[n-1]) and (rs[0], ..., rs[n-1]) are instances of T, and
+      c) the result is not none
+    Then
+      the result is type correct.
+
+    TODO(Leo): move to a different file?
+*/
+optional<expr> replace(fun_info_manager & infom, expr const & e, unsigned n, expr const * ts, expr const * rs);
+inline optional<expr> replace(fun_info_manager & infom, expr const & e, buffer<expr> const & ts, buffer<expr> const & rs) {
+    lean_assert(ts.size() == rs.size());
+    return replace(infom, e, ts.size(), ts.data(), rs.data());
+}
 }