From 67e9c83f54b44d72fb37840aa6bb8a45e6bb85a6 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 12 Oct 2020 10:38:32 -0700
Subject: [PATCH] fix: `for` result type

---
 src/Lean/Elab/Do.lean                      |  4 ++--
 tests/lean/run/forBodyResultTypeIssue.lean | 17 +++++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)
 create mode 100644 tests/lean/run/forBodyResultTypeIssue.lean

diff --git a/src/Lean/Elab/Do.lean b/src/Lean/Elab/Do.lean
index b07e743640..9219297de2 100644
--- a/src/Lean/Elab/Do.lean
+++ b/src/Lean/Elab/Do.lean
@@ -843,8 +843,8 @@ ctx ← read;
 u ← mkUVarTuple ref;
 match ctx.kind with
 | Kind.regular         => if ctx.uvars.isEmpty then pure action else `(HasBind.bind $action fun y => HasPure.pure (y, $u))
-| Kind.forIn           => `(HasBind.bind $action fun _ => HasPure.pure (ForInStep.yield $u))
-| Kind.forInWithReturn => `(HasBind.bind $action fun _ => HasPure.pure (ForInStep.yield (none, $u)))
+| Kind.forIn           => `(HasBind.bind $action fun (_ : PUnit) => HasPure.pure (ForInStep.yield $u))
+| Kind.forInWithReturn => `(HasBind.bind $action fun (_ : PUnit) => HasPure.pure (ForInStep.yield (none, $u)))
 | Kind.nestedBC        => unreachable!
 | Kind.nestedPR        => `(HasBind.bind $action fun y => (HasPure.pure (DoResultPR.«pure» y $u)))
 | Kind.nestedSBC       => `(HasBind.bind $action fun y => (HasPure.pure (DoResultSBC.«pureReturn» y $u)))
diff --git a/tests/lean/run/forBodyResultTypeIssue.lean b/tests/lean/run/forBodyResultTypeIssue.lean
new file mode 100644
index 0000000000..61accb9332
--- /dev/null
+++ b/tests/lean/run/forBodyResultTypeIssue.lean
@@ -0,0 +1,17 @@
+new_frontend
+
+abbrev M := ExceptT String $ StateT Nat Id
+
+def f (xs : List Nat) : M Unit := do
+for x in xs do
+  if x == 0 then
+    throw "contains zero"
+
+#eval f [1, 2, 3] $.run' 0
+#eval f [1, 0, 3] $.run' 0
+
+theorem ex1 : f [1, 2, 3] $.run' 0 = Except.ok () :=
+rfl
+
+theorem ex2 : f [1, 0, 3] $.run' 0 = Except.error "contains zero" :=
+rfl