feat(library/compiler): add cse to new compiler stack

src/library/compiler/CMakeLists.txt

@@ -4,5 +4,5 @@ add_library(compiler OBJECT util.cpp eta_expansion.cpp preprocess.cpp
   lambda_lifting.cpp simp_inductive.cpp nat_value.cpp
   vm_compiler.cpp old_cse.cpp elim_unused_lets.cpp extract_values.cpp init_module.cpp
   ## New compiler
-  lc_util.cpp lcnf.cpp csimp.cpp elim_dead_let.cpp
+  lc_util.cpp lcnf.cpp csimp.cpp elim_dead_let.cpp cse.cpp
   )

src/library/compiler/cse.cpp

@@ -0,0 +1,141 @@
+/*
+Copyright (c) 2018 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Author: Leonardo de Moura
+*/
+#include <algorithm>
+#include "util/name_generator.h"
+#include "kernel/environment.h"
+#include "kernel/instantiate.h"
+#include "kernel/abstract.h"
+#include "kernel/expr_maps.h"
+#include "library/compiler/lc_util.h"
+
+namespace lean {
+static name * g_cse_fresh = nullptr;
+
+class cse_fn {
+    environment       m_env;
+    name_generator    m_ngen;
+    expr_map<expr>    m_lval2fvar;
+    std::vector<expr> m_lvals;
+public:
+
+    expr visit_let(expr e) {
+        unsigned lvals_size = m_lvals.size();
+        buffer<expr> fvars;
+        buffer<expr> to_keep_fvars;
+        buffer<std::tuple<name, expr, expr>> entries;
+        while (is_let(e)) {
+            expr value = instantiate_rev(let_value(e), fvars.size(), fvars.data());
+            auto it    = m_lval2fvar.find(value);
+            if (it != m_lval2fvar.end()) {
+                lean_assert(is_fvar(it->second));
+                fvars.push_back(it->second);
+            } else {
+                expr type      = instantiate_rev(let_type(e), fvars.size(), fvars.data());
+                expr new_value = visit_let_value(value);
+                expr fvar      = mk_fvar(m_ngen.next());
+                fvars.push_back(fvar);
+                to_keep_fvars.push_back(fvar);
+                entries.emplace_back(let_name(e), type, new_value);
+                if (!is_cases_on_app(m_env, new_value)) {
+                    m_lval2fvar.insert(mk_pair(new_value, fvar));
+                    m_lvals.push_back(new_value);
+                }
+            }
+            e = let_body(e);
+        }
+        e = instantiate_rev(e, fvars.size(), fvars.data());
+        e = abstract(e, to_keep_fvars.size(), to_keep_fvars.data());
+        lean_assert(entries.size() == to_keep_fvars.size());
+        unsigned i = entries.size();
+        while (i > 0) {
+            --i;
+            expr new_type  = abstract(std::get<1>(entries[i]), i, to_keep_fvars.data());
+            expr new_value = abstract(std::get<2>(entries[i]), i, to_keep_fvars.data());
+            e = mk_let(std::get<0>(entries[i]), new_type, new_value, e);
+        }
+        /* Restore m_lval2fvar */
+        for (unsigned i = lvals_size; i < m_lvals.size(); i++) {
+            m_lval2fvar.erase(m_lvals[i]);
+        }
+        m_lvals.resize(lvals_size);
+        return e;
+    }
+
+    expr visit_lambda(expr e) {
+        buffer<expr> fvars;
+        buffer<std::tuple<name, expr, binder_info>> entries;
+        while (is_lambda(e)) {
+            expr domain = instantiate_rev(binding_domain(e), fvars.size(), fvars.data());
+            expr fvar   = mk_fvar(m_ngen.next());
+            entries.emplace_back(binding_name(e), domain, binding_info(e));
+            fvars.push_back(fvar);
+            e = binding_body(e);
+        }
+        e = visit(instantiate_rev(e, fvars.size(), fvars.data()));
+        e = abstract(e, fvars.size(), fvars.data());
+        unsigned i = entries.size();
+        while (i > 0) {
+            --i;
+            expr new_domain = abstract(std::get<1>(entries[i]), i, fvars.data());
+            e = mk_lambda(std::get<0>(entries[i]), new_domain, e, std::get<2>(entries[i]));
+        }
+        return e;
+    }
+
+    expr visit_app(expr const & e) {
+        if (is_cases_on_app(m_env, e)) {
+            buffer<expr> args;
+            expr const & c = get_app_args(e, args);
+            lean_assert(is_constant(c));
+            inductive_val I_val      = m_env.get(const_name(c).get_prefix()).to_inductive_val();
+            unsigned first_minor_idx = I_val.get_nparams() + 1 /* typeformer/motive */ + I_val.get_nindices() + 1;
+            for (unsigned i = first_minor_idx; i < args.size(); i++) {
+
+                args[i] = visit(args[i]);
+            }
+            return mk_app(c, args);
+        } else {
+            return e;
+        }
+    }
+
+    expr visit_let_value(expr const & e) {
+        switch (e.kind()) {
+        case expr_kind::Lambda: return visit_lambda(e);
+        case expr_kind::App:    return visit_app(e);
+        default:                return e;
+        }
+    }
+
+    expr visit(expr const & e) {
+        switch (e.kind()) {
+        case expr_kind::Lambda: return visit_lambda(e);
+        case expr_kind::Let:    return visit_let(e);
+        default:                return e;
+        }
+    }
+
+public:
+    cse_fn(environment const & env):
+        m_env(env), m_ngen(*g_cse_fresh) {
+    }
+
+    expr operator()(expr const & e) { return visit(e); }
+};
+
+expr cse(environment const & env, expr const & e) {
+    return cse_fn(env)(e);
+}
+
+void initialize_cse() {
+    g_cse_fresh = new name("_cse_fresh");
+    register_name_generator_prefix(*g_cse_fresh);
+}
+void finalize_cse() {
+    delete g_cse_fresh;
+}
+}

src/library/compiler/cse.h

@@ -0,0 +1,13 @@
+/*
+Copyright (c) 2018 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Author: Leonardo de Moura
+*/
+#pragma once
+#include "kernel/environment.h"
+namespace lean {
+expr cse(environment const & env, expr const & e);
+void initialize_cse();
+void finalize_cse();
+}

src/library/compiler/elim_dead_let.cpp

@@ -128,7 +128,7 @@ public:
 expr elim_dead_let(expr const & e) { return elim_dead_let_fn()(e); }
 
 void initialize_elim_dead_let() {
-    g_elim_dead_let_fresh = new name("_elim_dead_let");
+    g_elim_dead_let_fresh = new name("_elim_dead_let_fresh");
     register_name_generator_prefix(*g_elim_dead_let_fresh);
 }
 void finalize_elim_dead_let() {

src/library/compiler/init_module.cpp

@@ -15,6 +15,7 @@ Author: Leonardo de Moura
 
 #include "library/compiler/lcnf.h"
 #include "library/compiler/elim_dead_let.h"
+#include "library/compiler/cse.h"
 
 namespace lean {
 void initialize_compiler_module() {
@@ -29,9 +30,11 @@ void initialize_compiler_module() {
     //======
     initialize_lcnf();
     initialize_elim_dead_let();
+    initialize_cse();
 }
 
 void finalize_compiler_module() {
+    finalize_cse();
     finalize_elim_dead_let();
     finalize_lcnf();
     //======

src/library/compiler/lc_util.h

@@ -8,6 +8,7 @@ Author: Leonardo de Moura
 #include "kernel/expr.h"
 #include "kernel/type_checker.h"
 #include "library/constants.h"
+#include "library/util.h"
 
 namespace lean {
 /* Return true if the given argument is mdata relevant to the compiler

src/library/compiler/preprocess.cpp

@@ -39,6 +39,7 @@ Author: Leonardo de Moura
 #include "library/compiler/lcnf.h"
 #include "library/compiler/csimp.h"
 #include "library/compiler/elim_dead_let.h"
+#include "library/compiler/cse.h"
 
 namespace lean {
 class expand_aux_fn : public compiler_step_visitor {
@@ -250,6 +251,10 @@ public:
                    expr r3 = elim_dead_let(r2);
                    tout() << r3 << "\n";
                    check(d, r3);
+                   tout() << ">> CSE\n";
+                   expr r4 = cse(m_env, r3);
+                   tout() << r4 << "\n";
+                   check(d, r4);
             });
         v = inline_simple_definitions(m_env, m_cache, v);
         lean_cond_assert("compiler", check(d, v));

tests/lean/run/new_compiler.lean

@@ -21,6 +21,11 @@ let y_1 := x in
 let y_2 := y_1 in
 y_2 + n
 
+def cse_tst (n : nat) : nat :=
+let y := nat.succ ((λ x, x) n) in
+let z := nat.succ n in
+y + z
+
 def tst1 (n : nat) : nat :=
 let p := (nat.succ n, n) in
 let q := (p, p) in