fix: improve error for result type mismatch in do blocks

This PR fixes #12846, where a `let x ← value` as the last do element in a block with non-Unit return type produced a confusing error on the value about `()` not having the right type. The error was misleading both in location (pointing at the value rather than the let statement) and in content (mentioning `PUnit.unit` which the user never wrote). The fix uses `withRef` to point the error at the enclosing statement and introduces `checkMonadicResultTypeMatches` to report the mismatch directly in terms of monadic result types.

src/Lean/Elab/BuiltinDo/Basic.lean

@@ -21,7 +21,8 @@ def elabDoIdDecl (x : Ident) (xType? : Option Term) (rhs : TSyntax `doElem) (k :
   let xType ← Term.elabType (xType?.getD (mkHole x))
   let lctx ← getLCtx
   let ctx ← read
-  elabDoElem rhs <| .mk (kind := kind) x.getId xType do
+  let ref ← getRef -- store the surrounding reference for error messages in `k`
+  elabDoElem rhs <| .mk (kind := kind) x.getId xType do withRef ref do
     withLCtxKeepingMutVarDefs lctx ctx x.getId do
       Term.addLocalVarInfo x (← getFVarFromUserName x.getId)
       k

src/Lean/Elab/Do/Basic.lean

@@ -416,14 +416,18 @@ def DoElemCont.mkBindUnlessPure (dec : DoElemCont) (e : Expr) : DoElabM Expr :=
     let k ← mkLambdaFVars #[xFVar] body
     mkBindApp eResultTy kResultTy e k
 
+def checkMonadicResultTypeMatches (resultType : Expr) (expectedType : Expr) : DoElabM Unit := do
+  unless ← isDefEqGuarded resultType expectedType do
+    throwError "Type mismatch. The rest of the `do` block has monadic result type{indentExpr resultType}\n\
+      but is expected to have monadic result type{indentExpr expectedType}"
+
 /--
 Return `let $k.resultName : PUnit := PUnit.unit; $(← k.k)`, ensuring that the result type of `k.k`
 is `PUnit` and then immediately zeta-reduce the `let`.
 -/
 def DoElemCont.continueWithUnit (dec : DoElemCont) : DoElabM Expr := do
-  let unit ← mkPUnitUnit
-  discard <| Term.ensureHasType dec.resultType unit
-  mapLetDeclZeta dec.resultName (← mkPUnit) unit (nondep := true) (kind := .ofBinderName dec.resultName) fun _ =>
+  checkMonadicResultTypeMatches dec.resultType (← mkPUnit)
+  mapLetDeclZeta dec.resultName (← mkPUnit) (← mkPUnitUnit) (nondep := true) (kind := .ofBinderName dec.resultName) fun _ =>
     dec.k
 
 /-- Elaborate the `DoElemCont` with the `deadCode` flag set to `deadSyntactically` to emit warnings. -/

tests/elab/issue12846.lean

@@ -0,0 +1,23 @@
+module
+
+set_option backward.do.legacy false
+
+/--
+error: Type mismatch. The rest of the `do` block has monadic result type
+  Bool
+but is expected to have monadic result type
+  Unit
+-/
+#guard_msgs in
+def test : IO Bool := do
+  let a ← pure 25
+
+/--
+error: Type mismatch. The rest of the `do` block has monadic result type
+  Bool
+but is expected to have monadic result type
+  Unit
+-/
+#guard_msgs in
+def test2 : IO Bool := do
+  let a := 25

tests/elab_fail/doNotation1.lean.out.expected

@@ -58,9 +58,7 @@ but is expected to have type
   List (Nat × Nat)
 in the application
   pure ()
-doNotation1.lean:82:0-83:9: error: Type mismatch
-  PUnit.unit
-has type
-  PUnit
-but is expected to have type
+doNotation1.lean:82:0-83:9: error: Type mismatch. The rest of the `do` block has monadic result type
   List (Nat × Nat)
+but is expected to have monadic result type
+  PUnit