From ea91317f1a29c658ca1b587f4f2ab156efdc07f3 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 15 Mar 2021 19:12:57 -0700
Subject: [PATCH] fix: avoid nontermination due to respecialization

---
 src/library/compiler/specialize.cpp |  7 ++++++-
 tests/lean/run/deep1.lean           | 12 ++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)
 create mode 100644 tests/lean/run/deep1.lean

diff --git a/src/library/compiler/specialize.cpp b/src/library/compiler/specialize.cpp
index fac7e4a130..d45033565d 100644
--- a/src/library/compiler/specialize.cpp
+++ b/src/library/compiler/specialize.cpp
@@ -984,6 +984,7 @@ class specialize_fn {
         if (!is_specialize_candidate(fn, args))
             return none_expr();
         // lean_trace(name("compiler", "specialize"), tout() << "specialize: " << fn << "\n";);
+        bool has_attr = has_specialize_attribute(const_name(fn));
         specialize_init_deps(fn, args, ctx);
         buffer<bool> bmask;
         get_bool_mask(const_name(fn), args.size(), bmask);
@@ -1060,7 +1061,11 @@ class specialize_fn {
                     return none_expr();
                 }
             }
-            if (respecialize) {
+            /* We should only re-specialize if the original function was marked with `[specialize]` attribute.
+               Recall that we always specialize functions containing instance implicit arguments.
+               This is a temporary workaround until we implement a proper code specializer.
+            */
+            if (has_attr && respecialize) {
                 for (comp_decl const & new_decl : new_decls) {
                     m_to_respecialize.insert(new_decl.fst());
                 }
diff --git a/tests/lean/run/deep1.lean b/tests/lean/run/deep1.lean
new file mode 100644
index 0000000000..26748d10d6
--- /dev/null
+++ b/tests/lean/run/deep1.lean
@@ -0,0 +1,12 @@
+partial def recurseM [Monad μ] (curr: α) (action: α -> μ (List α)) : μ PUnit := do
+  let children ← action curr
+  children.forM (recurseM · action)
+
+def specificTraverseList : Option Unit := recurseM () (fun _ => some [])
+
+partial def recurseM2 [Monad μ] (curr: α) (action: α -> μ (Array α)) : μ PUnit := do
+  let children ← action curr
+  children.forM (recurseM2 · action)
+
+def specificTraverseArray : Option Unit :=
+  recurseM2 () (fun _ => some #[])