From 0bcd07f0761a18dd2f8defe966e0064001da2b99 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 22 Oct 2018 10:19:39 -0700
Subject: [PATCH] feat(library/compiler/compiler): cache "stage2"

---
 src/library/compiler/compiler.cpp | 16 ++++++++++++++++
 src/library/compiler/util.h       |  6 ++++++
 2 files changed, 22 insertions(+)

diff --git a/src/library/compiler/compiler.cpp b/src/library/compiler/compiler.cpp
index 08bd7a44aa..5fc98ee81a 100644
--- a/src/library/compiler/compiler.cpp
+++ b/src/library/compiler/compiler.cpp
@@ -84,6 +84,19 @@ static environment cache_stage1(environment env, comp_decls const & ds) {
     return env;
 }
 
+static environment cache_stage2(environment env, comp_decls const & ds) {
+    for (comp_decl const & d : ds) {
+        name n = d.fst();
+        expr v = d.snd();
+        /* This a temporary hack to store Stage2 intermediate result.
+           We should not store this information as a declaration. */
+        declaration aux_decl = mk_definition(mk_cstage2_name(n), names(), mk_enf_object_type(),
+                                             v, reducibility_hints::mk_opaque(), true);
+        env = module::add(env, aux_decl, false);
+    }
+    return env;
+}
+
 #define trace_compiler(k, ds) lean_trace(k, trace(ds);)
 
 environment compile(environment const & env, options const & opts, names const & cs) {
@@ -134,6 +147,8 @@ environment compile(environment const & env, options const & opts, names const &
     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(simp_inductive, new_env, ds);
     trace_compiler(name({"compiler", "simplify_inductive"}), ds);
     new_env = emit_bytecode(new_env, ds);
@@ -155,6 +170,7 @@ void initialize_compiler() {
     register_trace_class({"compiler", "cse"});
     register_trace_class({"compiler", "specialize"});
     register_trace_class({"compiler", "stage1"});
+    register_trace_class({"compiler", "stage2"});
     register_trace_class({"compiler", "erase_irrelevant"});
     register_trace_class({"compiler", "lambda_lifting"});
     register_trace_class({"compiler", "extract_closed"});
diff --git a/src/library/compiler/util.h b/src/library/compiler/util.h
index 8bc014285f..0568d6df65 100644
--- a/src/library/compiler/util.h
+++ b/src/library/compiler/util.h
@@ -83,6 +83,12 @@ inline name mk_cstage1_name(name const & decl_name) {
     return name(decl_name, "_cstage1");
 }
 
+/* Create an auxiliary names for a declaration that saves the result of the compilation
+   after step erasure. */
+inline name mk_cstage2_name(name const & decl_name) {
+    return name(decl_name, "_cstage2");
+}
+
 /* Set `used[i] = true` if `fvars[i]` occurs in `e` */
 void mark_used_fvars(expr const & e, buffer<expr> const & fvars, buffer<bool> & used);