320 lines
12 KiB
C++
320 lines
12 KiB
C++
/*
|
|
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
|
|
Released under Apache 2.0 license as described in the file LICENSE.
|
|
|
|
Author: Leonardo de Moura
|
|
*/
|
|
#include "kernel/instantiate.h"
|
|
#include "library/module.h"
|
|
#include "library/trace.h"
|
|
#include "library/annotation.h"
|
|
#include "library/compiler/util.h"
|
|
#include "library/compiler/llnf.h"
|
|
#include "library/compiler/extern_attribute.h"
|
|
#include "library/compiler/export_attribute.h"
|
|
|
|
namespace lean {
|
|
static name * g_borrowed = nullptr;
|
|
|
|
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<bool> m_borrowed_args;
|
|
bool m_borrowed_res;
|
|
};
|
|
|
|
/* Inferred borrowed annotations cache. */
|
|
struct borrowed_ext : public environment_extension {
|
|
typedef name_map<borrowed_info> cache;
|
|
cache m_cache;
|
|
};
|
|
|
|
struct borrowed_ext_reg {
|
|
unsigned m_ext_id;
|
|
borrowed_ext_reg() { m_ext_id = environment::register_extension(std::make_shared<borrowed_ext>()); }
|
|
};
|
|
|
|
static borrowed_ext_reg * g_ext = nullptr;
|
|
static borrowed_ext const & get_extension(environment const & env) {
|
|
return static_cast<borrowed_ext const &>(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<borrowed_ext>(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<modification const> deserialize(deserializer & d) {
|
|
name fn; borrowed_info info;
|
|
d >> fn >> info.m_borrowed_res;
|
|
info.m_borrowed_args = read_list<bool>(d);
|
|
return std::make_shared<borrowed_modification>(fn, info);
|
|
}
|
|
};
|
|
|
|
bool get_inferred_borrowed_info(environment const & env, name const & fn, buffer<bool> & 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;
|
|
}
|
|
}
|
|
|
|
class borrow_inference_fn {
|
|
environment m_env;
|
|
name_generator m_ngen;
|
|
buffer<comp_decl> const & m_decls;
|
|
buffer<buffer<bool>> m_borrowed_args;
|
|
bool m_changed;
|
|
unsigned m_curr_fidx;
|
|
name_set m_owned_vars;
|
|
local_ctx m_lctx;
|
|
|
|
name_generator & ngen() { return m_ngen; }
|
|
|
|
void get_borrowed_info(name const & fn, buffer<bool> & borrowed_args, bool & borrowed_res) {
|
|
if (get_extern_borrowed_info(m_env, fn, borrowed_args, borrowed_res))
|
|
return;
|
|
if (get_inferred_borrowed_info(m_env, fn, borrowed_args, borrowed_res))
|
|
return;
|
|
for (unsigned i = 0; i < m_decls.size(); i++) {
|
|
if (m_decls[i].fst() == fn) {
|
|
borrowed_args.clear();
|
|
borrowed_args.append(m_borrowed_args[i]);
|
|
borrowed_res = false;
|
|
return;
|
|
}
|
|
}
|
|
unsigned arity = get_llnf_arity(m_env, fn);
|
|
borrowed_args.clear();
|
|
borrowed_args.resize(arity, false);
|
|
borrowed_res = false;
|
|
}
|
|
|
|
bool has_borrowed_result(name const & fn) {
|
|
buffer<bool> borrowed_args;
|
|
bool borrowed_res;
|
|
get_borrowed_info(fn, borrowed_args, borrowed_res);
|
|
return borrowed_res;
|
|
}
|
|
|
|
static bool is_proj(expr const & e) { return is_llnf_proj(get_app_fn(e)); }
|
|
|
|
bool is_owned(expr const & e) const {
|
|
return is_fvar(e) && m_owned_vars.contains(fvar_name(e));
|
|
}
|
|
|
|
void mark_owned(expr const & e) {
|
|
if (is_fvar(e))
|
|
m_owned_vars.insert(fvar_name(e));
|
|
}
|
|
|
|
void mark_owned(buffer<expr> const & es) {
|
|
for (expr const & e : es) mark_owned(e);
|
|
}
|
|
|
|
void collect_app(expr const & e, bool terminal) {
|
|
buffer<expr> args;
|
|
expr const & fn = get_app_args(e, args);
|
|
if (is_cases_on_app(m_env, e)) {
|
|
for (unsigned i = 1; i < args.size(); i++)
|
|
collect(args[i]);
|
|
} else if (is_llnf_apply(fn) || is_llnf_closure(fn) || is_llnf_reset(fn) || is_llnf_cnstr(fn) || is_llnf_reuse(fn)) {
|
|
mark_owned(args);
|
|
} else if (!is_llnf_op(fn)) {
|
|
bool tail_call = terminal && (m_decls[m_curr_fidx].fst() == const_name(fn));
|
|
buffer<bool> borrowed_args;
|
|
bool borrowed_res;
|
|
get_borrowed_info(const_name(fn), borrowed_args, borrowed_res);
|
|
for (unsigned i = 0; i < args.size(); i++) {
|
|
if (borrowed_args[i]) {
|
|
if (tail_call && is_owned(args[i])) {
|
|
/* We want to preserve tail calls.
|
|
Recall that given a call `f a`, if `a` is owned, but `f`'s argument is borrowed,
|
|
we must perform `dec a` after `f a`. */
|
|
lean_assert(const_name(fn) == m_decls[m_curr_fidx].fst());
|
|
lean_assert(borrowed_args[i] == m_borrowed_args[m_curr_fidx][i]);
|
|
lean_assert(m_borrowed_args[m_curr_fidx][i]);
|
|
m_borrowed_args[m_curr_fidx][i] = false;
|
|
m_changed = true;
|
|
}
|
|
} else {
|
|
mark_owned(args[i]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
expr get_lambda_body(expr e, buffer<expr> & fvars) {
|
|
while (is_lambda(e)) {
|
|
lean_assert(!has_loose_bvars(binding_domain(e)));
|
|
expr type = binding_domain(e);
|
|
expr new_fvar = m_lctx.mk_local_decl(ngen(), binding_name(e), type);
|
|
fvars.push_back(new_fvar);
|
|
e = binding_body(e);
|
|
}
|
|
return instantiate_rev(e, fvars.size(), fvars.data());
|
|
}
|
|
|
|
void collect_lambda(expr const & e) {
|
|
/* Visiting join point */
|
|
buffer<expr> fvars;
|
|
expr body = get_lambda_body(e, fvars);
|
|
for (expr const & x : fvars) {
|
|
/* All join-point arguments are considered to be owned. */
|
|
mark_owned(x);
|
|
}
|
|
return collect(body);
|
|
}
|
|
|
|
void collect_let(expr e) {
|
|
buffer<expr> fvars;
|
|
while (is_let(e)) {
|
|
lean_assert(!has_loose_bvars(let_type(e)));
|
|
expr type = let_type(e);
|
|
expr val = instantiate_rev(let_value(e), fvars.size(), fvars.data());
|
|
/* if `e` is of ther form `let x := val in x`, process `val` as terminal. */
|
|
bool terminal = is_bvar(let_body(e), 0);
|
|
collect(val, terminal);
|
|
expr new_fvar = m_lctx.mk_local_decl(ngen(), let_name(e), type, val);
|
|
if (is_app(val)) {
|
|
expr const & fn = get_app_fn(val);
|
|
if (is_llnf_apply(fn) || is_llnf_closure(fn) || is_llnf_cnstr(fn) || is_llnf_reuse(fn) || is_llnf_reset(fn)) {
|
|
mark_owned(new_fvar);
|
|
} else if (is_llnf_proj(fn)) {
|
|
if (is_owned(app_arg(val)))
|
|
mark_owned(new_fvar);
|
|
} else if (!is_llnf_op(fn) && !has_borrowed_result(const_name(fn))) {
|
|
mark_owned(new_fvar);
|
|
}
|
|
}
|
|
fvars.push_back(new_fvar);
|
|
e = let_body(e);
|
|
}
|
|
collect(instantiate_rev(e, fvars.size(), fvars.data()));
|
|
unsigned i = fvars.size();
|
|
while (i > 0) {
|
|
--i;
|
|
expr z = fvars[i];
|
|
expr z_val = *m_lctx.get_local_decl(z).get_value();
|
|
if (is_owned(z) && is_proj(z_val)) {
|
|
mark_owned(app_arg(z_val));
|
|
}
|
|
}
|
|
}
|
|
|
|
void collect(expr const & e, bool terminal = true) {
|
|
switch (e.kind()) {
|
|
case expr_kind::App: collect_app(e, terminal); return;
|
|
case expr_kind::Let: collect_let(e); return;
|
|
case expr_kind::Lambda: collect_lambda(e); return;
|
|
default: return;
|
|
}
|
|
}
|
|
|
|
void init_borrowed_args() {
|
|
m_borrowed_args.resize(m_decls.size());
|
|
for (unsigned i = 0; i < m_decls.size(); i++) {
|
|
bool is_exported = has_export_name(m_env, m_decls[i].fst());
|
|
/* For now, we assume that all exported functions all arguments are *not* borrowed */
|
|
m_borrowed_args[i].resize(get_num_nested_lambdas(m_decls[i].snd()), !is_exported);
|
|
}
|
|
}
|
|
|
|
public:
|
|
borrow_inference_fn(environment const & env, buffer<comp_decl> const & decls):
|
|
m_env(env), m_decls(decls) {
|
|
}
|
|
|
|
environment operator()() {
|
|
init_borrowed_args();
|
|
m_changed = true;
|
|
while (m_changed) {
|
|
m_changed = false;
|
|
for (unsigned fidx = 0; fidx < m_decls.size(); fidx++) {
|
|
comp_decl const & d = m_decls[fidx];
|
|
m_curr_fidx = fidx;
|
|
m_owned_vars.clear();
|
|
m_lctx = local_ctx();
|
|
expr const & code = d.snd();
|
|
if (is_lambda(code)) {
|
|
buffer<expr> args;
|
|
expr body = get_lambda_body(code, args);
|
|
for (unsigned i = 0; i < m_borrowed_args[fidx].size(); i++) {
|
|
if (!m_borrowed_args[fidx][i])
|
|
mark_owned(args[i]);
|
|
}
|
|
collect(body);
|
|
for (unsigned i = 0; i < args.size(); i++) {
|
|
expr arg_type = m_lctx.get_local_decl(args[i]).get_type();
|
|
if (!is_llnf_unboxed_type(arg_type) && is_owned(args[i]) && m_borrowed_args[fidx][i]) {
|
|
m_changed = true;
|
|
m_borrowed_args[fidx][i] = false;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
lean_trace(name({"compiler", "borrowed_inference"}),
|
|
for (unsigned fidx = 0; fidx < m_decls.size(); fidx++) {
|
|
comp_decl const & d = m_decls[fidx];
|
|
tout() << d.fst();
|
|
for (unsigned i = 0; i < m_borrowed_args[fidx].size(); i++) {
|
|
tout() << " " << m_borrowed_args[fidx][i];
|
|
}
|
|
tout() << "\n";
|
|
});
|
|
borrowed_ext ext = get_extension(m_env);
|
|
environment new_env = m_env;
|
|
for (unsigned fidx = 0; fidx < m_decls.size(); fidx++) {
|
|
name fn = m_decls[fidx].fst();
|
|
borrowed_info info{to_list(m_borrowed_args[fidx]), false};
|
|
new_env = module::add(new_env, std::make_shared<borrowed_modification>(fn, info));
|
|
ext.m_cache.insert(fn, info);
|
|
}
|
|
return update(new_env, ext);
|
|
}
|
|
};
|
|
|
|
environment infer_borrowed_annotations(environment const & env, buffer<comp_decl> const & ds) {
|
|
return borrow_inference_fn(env, ds)();
|
|
}
|
|
|
|
void initialize_borrowed_annotation() {
|
|
g_borrowed = new name("borrowed");
|
|
register_annotation(*g_borrowed);
|
|
g_ext = new borrowed_ext_reg();
|
|
borrowed_modification::init();
|
|
register_trace_class({"compiler", "borrowed_inference"});
|
|
}
|
|
|
|
void finalize_borrowed_annotation() {
|
|
delete g_borrowed;
|
|
delete g_ext;
|
|
borrowed_modification::finalize();
|
|
}
|
|
}
|