fix: improve result type mismatch errors and locations in new do elaborator (#13404)

This PR fixes #12846, where the new do elaborator produced confusing
errors when a do element's continuation had a mismatched monadic result
type. The errors were misleading both in location (e.g., pointing at the
value of `let x ← value` rather than the `let` keyword) and in content
(e.g., mentioning `PUnit.unit` which the user never wrote).

The fix introduces `DoElemCont.ensureUnitAt`/`ensureHasTypeAt`, which
check the continuation result type early and report mismatches with a
clear message ("The `do` element has monadic result type ... but the
rest of the `do` block has monadic result type ..."). Each do-element
elaborator (`let`, `have`, `let rec`, `for`, `unless`, `dbg_trace`,
`assert!`, `idbg`, etc.) now captures its keyword token via `%$tk` and
passes it to `ensureUnitAt` so that the error points at the do element
rather than at an internal elaboration artifact. The old ad-hoc type
check in `for` and the confusing `ensureHasType` call in
`continueWithUnit` are replaced by this uniform mechanism. Additionally,
`extractMonadInfo` now calls `instantiateMVars` on the expected type,
and `While.lean`/`If.lean` macros propagate token info through their
expansions.

Closes #12846

---------

Co-authored-by: Rob23oba <robin.arnez@web.de>

src/Init/While.lean

@@ -35,21 +35,23 @@ instance [Monad m] : ForIn m Loop Unit where
 syntax "repeat " doSeq : doElem
 
 macro_rules
-  | `(doElem| repeat $seq) => `(doElem| for _ in Loop.mk do $seq)
+  | `(doElem| repeat%$tk $seq) => `(doElem| for%$tk _ in Loop.mk do $seq)
 
 syntax "while " ident " : " termBeforeDo " do " doSeq : doElem
 
 macro_rules
-  | `(doElem| while $h : $cond do $seq) => `(doElem| repeat if $h:ident : $cond then $seq else break)
+  | `(doElem| while%$tk $h : $cond do $seq) =>
+    `(doElem| repeat%$tk if $h:ident : $cond then $seq else break)
 
 syntax "while " termBeforeDo " do " doSeq : doElem
 
 macro_rules
-  | `(doElem| while $cond do $seq) => `(doElem| repeat if $cond then $seq else break)
+  | `(doElem| while%$tk $cond do $seq) => `(doElem| repeat%$tk if $cond then $seq else break)
 
 syntax "repeat " doSeq ppDedent(ppLine) "until " term : doElem
 
 macro_rules
-  | `(doElem| repeat $seq until $cond) => `(doElem| repeat do $seq:doSeq; if $cond then break)
+  | `(doElem| repeat%$tk $seq until $cond) =>
+    `(doElem| repeat%$tk do $seq:doSeq; if $cond then break)
 
 end Lean

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/BuiltinDo/For.lean

@@ -23,7 +23,7 @@ open Lean.Meta
   | `(doFor| for $[$_ : ]? $_:ident in $_ do $_) =>
     -- This is the target form of the expander, handled by `elabDoFor` below.
     Macro.throwUnsupported
-  | `(doFor| for $decls:doForDecl,* do $body) =>
+  | `(doFor| for%$tk $decls:doForDecl,* do $body) =>
     let decls := decls.getElems
     let `(doForDecl| $[$h? : ]? $pattern in $xs) := decls[0]! | Macro.throwUnsupported
     let mut doElems := #[]
@@ -74,12 +74,13 @@ open Lean.Meta
           | some ($y, s') =>
             $s:ident := s'
             do $body)
-    doElems := doElems.push (← `(doSeqItem| for $[$h? : ]? $x:ident in $xs do $body))
+    doElems := doElems.push (← `(doSeqItem| for%$tk $[$h? : ]? $x:ident in $xs do $body))
     `(doElem| do $doElems*)
   | _ => Macro.throwUnsupported
 
 @[builtin_doElem_elab Lean.Parser.Term.doFor] def elabDoFor : DoElab := fun stx dec => do
-  let `(doFor| for $[$h? : ]? $x:ident in $xs do $body) := stx | throwUnsupportedSyntax
+  let `(doFor| for%$tk $[$h? : ]? $x:ident in $xs do $body) := stx | throwUnsupportedSyntax
+  let dec ← dec.ensureUnitAt tk
   checkMutVarsForShadowing #[x]
   let uα ← mkFreshLevelMVar
   let uρ ← mkFreshLevelMVar
@@ -124,9 +125,6 @@ open Lean.Meta
       defs := defs.push (mkConst ``Unit.unit)
     return defs
 
-  unless ← isDefEq dec.resultType (← mkPUnit) do
-    logError m!"Type mismatch. `for` loops have result type {← mkPUnit}, but the rest of the `do` sequence expected {dec.resultType}."
-
   let (preS, σ) ← mkProdMkN (← useLoopMutVars none) mi.u
 
   let (app, p?) ← match h? with

src/Lean/Elab/BuiltinDo/If.lean

@@ -17,6 +17,7 @@ namespace Lean.Elab.Do
 open Lean.Parser.Term
 open Lean.Meta
 
+open InternalSyntax in
 /--
 If the given syntax is a `doIf`, return an equivalent `doIf` that has an `else` but no `else if`s or
 `if let`s.
@@ -25,8 +26,8 @@ If the given syntax is a `doIf`, return an equivalent `doIf` that has an `else`
   match stx with
   | `(doElem|if $_:doIfProp then $_ else $_) =>
     Macro.throwUnsupported
-  | `(doElem|if $cond:doIfCond then $t $[else if $conds:doIfCond then $ts]* $[else $e?]?) => do
-    let mut e : Syntax ← e?.getDM `(doSeq|pure PUnit.unit)
+  | `(doElem|if%$tk $cond:doIfCond then $t $[else if%$tks $conds:doIfCond then $ts]* $[else $e?]?) => do
+    let mut e : Syntax ← e?.getDM `(doSeq| skip%$tk)
     let mut eIsSeq := true
     for (cond, t) in Array.zip (conds.reverse.push cond) (ts.reverse.push t) do
       e ← if eIsSeq then pure e else `(doSeq|$(⟨e⟩):doElem)

src/Lean/Elab/BuiltinDo/Let.lean

@@ -88,17 +88,18 @@ private def checkLetConfigInDo (config : Term.LetConfig) : DoElabM Unit := do
     throwError "`+generalize` is not supported in `do` blocks"
 
 partial def elabDoLetOrReassign (config : Term.LetConfig) (letOrReassign : LetOrReassign) (decl : TSyntax ``letDecl)
-    (dec : DoElemCont) : DoElabM Expr := do
+    (tk : Syntax) (dec : DoElemCont) : DoElabM Expr := do
   checkLetConfigInDo config
   let vars ← getLetDeclVars decl
   letOrReassign.checkMutVars vars
+  let dec ← dec.ensureUnitAt tk
   -- Some decl preprocessing on the patterns and expected types:
   let decl ← pushTypeIntoReassignment letOrReassign decl
   let mγ ← mkMonadicType (← read).doBlockResultType
   match decl with
   | `(letDecl| $decl:letEqnsDecl) =>
     let declNew ← `(letDecl| $(⟨← liftMacroM <| Term.expandLetEqnsDecl decl⟩):letIdDecl)
-    return ← Term.withMacroExpansion decl declNew <| elabDoLetOrReassign config letOrReassign declNew dec
+    return ← Term.withMacroExpansion decl declNew <| elabDoLetOrReassign config letOrReassign declNew tk dec
   | `(letDecl| $pattern:term $[: $xType?]? := $rhs) =>
     let rhs ← match xType? with | some xType => `(($rhs : $xType)) | none => pure rhs
     let contElab : DoElabM Expr := elabWithReassignments letOrReassign vars dec.continueWithUnit
@@ -162,10 +163,11 @@ partial def elabDoLetOrReassign (config : Term.LetConfig) (letOrReassign : LetOr
             mkLetFVars #[x, h'] body (usedLetOnly := config.usedOnly) (generalizeNondepLet := false)
   | _ => throwUnsupportedSyntax
 
-def elabDoArrow (letOrReassign : LetOrReassign) (stx : TSyntax [``doIdDecl, ``doPatDecl]) (dec : DoElemCont) : DoElabM Expr := do
+def elabDoArrow (letOrReassign : LetOrReassign) (stx : TSyntax [``doIdDecl, ``doPatDecl]) (tk : Syntax) (dec : DoElemCont) : DoElabM Expr := do
   match stx with
   | `(doIdDecl| $x:ident $[: $xType?]? ← $rhs) =>
     letOrReassign.checkMutVars #[x]
+    let dec ← dec.ensureUnitAt tk
     -- For plain variable reassignment, we know the expected type of the reassigned variable and
     -- propagate it eagerly via type ascription if the user hasn't provided one themselves:
     let xType? ← match letOrReassign, xType? with
@@ -177,6 +179,7 @@ def elabDoArrow (letOrReassign : LetOrReassign) (stx : TSyntax [``doIdDecl, ``do
       (kind := dec.kind)
   | `(doPatDecl| _%$pattern $[: $patType?]? ← $rhs) =>
     let x := mkIdentFrom pattern (← mkFreshUserName `__x)
+    let dec ← dec.ensureUnitAt tk
     elabDoIdDecl x patType? rhs dec.continueWithUnit (kind := dec.kind)
   | `(doPatDecl| $pattern:term $[: $patType?]? ← $rhs $[| $otherwise? $(rest?)?]?) =>
     let rest? := rest?.join
@@ -205,17 +208,18 @@ private def getLetConfigAndCheckMut (letConfigStx : TSyntax ``Parser.Term.letCon
   Term.mkLetConfig letConfigStx initConfig
 
 @[builtin_doElem_elab Lean.Parser.Term.doLet] def elabDoLet : DoElab := fun stx dec => do
-  let `(doLet| let $[mut%$mutTk?]? $config:letConfig $decl:letDecl) := stx | throwUnsupportedSyntax
+  let `(doLet| let%$tk $[mut%$mutTk?]? $config:letConfig $decl:letDecl) := stx | throwUnsupportedSyntax
   let config ← getLetConfigAndCheckMut config mutTk?
-  elabDoLetOrReassign config (.let mutTk?) decl dec
+  elabDoLetOrReassign config (.let mutTk?) decl tk dec
 
 @[builtin_doElem_elab Lean.Parser.Term.doHave] def elabDoHave : DoElab := fun stx dec => do
-  let `(doHave| have $config:letConfig $decl:letDecl) := stx | throwUnsupportedSyntax
+  let `(doHave| have%$tk $config:letConfig $decl:letDecl) := stx | throwUnsupportedSyntax
   let config ← Term.mkLetConfig config { nondep := true }
-  elabDoLetOrReassign config .have decl dec
+  elabDoLetOrReassign config .have decl tk dec
 
 @[builtin_doElem_elab Lean.Parser.Term.doLetRec] def elabDoLetRec : DoElab := fun stx dec => do
-  let `(doLetRec| let rec $decls:letRecDecls) := stx | throwUnsupportedSyntax
+  let `(doLetRec| let%$tk rec $decls:letRecDecls) := stx | throwUnsupportedSyntax
+  let dec ← dec.ensureUnitAt tk
   let vars ← getLetRecDeclsVars decls
   let mγ ← mkMonadicType (← read).doBlockResultType
   doElabToSyntax m!"let rec body of group {vars}" dec.continueWithUnit fun body => do
@@ -227,13 +231,13 @@ private def getLetConfigAndCheckMut (letConfigStx : TSyntax ``Parser.Term.letCon
 @[builtin_doElem_elab Lean.Parser.Term.doReassign] def elabDoReassign : DoElab := fun stx dec => do
   -- def doReassign := letIdDeclNoBinders <|> letPatDecl
   match stx with
-  | `(doReassign| $x:ident $[: $xType?]? := $rhs) =>
+  | `(doReassign| $x:ident $[: $xType?]? :=%$tk $rhs) =>
     let decl : TSyntax ``letIdDecl ← `(letIdDecl| $x:ident $[: $xType?]? := $rhs)
     let decl : TSyntax ``letDecl := ⟨mkNode ``letDecl #[decl]⟩
-    elabDoLetOrReassign {} .reassign decl dec
+    elabDoLetOrReassign {} .reassign decl tk dec
   | `(doReassign| $decl:letPatDecl) =>
     let decl : TSyntax ``letDecl := ⟨mkNode ``letDecl #[decl]⟩
-    elabDoLetOrReassign {} .reassign decl dec
+    elabDoLetOrReassign {} .reassign decl decl dec
   | _ => throwUnsupportedSyntax
 
 @[builtin_doElem_elab Lean.Parser.Term.doLetElse] def elabDoLetElse : DoElab := fun stx dec => do
@@ -255,17 +259,17 @@ private def getLetConfigAndCheckMut (letConfigStx : TSyntax ``Parser.Term.letCon
     elabDoElem (← `(doElem| match $rhs:term with | $pattern => $body:doSeqIndent | _ => $otherwise:doSeqIndent)) dec
 
 @[builtin_doElem_elab Lean.Parser.Term.doLetArrow] def elabDoLetArrow : DoElab := fun stx dec => do
-  let `(doLetArrow| let $[mut%$mutTk?]? $cfg:letConfig $decl) := stx | throwUnsupportedSyntax
+  let `(doLetArrow| let%$tk $[mut%$mutTk?]? $cfg:letConfig $decl) := stx | throwUnsupportedSyntax
   let config ← getLetConfigAndCheckMut cfg mutTk?
   checkLetConfigInDo config
   if config.nondep || config.usedOnly || config.zeta || config.eq?.isSome then
     throwErrorAt cfg "configuration options are not supported with `←`"
-  elabDoArrow (.let mutTk?) decl dec
+  elabDoArrow (.let mutTk?) decl tk dec
 
 @[builtin_doElem_elab Lean.Parser.Term.doReassignArrow] def elabDoReassignArrow : DoElab := fun stx dec => do
   match stx with
   | `(doReassignArrow| $decl:doIdDecl) =>
-    elabDoArrow .reassign decl dec
+    elabDoArrow .reassign decl decl dec
   | `(doReassignArrow| $decl:doPatDecl) =>
-    elabDoArrow .reassign decl dec
+    elabDoArrow .reassign decl decl dec
   | _ => throwUnsupportedSyntax

src/Lean/Elab/BuiltinDo/Misc.lean

@@ -16,6 +16,12 @@ namespace Lean.Elab.Do
 open Lean.Parser.Term
 open Lean.Meta
 
+open InternalSyntax in
+@[builtin_doElem_elab Lean.Parser.Term.InternalSyntax.doSkip] def elabDoSkip : DoElab := fun stx dec => do
+  let `(doSkip| skip%$tk) := stx | throwUnsupportedSyntax
+  let dec ← dec.ensureUnitAt tk
+  dec.continueWithUnit
+
 @[builtin_doElem_elab Lean.Parser.Term.doExpr] def elabDoExpr : DoElab := fun stx dec => do
   let `(doExpr| $e:term) := stx | throwUnsupportedSyntax
   let mα ← mkMonadicType dec.resultType
@@ -26,24 +32,28 @@ open Lean.Meta
   let `(doNested| do $doSeq) := stx | throwUnsupportedSyntax
   elabDoSeq ⟨doSeq.raw⟩ dec
 
+open InternalSyntax in
 @[builtin_doElem_elab Lean.Parser.Term.doUnless] def elabDoUnless : DoElab := fun stx dec => do
-  let `(doUnless| unless $cond do $body) := stx | throwUnsupportedSyntax
-  elabDoElem (← `(doElem| if $cond then pure PUnit.unit else $body)) dec
+  let `(doUnless| unless%$tk $cond do $body) := stx | throwUnsupportedSyntax
+  elabDoElem (← `(doElem| if $cond then skip%$tk else $body)) dec
 
 @[builtin_doElem_elab Lean.Parser.Term.doDbgTrace] def elabDoDbgTrace : DoElab := fun stx dec => do
-  let `(doDbgTrace| dbg_trace $msg:term) := stx | throwUnsupportedSyntax
+  let `(doDbgTrace| dbg_trace%$tk $msg:term) := stx | throwUnsupportedSyntax
   let mγ ← mkMonadicType (← read).doBlockResultType
+  let dec ← dec.ensureUnitAt tk
   doElabToSyntax "dbg_trace body" dec.continueWithUnit fun body => do
   Term.elabTerm (← `(dbg_trace $msg; $body)) mγ
 
 @[builtin_doElem_elab Lean.Parser.Term.doAssert] def elabDoAssert : DoElab := fun stx dec => do
-  let `(doAssert| assert! $cond) := stx | throwUnsupportedSyntax
+  let `(doAssert| assert!%$tk $cond) := stx | throwUnsupportedSyntax
   let mγ ← mkMonadicType (← read).doBlockResultType
+  let dec ← dec.ensureUnitAt tk
   doElabToSyntax "assert! body" dec.continueWithUnit fun body => do
   Term.elabTerm (← `(assert! $cond; $body)) mγ
 
 @[builtin_doElem_elab Lean.Parser.Term.doDebugAssert] def elabDoDebugAssert : DoElab := fun stx dec => do
-  let `(doDebugAssert| debug_assert! $cond) := stx | throwUnsupportedSyntax
+  let `(doDebugAssert| debug_assert!%$tk $cond) := stx | throwUnsupportedSyntax
   let mγ ← mkMonadicType (← read).doBlockResultType
+  let dec ← dec.ensureUnitAt tk
   doElabToSyntax "debug_assert! body" dec.continueWithUnit fun body => do
   Term.elabTerm (← `(debug_assert! $cond; $body)) mγ

src/Lean/Elab/Do/Basic.lean

@@ -374,14 +374,60 @@ def withLCtxKeepingMutVarDefs (oldLCtx : LocalContext) (oldCtx : Context) (resul
     mutVarDefs := oldMutVarDefs
   }) k
 
+def mkMonadicResultTypeMismatchError (contType : Expr) (elementType : Expr) : MessageData :=
+  m!"Type mismatch. The `do` element has monadic result type{indentExpr elementType}\n\
+    but the rest of the `do` block has monadic result type{indentExpr contType}"
+
+/--
+Given a continuation `dec`, a reference `ref`, and an element result type `elementType`, returns a
+continuation derived from `dec` with result type `elementType`.
+If `dec` already has result type `elementType`, simply returns `dec`.
+Otherwise, an error is logged and a new continuation is returned that calls `dec` with `sorry` as a
+result. The error is reported at `ref`.
+-/
+def DoElemCont.ensureHasTypeAt (dec : DoElemCont) (ref : Syntax) (elementType : Expr) : DoElabM DoElemCont := do
+  if ← isDefEqGuarded dec.resultType elementType then
+    return dec
+  let errMessage := mkMonadicResultTypeMismatchError dec.resultType elementType
+  unless (← readThe Term.Context).errToSorry do
+    throwErrorAt ref errMessage
+  logErrorAt ref errMessage
+  return {
+    resultName := ← mkFreshUserName `__r
+    resultType := elementType
+    k := do
+      mapLetDecl dec.resultName dec.resultType (← mkSorry dec.resultType true)
+          (nondep := true) (kind := .implDetail) fun _ => dec.k
+    kind := dec.kind
+  }
+
+/--
+Given a continuation `dec` and a reference `ref`, returns a continuation derived from `dec` with result type `PUnit`.
+If `dec` already has result type `PUnit`, simply returns `dec`. Otherwise, an error is logged and a
+new continuation is returned that calls `dec` with `sorry` as a result. The error is reported at `ref`.
+-/
+def DoElemCont.ensureUnitAt (dec : DoElemCont) (ref : Syntax) : DoElabM DoElemCont := do
+  dec.ensureHasTypeAt ref (← mkPUnit)
+
+/--
+Given a continuation `dec`, returns a continuation derived from `dec` with result type `PUnit`.
+If `dec` already has result type `PUnit`, simply returns `dec`. Otherwise, an error is logged and a
+new continuation is returned that calls `dec` with `sorry` as a result.
+-/
+def DoElemCont.ensureUnit (dec : DoElemCont) : DoElabM DoElemCont := do
+  dec.ensureUnitAt (← getRef)
+
 /--
 Return `$e >>= fun ($dec.resultName : $dec.resultType) => $(← dec.k)`, cancelling
 the bind if `$(← dec.k)` is `pure $dec.resultName` or `e` is some `pure` computation.
 -/
 def DoElemCont.mkBindUnlessPure (dec : DoElemCont) (e : Expr) : DoElabM Expr := do
+  -- let eResultTy ← mkFreshResultType
+  -- let e ← Term.ensureHasType (← mkMonadicType eResultTy) e
+  -- let dec ← dec.ensureHasType eResultTy
   let x := dec.resultName
-  let eResultTy := dec.resultType
   let k := dec.k
+  let eResultTy := dec.resultType
   -- The .ofBinderName below is mainly to interpret `__do_lift` binders as implementation details.
   let declKind := .ofBinderName x
   let kResultTy ← mkFreshResultType `kResultTy
@@ -421,9 +467,8 @@ Return `let $k.resultName : PUnit := PUnit.unit; $(← k.k)`, ensuring that the
 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 _ =>
+  let dec ← dec.ensureUnit
+  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. -/
@@ -604,6 +649,7 @@ def enterFinally (resultType : Expr) (k : DoElabM Expr) : DoElabM Expr := do
 /-- Extracts `MonadInfo` and monadic result type `α` from the expected type of a `do` block `m α`. -/
 private partial def extractMonadInfo (expectedType? : Option Expr) : Term.TermElabM (MonadInfo × Expr) := do
   let some expectedType := expectedType? | mkUnknownMonadResult
+  let expectedType ← instantiateMVars expectedType
   let extractStep? (type : Expr) : Term.TermElabM (Option (MonadInfo × Expr)) := do
     let .app m resultType := type.consumeMData | return none
     unless ← isType resultType do return none

src/Lean/Elab/Do/InferControlInfo.lean

@@ -79,6 +79,7 @@ builtin_initialize controlInfoElemAttribute : KeyedDeclsAttribute ControlInfoHan
 
 namespace InferControlInfo
 
+open InternalSyntax in
 mutual
 
 partial def ofElem (stx : TSyntax `doElem) : TermElabM ControlInfo := do
@@ -152,6 +153,7 @@ partial def ofElem (stx : TSyntax `doElem) : TermElabM ControlInfo := do
     let finInfo ← ofOptionSeq finSeq?
     return info.sequence finInfo
   -- Misc
+  | `(doElem| skip) => return .pure
   | `(doElem| dbg_trace $_) => return .pure
   | `(doElem| assert! $_) => return .pure
   | `(doElem| debug_assert! $_) => return .pure

src/Lean/Elab/Do/Legacy.lean

@@ -1815,6 +1815,13 @@ mutual
               return mkTerminalAction term
             else
               return mkSeq term (← doSeqToCode doElems)
+          else if k == ``Parser.Term.InternalSyntax.doSkip then
+            -- In the legacy elaborator, `skip` is treated as `pure PUnit.unit`.
+            let term ← withRef doElem `(pure PUnit.unit)
+            if doElems.isEmpty then
+              return mkTerminalAction term
+            else
+              return mkSeq term (← doSeqToCode doElems)
           else
             throwError "unexpected do-element of kind {doElem.getKind}:\n{doElem}"
 end

src/Lean/Elab/Idbg.lean

@@ -364,8 +364,9 @@ def elabIdbgTerm : TermElab := fun stx expectedType? => do
 
 @[builtin_doElem_elab Lean.Parser.Term.doIdbg]
 def elabDoIdbg : DoElab := fun stx dec => do
-  let `(Lean.Parser.Term.doIdbg| idbg $e) := stx | throwUnsupportedSyntax
+  let `(Lean.Parser.Term.doIdbg| idbg%$tk $e) := stx | throwUnsupportedSyntax
   let mγ ← mkMonadicType (← read).doBlockResultType
+  let dec ← dec.ensureUnitAt tk
   doElabToSyntax "idbg body" dec.continueWithUnit fun body => do
     elabIdbgCore (e := e) (body := body) (ref := stx) mγ
 

tests/elab/issue12846.lean

@@ -0,0 +1,108 @@
+module
+
+set_option backward.do.legacy false
+
+-- Original issue: `let x ← value` as last element in non-Unit do block
+/--
+error: Type mismatch. The `do` element has monadic result type
+  Unit
+but the rest of the `do` block has monadic result type
+  Bool
+-/
+#guard_msgs in
+def test_letArrow : IO Bool := do
+  let a ← pure 25
+
+/--
+error: Type mismatch. The `do` element has monadic result type
+  Unit
+but the rest of the `do` block has monadic result type
+  Bool
+-/
+#guard_msgs in
+def test_let : IO Bool := do
+  let a := 25
+
+-- `have` as last element
+/--
+error: Type mismatch. The `do` element has monadic result type
+  Unit
+but the rest of the `do` block has monadic result type
+  Bool
+-/
+#guard_msgs in
+def test_have : IO Bool := do
+  have a := 25
+
+-- `let rec` as last element
+/--
+error: Type mismatch. The `do` element has monadic result type
+  Unit
+but the rest of the `do` block has monadic result type
+  Bool
+-/
+#guard_msgs in
+def test_letRec : IO Bool := do
+  let rec f : Nat → Nat
+    | 0 => 0
+    | n + 1 => f n
+
+-- `for` as last element
+/--
+error: Type mismatch. The `do` element has monadic result type
+  Unit
+but the rest of the `do` block has monadic result type
+  Bool
+-/
+#guard_msgs in
+def test_for : IO Bool := do
+  for _ in [1, 2, 3] do
+    pure ()
+
+-- `dbg_trace` as last element
+/--
+error: Type mismatch. The `do` element has monadic result type
+  Unit
+but the rest of the `do` block has monadic result type
+  Bool
+-/
+#guard_msgs in
+def test_dbgTrace : IO Bool := do
+  dbg_trace "hello"
+
+-- `assert!` as last element
+/--
+error: Type mismatch. The `do` element has monadic result type
+  Unit
+but the rest of the `do` block has monadic result type
+  Bool
+-/
+#guard_msgs in
+def test_assert : IO Bool := do
+  assert! true
+
+-- `if` without else as last element
+/--
+error: Application type mismatch: The argument
+  ()
+has type
+  Unit
+but is expected to have type
+  Bool
+in the application
+  pure ()
+---
+error: Type mismatch. The `do` element has monadic result type
+  Unit
+but the rest of the `do` block has monadic result type
+  Bool
+-/
+#guard_msgs in
+def test_if_no_else : IO Bool := do
+  if true then
+    pure ()
+
+-- `if` with else works fine when branches match the result type
+#guard_msgs in
+def test_if_else_ok : IO Bool := do
+  if true then pure true else pure false

tests/elab/newdo.lean

@@ -26,11 +26,7 @@ Many of these are extracted from our code base.
     x := x + 1
   return ⟨3, by decide⟩
 
--- Regression test cases of what's broken in the legacy do elaborator:
-example : Unit := (Id.run do  let n ← if true then pure 3 else pure 42)
 example : Unit := (Id.run do let n ← if true then pure 3 else pure 42)
-example := (Id.run do  let mut x := 0; x ← return 10)
-example := (Id.run do let mut x := 0; x ← return 10)
 
 -- Another complicated `match` that would need to generalize the join point type if it was dependent
 example (x : Nat) : Id (Fin (x + 2)) := do
@@ -211,8 +207,8 @@ trace: [Elab.do] let x := 42;
               else
                 let x := x + i;
                 pure (ForInStep.yield (none, x))
-    let __r : Option ?m.185 := __s.fst
-    let x : ?m.185 := __s.snd
+    let __r : Option ?m.170 := __s.fst
+    let x : ?m.170 := __s.snd
     match __r with
       | some r => pure r
       | none =>

tests/elab_fail/doNotation1.lean.out.expected

@@ -49,18 +49,7 @@ doNotation1.lean:78:21-78:31: error: typeclass instance problem is stuck
 Note: Lean will not try to resolve this typeclass instance problem because the type argument to `ToString` is a metavariable. This argument must be fully determined before Lean will try to resolve the typeclass.
 
 Hint: Adding type annotations and supplying implicit arguments to functions can give Lean more information for typeclass resolution. For example, if you have a variable `x` that you intend to be a `Nat`, but Lean reports it as having an unresolved type like `?m`, replacing `x` with `(x : Nat)` can get typeclass resolution un-stuck.
-doNotation1.lean:82:0-83:9: error: Type mismatch. `for` loops have result type Unit, but the rest of the `do` sequence expected List (Nat × Nat).
-doNotation1.lean:83:7-83:9: error: Application type mismatch: The argument
-  ()
-has type
+doNotation1.lean:82:0-82:3: error: Type mismatch. The `do` element has monadic result type
   Unit
-but is expected to have type
-  List (Nat × Nat)
-in the application
-  pure ()
-doNotation1.lean:82:0-83:9: error: Type mismatch
-  ()
-has type
-  Unit
-but is expected to have type
+but the rest of the `do` block has monadic result type
   List (Nat × Nat)