From cda4de474b6837a4601b764f224f481a57fc1055 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Tue, 6 Oct 2020 15:11:27 -0700
Subject: [PATCH] fix: weird error message in `do` notation

cc @Kha
---
 src/Lean/Elab/Do.lean                    | 2 +-
 tests/lean/doNotation1.lean              | 8 ++++++++
 tests/lean/doNotation1.lean.expected.out | 4 ++++
 3 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/Lean/Elab/Do.lean b/src/Lean/Elab/Do.lean
index 3b8d25f7e3..c02dfc3da8 100644
--- a/src/Lean/Elab/Do.lean
+++ b/src/Lean/Elab/Do.lean
@@ -205,7 +205,7 @@ partial def convertReturnIntoJmpAux (jp : Name) (xs : Array Name) : Code → Mac
 | Code.ite ref x? h c t e    => Code.ite ref x? h c <$> convertReturnIntoJmpAux t <*> convertReturnIntoJmpAux e
 | Code.«match» ref ds t alts => Code.«match» ref ds t <$> alts.mapM fun alt => do rhs ← convertReturnIntoJmpAux alt.rhs; pure { alt with rhs := rhs }
 | Code.returnAction e        => do c ← expandReturnAction e; convertReturnIntoJmpAux c
-| Code.«return» ref val      => pure $ Code.jmp ref jp ((xs.map $ mkIdentFrom ref).push val)
+| Code.«return» ref val      => do val ← `(ensureExpectedType! "type mismatch, returned value" $val); pure $ Code.jmp ref jp ((xs.map $ mkIdentFrom ref).push val)
 | c                          => pure c
 
 /- Convert `return _ x` instructions in `c` into `jmp _ jp xs`. -/
diff --git a/tests/lean/doNotation1.lean b/tests/lean/doNotation1.lean
index 698d5c65db..5b2774df8e 100644
--- a/tests/lean/doNotation1.lean
+++ b/tests/lean/doNotation1.lean
@@ -49,3 +49,11 @@ def f11 (x : Nat) : IO Unit := do
 if x > 0 then
   IO.println "x is not zero"
 IO.mkRef true -- error here as expected
+
+def f12 (x : Nat) : IO Unit := do
+if x > 0 then
+  return true
+else
+  x := x + 1
+  IO.println "hello" -- error here the other branch returns `Bool`
+IO.println x
diff --git a/tests/lean/doNotation1.lean.expected.out b/tests/lean/doNotation1.lean.expected.out
index be9ab12f52..76c4a2492a 100644
--- a/tests/lean/doNotation1.lean.expected.out
+++ b/tests/lean/doNotation1.lean.expected.out
@@ -27,3 +27,7 @@ has type
   EIO IO.Error (IO.Ref Bool)
 but is expected to have type
   EIO IO.Error Unit
+doNotation1.lean:58:2: error: type mismatch, returned value has type
+  Unit
+but is expected to have type
+  Bool