diff --git a/src/library/compiler/borrowed_annotation.cpp b/src/library/compiler/borrowed_annotation.cpp index cb2d4e6d27..ca9307a5b1 100644 --- a/src/library/compiler/borrowed_annotation.cpp +++ b/src/library/compiler/borrowed_annotation.cpp @@ -4,7 +4,9 @@ Released under Apache 2.0 license as described in the file LICENSE. Author: Leonardo de Moura */ +#include "library/module.h" #include "library/annotation.h" +#include "library/compiler/util.h" namespace lean { static name * g_borrowed = nullptr; @@ -13,12 +15,85 @@ expr mk_borrowed(expr const & e) { return mk_annotation(*g_borrowed, e); } bool is_borrowed(expr const & e) { return is_annotation(e, *g_borrowed); } expr const & get_borrowed_arg(expr const & e) { lean_assert(is_borrowed(e)); return get_annotation_arg(e); } +struct borrowed_info { + list m_borrowed_args; + bool m_borrowed_res; +}; + +/* Inferred borrowed annotations cache. */ +struct borrowed_ext : public environment_extension { + typedef name_map cache; + cache m_cache; +}; + +struct borrowed_ext_reg { + unsigned m_ext_id; + borrowed_ext_reg() { m_ext_id = environment::register_extension(std::make_shared()); } +}; + +static borrowed_ext_reg * g_ext = nullptr; +static borrowed_ext const & get_extension(environment const & env) { + return static_cast(env.get_extension(g_ext->m_ext_id)); +} +static environment update(environment const & env, borrowed_ext const & ext) { + return env.update(g_ext->m_ext_id, std::make_shared(ext)); +} + +/* Support for old module manager. + Remark: this code will be deleted in the future */ +struct borrowed_modification : public modification { + LEAN_MODIFICATION("borrowed"); + + name m_fn; + borrowed_info m_info; + + borrowed_modification(name const & fn, borrowed_info const & info):m_fn(fn), m_info(info) {} + + void perform(environment & env) const override { + borrowed_ext ext = get_extension(env); + ext.m_cache.insert(m_fn, m_info); + env = update(env, ext); + } + + void serialize(serializer & s) const override { + s << m_fn << m_info.m_borrowed_res; + write_list(s, m_info.m_borrowed_args); + } + + static std::shared_ptr deserialize(deserializer & d) { + name fn; borrowed_info info; + d >> fn >> info.m_borrowed_res; + info.m_borrowed_args = read_list(d); + return std::make_shared(fn, info); + } +}; + +bool get_inferred_borrowed_info(environment const & env, name const & fn, buffer & borrowed_args, bool & borrowed_res) { + borrowed_ext const & ext = get_extension(env); + if (borrowed_info const * info = ext.m_cache.find(fn)) { + to_buffer(info->m_borrowed_args, borrowed_args); + borrowed_res = info->m_borrowed_res; + return true; + } else { + return false; + } +} + +environment infer_borrowed_annotations(environment const & env, buffer const &) { + // TODO(Leo) + return env; +} + void initialize_borrowed_annotation() { g_borrowed = new name("borrowed"); register_annotation(*g_borrowed); + g_ext = new borrowed_ext_reg(); + borrowed_modification::init(); } void finalize_borrowed_annotation() { delete g_borrowed; + delete g_ext; + borrowed_modification::finalize(); } } diff --git a/src/library/compiler/borrowed_annotation.h b/src/library/compiler/borrowed_annotation.h index 14198fcc04..b5797d4678 100644 --- a/src/library/compiler/borrowed_annotation.h +++ b/src/library/compiler/borrowed_annotation.h @@ -5,11 +5,13 @@ Released under Apache 2.0 license as described in the file LICENSE. Author: Leonardo de Moura */ #pragma once -#include "kernel/expr.h" +#include "library/compiler/util.h" namespace lean { expr mk_borrowed(expr const & e); bool is_borrowed(expr const & e); expr const & get_borrowed_arg(expr const & e); +bool get_inferred_borrowed_info(environment const & env, name const & fn, buffer & borrowed_args, bool & borrowed_res); +environment infer_borrowed_annotations(environment const & env, buffer const &); void initialize_borrowed_annotation(); void finalize_borrowed_annotation(); } diff --git a/src/library/compiler/llnf.cpp b/src/library/compiler/llnf.cpp index db64c2cd87..2064e466e9 100644 --- a/src/library/compiler/llnf.cpp +++ b/src/library/compiler/llnf.cpp @@ -23,6 +23,7 @@ Author: Leonardo de Moura #include "library/compiler/cse.h" #include "library/compiler/elim_dead_let.h" #include "library/compiler/extern_attribute.h" +#include "library/compiler/borrowed_annotation.h" namespace lean { static expr * g_apply = nullptr; @@ -278,6 +279,8 @@ unsigned get_llnf_arity(environment const & env, name const & n) { static void get_borrowed_info(environment const & env, name const & n, buffer & borrowed_args, bool & borrowed_res) { if (get_extern_borrowed_info(env, n, borrowed_args, borrowed_res)) return; /* `n` is an extern/native function declaration. */ + if (get_inferred_borrowed_info(env, n, borrowed_args, borrowed_res)) + return; /* We currently do not support borrowed annotations in user declarations. */ unsigned arity = get_llnf_arity(env, n); borrowed_args.clear(); @@ -2386,14 +2389,15 @@ pair to_llnf(environment const & env, comp_decls const new_v = insert_reset_reuse_fn(new_env, unboxed)(new_v); new_v = simp_cases(new_env, new_v); rs.push_back(comp_decl(d.fst(), new_v)); - if (unboxed) { - if (optional> p = mk_boxed_version(new_env, d.fst(), get_num_nested_lambdas(d.snd()))) { + } + new_env = infer_borrowed_annotations(new_env, rs); + if (unboxed) { + for (comp_decl const & r : rs) { + if (optional> p = mk_boxed_version(new_env, r.fst(), get_num_nested_lambdas(r.snd()))) { new_env = p->first; bs.push_back(p->second); } } - } - if (unboxed) { for (comp_decl & r : rs) { expr new_code = explicit_boxing_fn(new_env)(r.fst(), r.snd()); new_code = ecse(new_env, new_code);