feat(library/compiler): add borrowed annotation inference skeleton

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<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;
+    }
+}
+
+environment infer_borrowed_annotations(environment const & env, buffer<comp_decl> 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();
 }
 }

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<bool> & borrowed_args, bool & borrowed_res);
+environment infer_borrowed_annotations(environment const & env, buffer<comp_decl> const &);
 void initialize_borrowed_annotation();
 void finalize_borrowed_annotation();
 }

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<bool> & 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<environment, comp_decls> 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<pair<environment, comp_decl>> 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<pair<environment, comp_decl>> 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);