From 6cfdfe99422b420281ddbfd8acc830e948aa0ac0 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sat, 21 Aug 2021 12:22:15 -0700
Subject: [PATCH] feat: apply `csimp` attribute constant replacements

---
 src/library/compiler/compiler.cpp      |  7 +++++++
 tests/lean/csimpAttr.lean              | 17 +++++++++++++++++
 tests/lean/csimpAttr.lean.expected.out |  6 ++++++
 3 files changed, 30 insertions(+)
 create mode 100644 tests/lean/csimpAttr.lean
 create mode 100644 tests/lean/csimpAttr.lean.expected.out

diff --git a/src/library/compiler/compiler.cpp b/src/library/compiler/compiler.cpp
index f4872e6ca8..050a6be395 100644
--- a/src/library/compiler/compiler.cpp
+++ b/src/library/compiler/compiler.cpp
@@ -159,6 +159,12 @@ bool is_main_fn_type(expr const & type) {
 
 #define trace_compiler(k, ds) lean_trace(k, trace_comp_decls(ds););
 
+extern "C" object* lean_csimp_replace_constants(object* env, object* n);
+
+expr csimp_replace_constants(environment const & env, expr const & e) {
+    return expr(lean_csimp_replace_constants(env.to_obj_arg(), e.to_obj_arg()));
+}
+
 environment compile(environment const & env, options const & opts, names cs) {
     if (!is_codegen_enabled(opts))
         return env;
@@ -206,6 +212,7 @@ environment compile(environment const & env, options const & opts, names cs) {
     // trace(ds);
     ds = apply(cce, env, ds);
     trace_compiler(name({"compiler", "cce"}), ds);
+    ds = apply(csimp_replace_constants, env, ds);
     ds = apply(simp, env, ds);
     trace_compiler(name({"compiler", "simp"}), ds);
     // trace(ds);
diff --git a/tests/lean/csimpAttr.lean b/tests/lean/csimpAttr.lean
new file mode 100644
index 0000000000..1a27ed0a58
--- /dev/null
+++ b/tests/lean/csimpAttr.lean
@@ -0,0 +1,17 @@
+def foo (x : Nat) :=
+  2*x
+
+def boo (x : Nat) :=
+  x + x
+
+@[csimp] theorem foo_eq_boo1 (x : Nat) : foo x = boo x := by -- Error
+  simp [foo, boo, Nat.mul_comm]
+  show (x * 1) + x = x + x
+  simp
+
+@[csimp] theorem foo_eq_boo2 : foo = boo :=
+  funext foo_eq_boo1
+
+set_option trace.compiler.ir.init true
+def f (x : Nat) : Nat :=
+  foo (foo x)
diff --git a/tests/lean/csimpAttr.lean.expected.out b/tests/lean/csimpAttr.lean.expected.out
new file mode 100644
index 0000000000..6cbefeae02
--- /dev/null
+++ b/tests/lean/csimpAttr.lean.expected.out
@@ -0,0 +1,6 @@
+
+[init]
+def f (x_1 : obj) : obj :=
+  let x_2 : obj := Nat.add x_1 x_1;
+  let x_3 : obj := Nat.add x_2 x_2;
+  ret x_3csimpAttr.lean:7:17-7:28: error: invalid 'csimp' theorem, only constant replacement theorems (e.g., `@f = @g`) are currently supported.