From 01229425a2c514b67090a3046be0050886798ac6 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 22 Oct 2018 15:49:07 -0700
Subject: [PATCH] feat(library/compiler/compiler): inline "small" stage2
 declarations at `csimp` after erasure

---
 src/library/compiler/compiler.cpp |  6 ++++--
 src/library/compiler/csimp.cpp    | 33 ++++++++++++++++++++-----------
 2 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/src/library/compiler/compiler.cpp b/src/library/compiler/compiler.cpp
index 5fc98ee81a..a42ac36328 100644
--- a/src/library/compiler/compiler.cpp
+++ b/src/library/compiler/compiler.cpp
@@ -134,7 +134,7 @@ environment compile(environment const & env, options const & opts, names const &
     trace_compiler(name({"compiler", "elim_dead_let"}), ds);
     ds = apply(erase_irrelevant, new_env, ds);
     trace_compiler(name({"compiler", "erase_irrelevant"}), ds);
-    ds = apply(esimp, env, ds);
+    ds = apply(esimp, new_env, ds);
     trace_compiler(name({"compiler", "simp"}), ds);
     ds = apply(elim_dead_let, ds);
     trace_compiler(name({"compiler", "elim_dead_let"}), ds);
@@ -142,13 +142,15 @@ environment compile(environment const & env, options const & opts, names const &
     trace_compiler(name({"compiler", "cse"}), ds);
     ds = lambda_lifting(new_env, ds);
     trace_compiler(name({"compiler", "lambda_lifting"}), ds);
-    ds = apply(esimp, env, ds);
+    ds = apply(esimp, new_env, ds);
     trace_compiler(name({"compiler", "simp"}), ds);
     std::tie(new_env, ds) = extract_closed(new_env, ds);
     ds = apply(elim_dead_let, ds);
     trace_compiler(name({"compiler", "extract_closed"}), ds);
     new_env = cache_stage2(new_env, ds);
     trace_compiler(name({"compiler", "stage2"}), ds);
+    ds = apply(esimp, new_env, ds);
+    trace_compiler(name({"compiler", "simp"}), ds);
     ds = apply(simp_inductive, new_env, ds);
     trace_compiler(name({"compiler", "simplify_inductive"}), ds);
     new_env = emit_bytecode(new_env, ds);
diff --git a/src/library/compiler/csimp.cpp b/src/library/compiler/csimp.cpp
index 8e3d2b7a9b..ba325d8518 100644
--- a/src/library/compiler/csimp.cpp
+++ b/src/library/compiler/csimp.cpp
@@ -1087,19 +1087,30 @@ class csimp_fn {
     expr try_inline(expr const & fn, expr const & e, bool is_let_val) {
         lean_assert(is_constant(fn));
         lean_assert(is_eqp(find(get_app_fn(e)), fn));
-        if (!m_before_erasure) return e;
         if (!m_cfg.m_inline) return e;
         if (has_noinline_attribute(env(), const_name(fn))) return e;
-        name c = mk_cstage1_name(const_name(fn));
-        optional<constant_info> info = env().find(c);
-        if (!info || !info->is_definition()) return e;
-        if (get_app_num_args(e) < get_num_nested_lambdas(info->get_value())) return e;
-        if (!has_inline_attribute(env(), const_name(fn)) &&
-            get_lcnf_size(env(), info->get_value()) > m_cfg.m_inline_threshold)
-            return e;
-        if (is_recursive(c)) return e;
-        expr new_fn = instantiate_value_lparams(*info, const_levels(fn));
-        return beta_reduce(new_fn, e, is_let_val);
+        if (m_before_erasure) {
+            name c = mk_cstage1_name(const_name(fn));
+            optional<constant_info> info = env().find(c);
+            if (!info || !info->is_definition()) return e;
+            if (get_app_num_args(e) < get_num_nested_lambdas(info->get_value())) return e;
+            if (!has_inline_attribute(env(), const_name(fn)) &&
+                get_lcnf_size(env(), info->get_value()) > m_cfg.m_inline_threshold)
+                return e;
+            if (is_recursive(c)) return e;
+            expr new_fn = instantiate_value_lparams(*info, const_levels(fn));
+            return beta_reduce(new_fn, e, is_let_val);
+        } else {
+            name c = mk_cstage2_name(const_name(fn));
+            optional<constant_info> info = env().find(c);
+            if (!info || !info->is_definition()) return e;
+            unsigned arity = get_num_nested_lambdas(info->get_value());
+            if (arity == 0 || get_app_num_args(e) < arity) return e;
+            if (get_lcnf_size(env(), info->get_value()) > m_cfg.m_inline_threshold)
+                return e;
+            if (is_recursive(c)) return e;
+            return beta_reduce(info->get_value(), e, is_let_val);
+        }
     }
 
     expr visit_inline_app(expr const & e, bool is_let_val) {