fix: weird error message in do notation

cc @Kha

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`. -/

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

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