From a516d384ae07fa38e31c8406b0cb4d23a7470f7f Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 2 Jan 2017 01:42:36 -0800
Subject: [PATCH] fix(library/compiler/erase_irrelevant): make sure io monad
 actions are not erased by dead code elimination optimization

---
 src/library/compiler/erase_irrelevant.cpp | 12 +++++-------
 tests/lean/io_bug1.lean                   | 22 ++++++++++++++++++++++
 tests/lean/io_bug1.lean.expected.out      | 22 ++++++++++++++++++++++
 3 files changed, 49 insertions(+), 7 deletions(-)
 create mode 100644 tests/lean/io_bug1.lean
 create mode 100644 tests/lean/io_bug1.lean.expected.out

diff --git a/src/library/compiler/erase_irrelevant.cpp b/src/library/compiler/erase_irrelevant.cpp
index 4944b3836d..4e8a2567ae 100644
--- a/src/library/compiler/erase_irrelevant.cpp
+++ b/src/library/compiler/erase_irrelevant.cpp
@@ -325,18 +325,16 @@ class erase_irrelevant_fn : public compiler_step_visitor {
               So, in this version, we have a simpler (monad.bind v b)
 
                bind v b :=
-               fun s, let a := v unit in
+               fun s, let a := v s in
                       b a a
 
-              We use (b a a) instead of (b a unit) to make sure the let-declaration
-              is not erased by the elim_unused_lets step. This is ok because
-              the second argument is not used during runtime.
-
-              We use a let-expression to make sure that `v unit` is not erased.
+              We use (b a a) instead of (b a unit), and (v s) instead of (v unit)
+              to make sure the let-declaration is not erased by the elim_unused_lets step.
+              For (b a a), this hack is ok because the second argument is not used during runtime.
             */
             expr v    = visit(args[4]);
             expr u    = mk_neutral_expr();
-            expr vu   = mk_app(v, u);
+            expr vu   = mk_app(v, mk_var(0));
             expr b    = visit(args[5]);
             expr bau  = beta_reduce(mk_app(b, mk_var(0), mk_var(0)));
             expr let  = mk_let("a", u, vu, bau);
diff --git a/tests/lean/io_bug1.lean b/tests/lean/io_bug1.lean
new file mode 100644
index 0000000000..27c90dc7ff
--- /dev/null
+++ b/tests/lean/io_bug1.lean
@@ -0,0 +1,22 @@
+import system.io
+
+def bar : io unit :=
+do put_str "one", put_str "two", put_str "three"
+
+vm_eval bar
+
+print "---------"
+
+def foo : ℕ → io unit
+| 0     := put_str "at zero\n"
+| (n+1) := do put_str "in\n", foo n, put_str "out\n"
+
+vm_eval foo 3
+
+print "---------"
+
+def foo2 : ℕ → io unit
+| 0     := put_str "at zero\n"
+| (n+1) := do put_str "in\n", foo2 n, put_str "out\n", put_str "out2\n"
+
+vm_eval foo2 3
diff --git a/tests/lean/io_bug1.lean.expected.out b/tests/lean/io_bug1.lean.expected.out
new file mode 100644
index 0000000000..a1207367d4
--- /dev/null
+++ b/tests/lean/io_bug1.lean.expected.out
@@ -0,0 +1,22 @@
+onetwothree
+---------
+in
+in
+in
+at zero
+out
+out
+out
+
+---------
+in
+in
+in
+at zero
+out
+out2
+out
+out2
+out
+out2
+