Vm override (#48)

* Introduce VM_override

* fix tests

* feat(vm_override): overrides are found at compile time

* merge with master

* expand overrides during inlining phase

* feat(vm_override): docs, tidy test

* disable overrides in inliner when flag is turned off

* fix `vm_override.enable` flag

* fix style

* fix test

* apply reviewer suggestions

* avoid redundant traversal

* move comment

* roll back change

* simplify inliner / overrider

Co-authored-by: E.W.Ayers <[email protected]>

library/init/category/option.lean

@@ -34,6 +34,9 @@ namespace option_t
   @[inline] protected def fail : option_t m α :=
   ⟨pure none⟩
 
+  @[inline] def of_option : option α → option_t m α
+  | o := ⟨pure o⟩
+
   instance : alternative (option_t m) :=
   { failure := @option_t.fail m _,
     orelse  := @option_t.orelse m _,

library/init/meta/float.lean

@@ -138,8 +138,10 @@ meta constant of_nat : nat → float
 meta constant of_int : int → float
 meta instance of_nat_coe : has_coe nat float := ⟨of_nat⟩
 meta instance of_int_coe : has_coe int float := ⟨of_int⟩
-meta instance : has_zero float := ⟨of_nat 0⟩
-meta instance : has_one float := ⟨of_nat 1⟩
+protected meta def zero : float := of_nat 0
+protected meta def one : float := of_nat 1
+meta instance : has_zero float := ⟨float.zero⟩
+meta instance : has_one float := ⟨float.one⟩
 
 meta constant to_repr : float → string
 meta instance : has_repr float := ⟨to_repr⟩
@@ -152,4 +154,4 @@ meta instance has_nat_pow : has_pow float nat :=
 ⟨λ a b, native.float.pow a (float.of_nat b)⟩
 
 end float
-end native
+end native

library/init/meta/vm.lean

@@ -67,6 +67,8 @@ meta constant pos       : vm_decl → option pos
 meta constant olean     : vm_decl → option string
 /-- Return names .olean file where the given VM declaration was imported from. -/
 meta constant args_info : vm_decl → list vm_local_info
+/-- If the given declaration is overridden by another declaration using the vm_override attribute, then this returns the overriding declaration.-/
+meta constant override_idx : vm_decl → option nat
 end vm_decl
 
 meta constant vm_core : Type → Type
@@ -81,7 +83,12 @@ meta instance : monad vm_core :=
 
 namespace vm
 meta constant get_env              : vm environment
+/-- Returns the vm declaration associated with the given name. Remark: does _not_ return the vm_override if present.-/
 meta constant get_decl             : name → vm vm_decl
+/-- Returns the vm declaration associated with the given index. Remark: does _not_ return the vm_override if present.-/
+meta constant decl_of_idx          : nat → vm vm_decl
+meta def get_override              : vm_decl → vm vm_decl
+| d := option_t.of_option d.override_idx >>= decl_of_idx
 meta constant get_options          : vm options
 meta constant stack_size           : vm nat
 /-- Return the vm_obj stored at the given position on the execution stack.

src/frontends/lean/builtin_cmds.cpp

@@ -442,7 +442,7 @@ static environment compile_cmd(parser & p) {
     declaration d = p.env().get(n);
     if (!d.is_definition())
         throw parser_error("invalid #compile command, declaration is not a definition", pos);
-    return vm_compile(p.env(), d);
+    return vm_compile(p.env(), p.get_options(), d);
 }
 
 static environment eval_cmd(parser & p) {
@@ -478,7 +478,7 @@ static environment eval_cmd(parser & p) {
     }
 
     name fn_name = "_main";
-    auto new_env = compile_expr(p.env(), fn_name, ls, type, e, pos);
+    auto new_env = compile_expr(p.env(), p.get_options(), fn_name, ls, type, e, pos);
 
     auto out = p.mk_message(p.cmd_pos(), p.pos(), INFORMATION);
     out.set_caption("eval result");

src/frontends/lean/decl_attributes.cpp

@@ -52,7 +52,7 @@ void decl_attributes::parse_core(parser & p, bool compact) {
             pre_val = mk_typed_expr(mk_constant(get_nat_name()), pre_val);
             expr nat = mk_constant(get_nat_name());
             expr val = p.elaborate("_attribute", list<expr>(), pre_val).first;
-            vm_obj prio = eval_closed_expr(p.env(), "_attribute", nat, val, pos);
+            vm_obj prio = eval_closed_expr(p.env(), p.get_options(), "_attribute", nat, val, pos);
             if (optional<unsigned> _prio = try_to_unsigned(prio)) {
                 m_prio = _prio;
             } else {
@@ -135,7 +135,7 @@ bool decl_attributes::ok_for_inductive_type() const {
     for (entry const & e : m_entries) {
         name const & n = e.m_attr->get_name();
         if (is_system_attribute(n)) {
-            if ((n != "class" && !is_class_symbol_tracking_attribute(n)) || e.deleted())
+            if ((n != "class" && n != "vm_override" && !is_class_symbol_tracking_attribute(n)) || e.deleted())
                 return false;
         }
     }

src/frontends/lean/definition_cmds.cpp

@@ -219,7 +219,7 @@ static environment compile_decl(parser_info const & p, environment const & env,
                                 name const & c_name, name const & c_real_name, pos_info const & pos) {
     try {
         time_task _("compilation", p.mk_message(message_severity::INFORMATION), p.get_options(), c_real_name);
-        return vm_compile(env, env.get(c_real_name));
+        return vm_compile(env, p.get_options(), env.get(c_real_name));
     } catch (exception & ex) {
         // FIXME(gabriel): use position from exception
         auto out = p.mk_message(pos, WARNING);

src/frontends/lean/elaborator.cpp

@@ -2311,7 +2311,7 @@ expr elaborator::mk_aux_meta_def(expr const & e, expr const & ref) {
     if (!is_constant(new_c)) {
         throw elaborator_exception(ref, "failed to create auxiliary definition");
     }
-    m_env = vm_compile(m_env, m_env.get(const_name(new_c)));
+    m_env = vm_compile(m_env, m_ctx.get_options(), m_env.get(const_name(new_c)));
     m_ctx.set_env(m_env);
     m_ctx.set_mctx(mctx);
     return new_c;

src/frontends/lean/notation_cmd.cpp

@@ -43,13 +43,14 @@ static unsigned parse_precedence_core(parser & p) {
         return p.parse_small_nat();
     } else {
         environment env = p.env();
+        options opts = p.get_options();
         env = open_prec_aliases(env);
         parser::local_scope scope(p, env);
         expr pre_val = p.parse_expr(get_max_prec());
         expr nat = mk_constant(get_nat_name());
         pre_val  = mk_typed_expr(nat, pre_val);
         expr val = p.elaborate("notation", list<expr>(), pre_val).first;
-        vm_obj p = eval_closed_expr(env, "_precedence", nat, val, pos);
+        vm_obj p = eval_closed_expr(env, opts, "_precedence", nat, val, pos);
         if (optional<unsigned> _p = try_to_unsigned(p)) {
             return *_p;
         } else {

src/frontends/lean/structure_cmd.cpp

@@ -1255,7 +1255,7 @@ struct structure_cmd_fn {
                                                                              coercion_type, coercion_value, use_conv_opt);
             m_env = module::add(m_env, check(m_env, coercion_decl));
             add_alias(coercion_name);
-            m_env = vm_compile(m_env, m_env.get(coercion_name));
+            m_env = vm_compile(m_env, m_p.get_options(), m_env.get(coercion_name));
             if (!m_private_parents[i]) {
                 if (m_meta_info.m_attrs.has_class() && is_class(m_env, parent_name)) {
                     // if both are classes, then we also mark coercion_name as an instance

src/frontends/lean/util.cpp

@@ -512,19 +512,22 @@ void initialize_frontend_lean_util() {
                                 });
 }
 
-environment compile_expr(environment const & env, name const & n, level_param_names const & ls, expr const & type, expr const & e, pos_info const & pos) {
+environment compile_expr(environment const & env, options const & opts,
+                         name const & n, level_param_names const & ls,
+                         expr const & type, expr const & e, pos_info const & pos) {
     environment new_env = env;
     bool use_conv_opt   = true;
     bool is_trusted     = false;
     auto cd = check(new_env, mk_definition(new_env, n, ls, type, e, use_conv_opt, is_trusted));
     new_env = new_env.add(cd);
     new_env = add_transient_decl_pos_info(new_env, n, pos);
-    return vm_compile(new_env, new_env.get(n));
+    return vm_compile(new_env, opts, new_env.get(n));
 }
 
-vm_obj eval_closed_expr(environment const & env, name const & n, expr const & type, expr const & e, pos_info const & pos) {
-    environment new_env = compile_expr(env, n, {}, type, e, pos);
-    vm_state vm(new_env, {});
+vm_obj eval_closed_expr(environment const & env, options const & opts,
+                        name const & n, expr const & type, expr const & e, pos_info const & pos) {
+    environment new_env = compile_expr(env, opts, n, {}, type, e, pos);
+    vm_state vm(new_env, opts);
     return vm.invoke(n, 0, nullptr);
 }
 

src/frontends/lean/util.h

@@ -134,8 +134,10 @@ bool is_anonymous_field_notation(expr const & e);
 name const & get_field_notation_field_name(expr const & e);
 unsigned get_field_notation_field_idx(expr const & e);
 
-environment compile_expr(environment const & env, name const & n, level_param_names const & ls, expr const & type, expr const & e, pos_info const & pos);
-vm_obj eval_closed_expr(environment const & env, name const & n, expr const & type, expr const & e, pos_info const & pos);
+environment compile_expr(environment const & env, options const & opts, name const & n,
+                         level_param_names const & ls, expr const & type, expr const & e, pos_info const & pos);
+vm_obj eval_closed_expr(environment const & env, options const & opts, name const & n,
+                        expr const & type, expr const & e, pos_info const & pos);
 
 expr mk_lean_list(parser & p, buffer<expr> const & es, pos_info const & pos);
 expr mk_lean_list(buffer<expr> const & es);

src/kernel/inductive/inductive.cpp

@@ -185,6 +185,11 @@ name get_elim_name(name const & n) {
     return n + name("rec");
 }
 
+/**\ brief Return recursor name */
+name get_cases_on_name(name const & n) {
+    return n + name("cases_on");
+}
+
 environment certified_inductive_decl::add_constant(environment const & env, name const & n, level_param_names const & ls,
                                                    expr const & t) const {
     return env.add(certify_unchecked::certify_or_check(env, mk_constant_assumption(n, ls, t, m_is_trusted)));

src/kernel/inductive/inductive.h

@@ -142,6 +142,7 @@ bool has_dep_elim(environment const & env, name const & n);
 
 /** \brief Return the eliminator/recursor associated with an inductive datatype */
 name get_elim_name(name const & n);
+name get_cases_on_name(name const & n);
 }
 void initialize_inductive_module();
 void finalize_inductive_module();

src/library/compiler/inliner.cpp

@@ -8,10 +8,10 @@ Author: Leonardo de Moura
 #include "kernel/inductive/inductive.h"
 #include "library/util.h"
 #include "library/module.h"
-#include "library/trace.h"
 #include "library/normalize.h"
 #include "library/attribute_manager.h"
 #include "library/vm/vm.h"
+#include "library/vm/vm_override.h"
 #include "library/compiler/util.h"
 #include "library/compiler/compiler_step_visitor.h"
 
@@ -43,6 +43,8 @@ void finalize_inliner() {
 }
 
 class inline_simple_definitions_fn : public compiler_step_visitor {
+    bool m_enable_overrides;
+
     /* Return true iff v is of the form (g y_1 ... y_n) where
        y_i is a constant or a variable.
        Moreover, y_i's variables are pairwise distinct. */
@@ -77,6 +79,14 @@ class inline_simple_definitions_fn : public compiler_step_visitor {
             return new_e;
     }
 
+    expr default_visit_constant(expr const & e) {
+        expr new_e = compiler_step_visitor::visit_constant(e);
+        if (new_e != e)
+            return visit(new_e);
+        else
+            return new_e;
+    }
+
     bool is_nonrecursive_recursor(name const & n) {
         if (auto I_name = inductive::is_elim_rule(env(), n)) {
             return !is_recursive_datatype(env(), *I_name);
@@ -129,11 +139,37 @@ class inline_simple_definitions_fn : public compiler_step_visitor {
         return ctx().reduce_projection(e);
     }
 
-    virtual expr visit_app(expr const & e) override {
-        expr const & fn = get_app_fn(e);
+    expr apply_overrides(expr const & e) {
+        if (m_enable_overrides) {
+            buffer<expr> args;
+            expr const & fn = get_app_args(e, args);
+            if (is_constant(fn)) {
+                auto n = const_name(fn);
+                if (auto override_name = get_vm_override_name(m_env, n, m_enable_overrides)) {
+                    return mk_app(mk_constant(*override_name, const_levels(fn)), args);
+                }
+            }
+        }
+        return e;
+    }
+
+    virtual expr visit_constant(expr const & e) override {
+        if (m_enable_overrides) {
+            name n = const_name(e);
+            if (auto override_name = get_vm_override_name(m_env, n, m_enable_overrides)) {
+                return default_visit_constant(mk_constant(*override_name, const_levels(e)));
+            }
+        }
+        return default_visit_constant(e);
+    }
+
+    virtual expr visit_app(expr const & arg) override {
+        expr e = apply_overrides(arg);
+        expr fn = get_app_fn(e);
         if (!is_constant(fn))
             return default_visit_app(e);
-        name const & n  = const_name(fn);
+        name const & n = const_name(fn);
+
         if (is_vm_builtin_function(n) || is_pack_unpack(env(), e))
             return default_visit_app(e);
         if (is_cases_on_recursor(env(), n) || is_nonrecursive_recursor(n))
@@ -164,11 +200,14 @@ class inline_simple_definitions_fn : public compiler_step_visitor {
     }
 
 public:
-    inline_simple_definitions_fn(environment const & env, abstract_context_cache & cache):
-        compiler_step_visitor(env, cache) {}
+    inline_simple_definitions_fn(environment const & env, options const opts,
+                                 abstract_context_cache & cache):
+        compiler_step_visitor(env, cache),
+        m_enable_overrides(get_vm_override_enabled(opts)) { }
 };
 
-expr inline_simple_definitions(environment const & env, abstract_context_cache & cache, expr const & e) {
-    return inline_simple_definitions_fn(env, cache)(e);
+expr inline_simple_definitions(environment const & env, options const & opts,
+                               abstract_context_cache & cache, expr const & e) {
+    return inline_simple_definitions_fn(env, opts, cache)(e);
 }
 }

src/library/compiler/inliner.h

@@ -22,7 +22,8 @@ bool is_inline(environment const & env, name const & n);
     This procedure also simplifies projection applications.
 
     Example: this procedure reduces (@add nat nat_has_add a b) into (nat.add a b). */
-expr inline_simple_definitions(environment const & env, abstract_context_cache & cache, expr const & e);
+expr inline_simple_definitions(environment const & env, options const & opts,
+                               abstract_context_cache & cache, expr const & e);
 
 void initialize_inliner();
 void finalize_inliner();

src/library/compiler/preprocess.cpp

@@ -161,6 +161,7 @@ static name * g_tmp_prefix = nullptr;
 
 class preprocess_fn {
     environment    m_env;
+    options        m_opts;
     context_cache  m_cache;
 
     bool check(declaration const & d, expr const & v) {
@@ -230,15 +231,15 @@ class preprocess_fn {
     }
 
 public:
-    preprocess_fn(environment const & env):
-        m_env(env) {}
+    preprocess_fn(environment const & env, options const & opts):
+        m_env(env), m_opts(opts) {}
 
     void operator()(declaration const & d, buffer<procedure> & procs) {
         if (compile_irrelevant(d, procs))
             return;
         expr v = d.get_value();
         lean_trace(name({"compiler", "input"}), tout() << "\n" << v << "\n";);
-        v = inline_simple_definitions(m_env, m_cache, v);
+        v = inline_simple_definitions(m_env, m_opts, m_cache, v);
         lean_cond_assert("compiler", check(d, v));
         lean_trace(name({"compiler", "inline"}), tout() << "\n" << v << "\n";);
         v = expand_aux(m_env, m_cache, v);
@@ -275,14 +276,15 @@ class preprocess_fn {
     }
 };
 
-void preprocess(environment const & env, declaration const & d, buffer<procedure> & result) {
-    return preprocess_fn(env)(d, result);
+void preprocess(environment const & env, options const & opts,
+                declaration const & d, buffer<procedure> & result) {
+    return preprocess_fn(env, opts)(d, result);
 }
 
-void preprocess(environment const & env, buffer<declaration> const & ds, buffer<procedure> & result) {
+void preprocess(environment const & env, options const & opts, buffer<declaration> const & ds, buffer<procedure> & result) {
     for (declaration const & d : ds) {
         buffer<procedure> procs;
-        preprocess(env, d, procs);
+        preprocess(env, opts, d, procs);
         result.append(procs);
     }
 }

src/library/compiler/preprocess.h

@@ -12,11 +12,11 @@ namespace lean {
     put definition in eta-expanded normal form, and
     eliminate nested (recursive) recursor applications.
     Nested recurse applications become new procedures. */
-void preprocess(environment const & env, declaration const & d, buffer<procedure> & result);
+void preprocess(environment const & env, options const & opts, declaration const & d, buffer<procedure> & result);
 
 /** \brief Similar to previous function, but supports a collection \c ds of potentially mutually recursive
     definitions. */
-void preprocess(environment const & env, buffer<declaration> const & ds, buffer<procedure> & result);
+void preprocess(environment const & env, options const & opts, buffer<declaration> const & ds, buffer<procedure> & result);
 
 void initialize_preprocess();
 void finalize_preprocess();