perf: template the kernel's `replace` on its callback This PR speeds up the kernel's expression traversals (`instantiate`, `instantiate_rev`, `abstract`, `lift_loose_bvars`, `lower_loose_bvars`, `instantiate_lparams`). `replace_tpl` is `replace` templated on its callback, so the callback is inlined instead of being called through a `std::function` at every visited node. It also takes a `skip` predicate: subterms the callback would return unchanged are returned before the cache lookup. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
diff --git a/src/kernel/abstract.cpp b/src/kernel/abstract.cpp index 9910fe3..8be1952 100644 --- a/src/kernel/abstract.cpp +++ b/src/kernel/abstract.cpp
@@ -15,9 +15,7 @@ lean_assert(std::all_of(subst, subst+n, [](expr const & e) { return !has_loose_bvars(e) && is_fvar(e); })); if (!has_fvar(e)) return e; - return replace(e, [=](expr const & m, unsigned offset) -> optional<expr> { - if (!has_fvar(m)) - return some_expr(m); // expression m does not contain free variables + return replace_tpl(e, [=](expr const & m, unsigned offset) -> optional<expr> { if (is_fvar(m)) { unsigned i = n; while (i > 0) { @@ -28,7 +26,7 @@ return none_expr(); } return none_expr(); - }); + }, true, [](expr const & m, unsigned) { return !has_fvar(m); }); } static object * lean_expr_abstract_core(object * e0, size_t n, object * subst) { @@ -38,9 +36,7 @@ lean_inc(e0); return e0; } - expr r = replace(e, [=](expr const & m, unsigned offset) -> optional<expr> { - if (!has_fvar(m) && !has_mvar(m)) - return some_expr(m); // expression m does not contain free/meta variables + expr r = replace_tpl(e, [=](expr const & m, unsigned offset) -> optional<expr> { bool fv = is_fvar(m); bool mv = is_mvar(m); if (fv || mv) { @@ -56,7 +52,7 @@ return none_expr(); } return none_expr(); - }); + }, true, [](expr const & m, unsigned) { return !has_fvar(m) && !has_mvar(m); }); return r.steal(); }
diff --git a/src/kernel/expr.cpp b/src/kernel/expr.cpp index b640897..e4c3476 100644 --- a/src/kernel/expr.cpp +++ b/src/kernel/expr.cpp
@@ -374,18 +374,17 @@ if (d == 0 || s >= get_loose_bvar_range(e)) return e; lean_assert(s >= d); - return replace(e, [=](expr const & e, unsigned offset) -> optional<expr> { + return replace_tpl(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_loose_bvar_range(e)) - return some_expr(e); // expression e does not contain bound variables with idx >= s1 if (is_bvar(e) && bvar_idx(e) >= s1) { lean_assert(bvar_idx(e) >= offset + d); return some_expr(mk_bvar(bvar_idx(e) - nat(d))); } else { return none_expr(); } + }, true, [=](expr const & e, unsigned offset) { + unsigned s1 = s + offset; + return s1 < s || s1 >= get_loose_bvar_range(e); }); } @@ -404,17 +403,15 @@ expr lift_loose_bvars(expr const & e, unsigned s, unsigned d) { if (d == 0 || s >= get_loose_bvar_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_loose_bvar_range(e)) - return some_expr(e); // expression e does not contain bound variables with idx >= s1 + return replace_tpl(e, [=](expr const & e, unsigned offset) -> optional<expr> { if (is_var(e) && bvar_idx(e) >= s + offset) { return some_expr(mk_bvar(bvar_idx(e) + nat(d))); } else { return none_expr(); } + }, true, [=](expr const & e, unsigned offset) { + unsigned s1 = s + offset; + return s1 < s || s1 >= get_loose_bvar_range(e); }); }
diff --git a/src/kernel/instantiate.cpp b/src/kernel/instantiate.cpp index 9624cc8..0022983 100644 --- a/src/kernel/instantiate.cpp +++ b/src/kernel/instantiate.cpp
@@ -15,12 +15,8 @@ expr instantiate(expr const & a, unsigned s, unsigned n, expr const * subst) { if (s >= get_loose_bvar_range(a) || n == 0) return a; - return replace(a, [=](expr const & m, unsigned offset) -> optional<expr> { + return replace_tpl(a, [=](expr const & m, unsigned offset) -> optional<expr> { unsigned s1 = s + offset; - if (s1 < s) - return some_expr(m); // overflow, vidx can't be >= max unsigned - if (s1 >= get_loose_bvar_range(m)) - return some_expr(m); // expression m does not contain loose bound variables with idx >= s1 if (is_bvar(m)) { nat const & vidx = bvar_idx(m); if (vidx >= s1) { @@ -33,6 +29,9 @@ } } return none_expr(); + }, true, [=](expr const & m, unsigned offset) { + unsigned s1 = s + offset; + return s1 < s || s1 >= get_loose_bvar_range(m); }); } @@ -57,9 +56,7 @@ lean_inc(a0); return a0; } - expr r = replace(a, [=](expr const & m, unsigned offset) -> optional<expr> { - if (offset >= get_loose_bvar_range(m)) - return some_expr(m); // expression m does not contain loose bound variables with idx >= offset + expr r = replace_tpl(a, [=](expr const & m, unsigned offset) -> optional<expr> { if (is_bvar(m)) { nat const & vidx = bvar_idx(m); if (vidx >= offset) { @@ -73,7 +70,7 @@ } } return none_expr(); - }); + }, true, [](expr const & m, unsigned offset) { return offset >= get_loose_bvar_range(m); }); return r.steal(); } @@ -98,9 +95,7 @@ expr instantiate_rev(expr const & a, unsigned n, expr const * subst) { if (!has_loose_bvars(a)) return a; - return replace(a, [=](expr const & m, unsigned offset) -> optional<expr> { - if (offset >= get_loose_bvar_range(m)) - return some_expr(m); // expression m does not contain loose bound variables with idx >= offset + return replace_tpl(a, [=](expr const & m, unsigned offset) -> optional<expr> { if (is_bvar(m)) { nat const & vidx = bvar_idx(m); if (vidx >= offset) { @@ -113,7 +108,7 @@ } } return none_expr(); - }); + }, true, [](expr const & m, unsigned offset) { return offset >= get_loose_bvar_range(m); }); } static object * lean_expr_instantiate_rev_core(object * a0, size_t n, object ** subst) { @@ -122,9 +117,7 @@ lean_inc(a0); return a0; } - expr r = replace(a, [=](expr const & m, unsigned offset) -> optional<expr> { - if (offset >= get_loose_bvar_range(m)) - return some_expr(m); // expression m does not contain loose bound variables with idx >= offset + expr r = replace_tpl(a, [=](expr const & m, unsigned offset) -> optional<expr> { if (is_bvar(m)) { nat const & vidx = bvar_idx(m); if (vidx >= offset) { @@ -138,7 +131,7 @@ } } return none_expr(); - }); + }, true, [](expr const & m, unsigned offset) { return offset >= get_loose_bvar_range(m); }); return r.steal(); } @@ -213,20 +206,24 @@ } } +// Keeps `map_reuse`'s buffer out of `replace_tpl`'s recursive stack frame. +__attribute__((noinline)) +static expr instantiate_const_lparams(expr const & e, names const & lps, levels const & ls) { + return update_constant(e, map_reuse(const_levels(e), [&](level const & l) { return instantiate(l, lps, ls); })); +} + expr instantiate_lparams(expr const & e, names const & lps, levels const & ls) { if (!has_param_univ(e)) return e; - return replace(e, [&](expr const & e) -> optional<expr> { - if (!has_param_univ(e)) - return some_expr(e); + return replace_tpl(e, [&](expr const & e, unsigned) -> optional<expr> { if (is_constant(e)) { - return some_expr(update_constant(e, map_reuse(const_levels(e), [&](level const & l) { return instantiate(l, lps, ls); }))); + return some_expr(instantiate_const_lparams(e, lps, ls)); } else if (is_sort(e)) { return some_expr(update_sort(e, instantiate(sort_level(e), lps, ls))); } else { return none_expr(); } - }); + }, true, [](expr const & e, unsigned) { return !has_param_univ(e); }); } expr instantiate_type_lparams(constant_info const & info, levels const & ls) {
diff --git a/src/kernel/replace_fn.cpp b/src/kernel/replace_fn.cpp index 8332023..cb1e431 100644 --- a/src/kernel/replace_fn.cpp +++ b/src/kernel/replace_fn.cpp
@@ -13,75 +13,8 @@ namespace lean { -class replace_rec_fn { - struct key_hasher { - std::size_t operator()(std::pair<lean_object *, unsigned> const & p) const { - return hash((size_t)p.first >> 3, p.second); - } - }; - lean::unordered_map<std::pair<lean_object *, unsigned>, expr, key_hasher> m_cache; - std::function<optional<expr>(expr const &, unsigned)> m_f; - bool m_use_cache; - - expr save_result(expr const & e, unsigned offset, expr r, bool shared) { - if (shared) - m_cache.insert(mk_pair(mk_pair(e.raw(), offset), r)); - return r; - } - - expr apply(expr const & e, unsigned offset) { - bool shared = false; - if (m_use_cache && !is_likely_unshared(e)) { - auto it = m_cache.find(mk_pair(e.raw(), offset)); - if (it != m_cache.end()) - return it->second; - shared = true; - } - if (optional<expr> r = m_f(e, offset)) { - return save_result(e, offset, std::move(*r), shared); - } else { - switch (e.kind()) { - case expr_kind::Const: case expr_kind::Sort: - case expr_kind::BVar: case expr_kind::Lit: - case expr_kind::MVar: case expr_kind::FVar: - return save_result(e, offset, e, shared); - case expr_kind::MData: { - expr new_e = apply(mdata_expr(e), offset); - return save_result(e, offset, update_mdata(e, new_e), shared); - } - case expr_kind::Proj: { - expr new_e = apply(proj_expr(e), offset); - return save_result(e, offset, update_proj(e, new_e), shared); - } - case expr_kind::App: { - expr new_f = apply(app_fn(e), offset); - expr new_a = apply(app_arg(e), offset); - return save_result(e, offset, update_app(e, new_f, new_a), shared); - } - case expr_kind::Pi: case expr_kind::Lambda: { - expr new_d = apply(binding_domain(e), offset); - expr new_b = apply(binding_body(e), offset+1); - return save_result(e, offset, update_binding(e, new_d, new_b), shared); - } - case expr_kind::Let: { - expr new_t = apply(let_type(e), offset); - expr new_v = apply(let_value(e), offset); - expr new_b = apply(let_body(e), offset+1); - return save_result(e, offset, update_let(e, new_t, new_v, new_b), shared); - } - } - lean_unreachable(); - } - } -public: - template<typename F> - replace_rec_fn(F const & f, bool use_cache):m_f(f), m_use_cache(use_cache) {} - - expr operator()(expr const & e) { return apply(e, 0); } -}; - expr replace(expr const & e, std::function<optional<expr>(expr const &, unsigned)> const & f, bool use_cache) { - return replace_rec_fn(f, use_cache)(e); + return replace_tpl(e, f, use_cache); } class replace_fn { @@ -204,6 +137,10 @@ } expr apply(expr const & e) { + // cheaper than consulting the cache + if (!has_param_univ(e)) + return e; + bool shared = false; if (is_shared(e)) { auto it = m_cache.find(e.raw()); @@ -212,9 +149,6 @@ shared = true; } - if (!has_param_univ(e)) - return save_result(e, e, shared); - switch (e.kind()) { case expr_kind::Const: return save_result(e, instantiate_const(e), shared);
diff --git a/src/kernel/replace_fn.h b/src/kernel/replace_fn.h index ed2d8f9..fef50ba 100644 --- a/src/kernel/replace_fn.h +++ b/src/kernel/replace_fn.h
@@ -9,8 +9,92 @@ #include "runtime/interrupt.h" #include "kernel/expr.h" #include "kernel/expr_maps.h" +#include "util/alloc.h" namespace lean { +/* The default `skip` of `replace_tpl`: no subterm is skipped. */ +struct replace_skip_none { + bool operator()(expr const &, unsigned) const { return false; } +}; + +template<typename F, typename S> +class replace_rec_fn { + struct key_hasher { + std::size_t operator()(std::pair<lean_object *, unsigned> const & p) const { + return hash((size_t)p.first >> 3, p.second); + } + }; + lean::unordered_map<std::pair<lean_object *, unsigned>, expr, key_hasher> m_cache; + F const & m_f; + S m_skip; + bool m_use_cache; + + expr save_result(expr const & e, unsigned offset, expr r, bool shared) { + if (shared) + m_cache.insert(mk_pair(mk_pair(e.raw(), offset), r)); + return r; + } + + expr apply(expr const & e, unsigned offset) { + if (m_skip(e, offset)) + return e; + bool shared = false; + if (m_use_cache && !is_likely_unshared(e)) { + auto it = m_cache.find(mk_pair(e.raw(), offset)); + if (it != m_cache.end()) + return it->second; + shared = true; + } + if (optional<expr> r = m_f(e, offset)) { + return save_result(e, offset, std::move(*r), shared); + } else { + switch (e.kind()) { + case expr_kind::Const: case expr_kind::Sort: + case expr_kind::BVar: case expr_kind::Lit: + case expr_kind::MVar: case expr_kind::FVar: + return save_result(e, offset, e, shared); + case expr_kind::MData: { + expr new_e = apply(mdata_expr(e), offset); + return save_result(e, offset, update_mdata(e, new_e), shared); + } + case expr_kind::Proj: { + expr new_e = apply(proj_expr(e), offset); + return save_result(e, offset, update_proj(e, new_e), shared); + } + case expr_kind::App: { + expr new_f = apply(app_fn(e), offset); + expr new_a = apply(app_arg(e), offset); + return save_result(e, offset, update_app(e, new_f, new_a), shared); + } + case expr_kind::Pi: case expr_kind::Lambda: { + expr new_d = apply(binding_domain(e), offset); + expr new_b = apply(binding_body(e), offset+1); + return save_result(e, offset, update_binding(e, new_d, new_b), shared); + } + case expr_kind::Let: { + expr new_t = apply(let_type(e), offset); + expr new_v = apply(let_value(e), offset); + expr new_b = apply(let_body(e), offset+1); + return save_result(e, offset, update_let(e, new_t, new_v, new_b), shared); + } + } + lean_unreachable(); + } + } +public: + replace_rec_fn(F const & f, S const & skip, bool use_cache):m_f(f), m_skip(skip), m_use_cache(use_cache) {} + + expr operator()(expr const & e) { return apply(e, 0); } +}; + +/** \brief `replace` for a statically known `f`, avoiding the `std::function` indirection at every visited node. + `skip(m, offset)` may only hold if `f(m, offset)` is `m` itself; such subterms are returned without + consulting or filling the cache, which costs more than `f`'s own early exit. */ +template<typename F, typename S = replace_skip_none> +expr replace_tpl(expr const & e, F const & f, bool use_cache = true, S const & skip = S()) { + return replace_rec_fn<F, S>(f, skip, use_cache)(e); +} + /** \brief Apply <tt>f</tt> to the subexpressions of a given expression.