From 340cade57541240e13bc63cdca9962cd7991272d Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sun, 20 Dec 2020 17:48:46 -0800
Subject: [PATCH] fix: bug at specialize.cpp

---
 src/library/compiler/specialize.cpp | 11 +++++++--
 tests/lean/run/specbug.lean         | 38 +++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 2 deletions(-)
 create mode 100644 tests/lean/run/specbug.lean

diff --git a/src/library/compiler/specialize.cpp b/src/library/compiler/specialize.cpp
index fe1261b64b..edb30840fd 100644
--- a/src/library/compiler/specialize.cpp
+++ b/src/library/compiler/specialize.cpp
@@ -893,7 +893,7 @@ class specialize_fn {
         // lean_trace(name("compiler", "spec_info"), tout() << "STEP 2 " << n << "\n" << code << "\n";);
         code = csimp(env(), code, m_cfg);
         code = visit(code);
-        // lean_trace(name("compiler", "spec_info"), tout() << "STEP 3 " << n << "\n" << code << "\n";);
+        // lean_trace(name("compiler", "specialize"), tout() << "new code " << n << "\n" << trace_pp_expr(code) << "\n";);
         m_new_decls.push_back(comp_decl(n, code));
     }
 
@@ -984,7 +984,13 @@ class specialize_fn {
         }
         optional<name> new_fn_name;
         expr key;
-        if (gcache_enabled) {
+        /* When `m_params.size > 1`, it is not safe to reuse cached specialization.
+           See test `tests/lean/run/specbug.lean`.
+           This is a bit hackish, but should not produce increase the generated code size too much.
+           On Dec 20 2020, before this fix, 5246 specializations were reused, but only 11 had `m_params.size > 1`.
+           This file will be deleted. So, it is not worth designing a better caching scheme.
+           TODO: when we reimplement this module in Lean, we should have a better caching heuristic. */
+        if (gcache_enabled && ctx.m_params.size() <= 1) {
             key = mk_app(fn, gcache_key_args);
             if (optional<name> it = get_cached_specialization(env(), key)) {
                 lean_trace(name({"compiler", "specialize"}), tout() << "get_cached_specialization [" << ctx.m_params.size() << "]: " << *it << "\n";
@@ -995,6 +1001,7 @@ class specialize_fn {
                                i++;
                            }
                            tout() << trace_pp_expr(key) << "\n";);
+                // std::cerr << *it << " " << ctx.m_vars.size() << " " << ctx.m_params.size() << "\n";
                 new_fn_name = *it;
             }
         }
diff --git a/tests/lean/run/specbug.lean b/tests/lean/run/specbug.lean
new file mode 100644
index 0000000000..e19a6912a7
--- /dev/null
+++ b/tests/lean/run/specbug.lean
@@ -0,0 +1,38 @@
+@[noinline] def f (x : Bool)   := x
+@[noinline] def g (x y : Bool) := x
+
+def h (x : Bool) (xs : List Nat) : List Bool :=
+  match x with
+  | true =>
+    let z := f true
+    let y := f false
+    xs.map fun x => g y z
+  | false =>
+    let y := f false
+    let z := f true
+    xs.map fun x => g y z
+
+theorem ex1 : h true [1] = h false [1] := rfl
+
+#eval h true [1]
+#eval h false [1]
+
+theorem ex2 : (h true [1] == h false [1]) = true :=
+  nativeDecide!
+
+@[noinline] def f2 (a : String) := a
+@[noinline] def g2 (a : String) (x : Bool) := a
+
+def h2 (x : Bool) (xs : List Nat) : List String :=
+  match x with
+  | false =>
+    let a := f2 "a"
+    let y := f false
+    xs.map fun x => g2 a y
+  | true =>
+    let y := f false
+    let a := f2 "a"
+    xs.map fun x => g2 a y
+
+#eval h2 true [1]
+#eval h2 false [1]