From 91dec25a353d4479fe849e3320ccd5b1df1f80aa Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sun, 6 Dec 2020 18:52:28 -0800
Subject: [PATCH] fix: bug at `runST` and `runEST`

---
 src/Init/System/ST.lean               |  6 ++++--
 tests/lean/runSTBug.lean              | 19 +++++++++++++++++++
 tests/lean/runSTBug.lean.expected.out |  3 +++
 3 files changed, 26 insertions(+), 2 deletions(-)
 create mode 100644 tests/lean/runSTBug.lean
 create mode 100644 tests/lean/runSTBug.lean.expected.out

diff --git a/src/Init/System/ST.lean b/src/Init/System/ST.lean
index 26bc65f582..a5596bb5a9 100644
--- a/src/Init/System/ST.lean
+++ b/src/Init/System/ST.lean
@@ -22,12 +22,14 @@ class STWorld (σ : outParam Type) (m : Type → Type)
 instance {σ m n} [STWorld σ m] [MonadLift m n] : STWorld σ n := ⟨⟩
 instance {ε σ} : STWorld σ (EST ε σ) := ⟨⟩
 
-@[inline] def runEST {ε α : Type} (x : forall (σ : Type), EST ε σ α) : Except ε α :=
+@[noinline nospecialize]
+def runEST {ε α : Type} (x : forall (σ : Type), EST ε σ α) : Except ε α :=
   match x Unit () with
   | EStateM.Result.ok a _     => Except.ok a
   | EStateM.Result.error ex _ => Except.error ex
 
-@[inline] def runST {α : Type} (x : forall (σ : Type), ST σ α) : α :=
+@[noinline nospecialize]
+def runST {α : Type} (x : forall (σ : Type), ST σ α) : α :=
   match x Unit () with
   | EStateM.Result.ok a _     => a
   | EStateM.Result.error ex _ => nomatch ex
diff --git a/tests/lean/runSTBug.lean b/tests/lean/runSTBug.lean
new file mode 100644
index 0000000000..18d1b6acdf
--- /dev/null
+++ b/tests/lean/runSTBug.lean
@@ -0,0 +1,19 @@
+import Lean
+
+open Lean
+
+def f (xs : List Nat) (delta : Nat) : List Nat :=
+  runST (fun ω => visit xs |>.run)
+where
+  visit {ω} : List Nat → MonadCacheT Nat Nat (ST ω) (List Nat)
+    | []    => []
+    | a::as => do
+      let b ← checkCache a fun _ => return a + delta
+      return b :: (← visit as)
+
+def tst (xs : List Nat) : IO Unit := do
+  IO.println (f xs 10)
+  IO.println (f xs 20)
+  IO.println (f xs 30)
+
+#eval tst [1, 2, 3, 1, 2]
diff --git a/tests/lean/runSTBug.lean.expected.out b/tests/lean/runSTBug.lean.expected.out
new file mode 100644
index 0000000000..1468d8ea9a
--- /dev/null
+++ b/tests/lean/runSTBug.lean.expected.out
@@ -0,0 +1,3 @@
+[11, 12, 13, 11, 12]
+[21, 22, 23, 21, 22]
+[31, 32, 33, 31, 32]