feat: add do← for effect forwarding (#13931)

This PR introduces the `do← body` marker (ASCII `do<- body`), which lets
ordinary continuation-taking wrappers like `withReader` or
`Meta.withLocalDecl` participate in the surrounding `do` block's control
flow. When `do← body` appears as the last argument of an application
inside a `do` block, the body's `return`, `break`, `continue`, and
`mut`-variable reassignments are forwarded out through the wrapper to
the enclosing block.

```lean
example : ReaderT Ctx IO Nat := do
  let mut x := 0
  withReader (·.bump) (do←
    if cond then return 0     -- early-returns from the outer `do`
    x := x + 1)               -- mutates the outer `x`
  return x
```

The syntax is reminiscent of a nested action `(← body)`, but unlike a
nested action, `body` is not run eagerly in the `do` block context
before the wrapping function is called. The wrapping function decides
when to run `body`, and code is inserted to forward `body`'s effects to
the outer `do` block.

Internally, `elabDoExpr` and the `ControlInfo` inferer recognise the
marker and route the body through the same `EffectForwarder` framework
that already powers `try`/`catch`. The framework now lets each caller
decide whether the surrounding continuation is dead: `try`/`catch`
propagates `info.noFallthrough`, while `do←` always treats the
continuation as live (an opaque wrapper might not invoke the body at
all). A new `DoElemCont.withDeadCodeFromInfo` helper captures the shared
pattern. A wrapper validator rejects shapes that don't fit `(… → m α) →
m α` with a polymorphic `α` (e.g. `withConst`, `foldlM`) up front with a
dedicated diagnostic, rather than letting a downstream type mismatch
surface. `dropParens` moves from `Lean.Elab.Term` to `Lean.Parser.Term`
since it is a generic term-syntax utility.

Author: sgraf812

src/Lean/Elab/BuiltinDo.lean

@@ -15,4 +15,5 @@ public import Lean.Elab.BuiltinDo.Jump
 public import Lean.Elab.BuiltinDo.Misc
 public import Lean.Elab.BuiltinDo.For
 public import Lean.Elab.BuiltinDo.TryCatch
+public import Lean.Elab.BuiltinDo.Forward
 public import Lean.Elab.BuiltinDo.Repeat

src/Lean/Elab/BuiltinDo/Forward.lean

@@ -0,0 +1,92 @@
+/-
+Copyright (c) 2026 Lean FRO LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Sebastian Graf
+-/
+module
+
+prelude
+public import Lean.Elab.Do.Basic
+meta import Lean.Parser.Do
+import Lean.Elab.Do.Control
+import Lean.Elab.Do.InferControlInfo
+import Lean.Elab.Binders
+
+namespace Lean.Elab.Do
+
+open Lean.Parser.Term
+open Lean.Meta
+
+@[builtin_term_elab Lean.Parser.Term.doForward]
+public def elabDoForward : Term.TermElab := fun _ _ =>
+  throwError "\
+    `do←` may only appear as the last argument of a function application \
+    inside an enclosing `do` block, optionally inside a `fun` binder"
+
+private def forwardHint (headApp : Term) (reason : MessageData) : MessageData :=
+  m!"\
+    `{headApp}` is not a valid `do←` wrapper: {reason}. The wrapper must have type \
+    `(… → m α) → m α` for some `α` that is universally quantified in the wrapper's signature and \
+    does not appear elsewhere."
+
+/--
+Check that `probeExpr` (the elaboration of `headApp ?forwarded`) has the shape required for
+`do←`: its result type is `m α` for an unassigned `α`, the `forwarded` slot is `… → m α` with
+the same `α` (and `α` not occurring in the slot's input types), and `α` does not occur in any
+*other* explicit argument of the elaborated probe.
+-/
+private def validateForwarder (headApp : Term) (forwarded probeExpr : Expr) : MetaM Unit := do
+  let reject {α} (reason : MessageData) : MetaM α := throwError forwardHint headApp reason
+  let probeType ← whnfD (← instantiateMVars (← inferType probeExpr))
+  let .app _ alphaRet := probeType
+    | reject m!"its return type `{probeType}` is not of the form `m α`"
+  unless alphaRet.isMVar do
+    reject m!"its return type pins `α` to a concrete type"
+  let argMentionsAlpha (arg : Expr) : MetaM (Option Expr) := do
+    let ty ← inferType arg
+    if (← instantiateMVars ty).findMVar? (· == alphaRet.mvarId!) |>.isSome then
+      return some ty
+    else return none
+  forallTelescopeReducing (← inferType forwarded) fun args body => do
+    unless ← isDefEq probeType (← whnfD (← instantiateMVars body)) do
+      reject m!"the forwarded body's `α` differs from the wrapper's return `α`"
+    for arg in args do
+      if let some ty ← argMentionsAlpha arg then
+        reject m!"`α` appears in the forwarded body's input type `{ty}`"
+  forallTelescopeReducing (← inferType probeExpr.getAppFn) fun fvars _ => do
+    for (fvar, appArg) in fvars.zip probeExpr.getAppArgs do
+      unless (← fvar.fvarId!.getDecl).binderInfo matches .default do continue
+      let appArgInst ← instantiateMVars appArg
+      if appArgInst == forwarded then continue
+      if let some ty ← argMentionsAlpha appArgInst then
+        reject m!"`α` appears in an applied explicit argument of type `{ty}`"
+
+/--
+Elaborate `e` as an application whose last argument is a `do←` body, performing the effect
+forwarding. Returns `none` if `e` doesn't have that shape.
+
+Strategy: build `headApp ?forwarded` with a fresh metavariable in the body slot and elaborate
+it. Lean's elaborator handles named args, defaults, and dot notation. We then check the
+wrapper's shape against the probe; on success, the lifted body is assigned into `?forwarded`
+and the probe expression itself is the result.
+-/
+public def tryElabForwardApp? (e : Term) (dec : DoElemCont) : DoElabM (Option Expr) := do
+  let some (headApp, arg) := Forward.matchApp? e | return none
+  let (forwarded, probeExpr) ← controlAtTermElabM fun _ => withFreshMacroScope do
+    let forwardedStx ← `(?forwarded)
+    let forwarded ← Term.elabTerm forwardedStx none
+    let probeExpr ← Term.elabTerm (← `($headApp $forwardedStx)) none
+    pure (forwarded, probeExpr)
+  validateForwarder headApp forwarded probeExpr
+  let bodyInfo ← InferControlInfo.ofSeq arg.body
+  let forwarder ← EffectForwarder.ofCont bodyInfo dec
+  let liftedBody ← controlAtTermElabM fun runInBase => do
+    Term.elabFunBinders (arg.binders.map (·.raw)) none fun bsExpr _ => runInBase do
+      mkLambdaFVars bsExpr (← forwarder.lift (elabDoSeq arg.body))
+  unless ← forwarded.mvarId!.checkedAssign liftedBody do
+    throwError "the lifted body's type does not match the wrapper's body slot type"
+  -- NB: No `withDeadCode` wrap on `restoreCont`'s result, because it is unclear whether `body` is
+  --     even called.
+  return some (← (← forwarder.restoreCont).mkBindUnlessPure probeExpr)
+
+end Lean.Elab.Do

src/Lean/Elab/BuiltinDo/Jump.lean

@@ -19,7 +19,7 @@ open Lean.Meta
 @[builtin_doElem_elab Lean.Parser.Term.doReturn] def elabDoReturn : DoElab := fun stx dec => do
   let `(doReturn| return $[$e?]?) := stx | throwUnsupportedSyntax
   let returnCont ← getReturnCont
-  -- When using the ControlLifter framework, `returnCont.resultType` can be different than the
+  -- When using the EffectForwarder framework, `returnCont.resultType` can be different than the
   -- result type of the `do` block. That's why we track it separately.
   -- trace[Elab.do] "return e: {e} with type {returnCont.resultType}"
   let e ← match e? with

src/Lean/Elab/BuiltinDo/Misc.lean

@@ -7,6 +7,7 @@ module
 
 prelude
 public import Lean.Elab.Do.Basic
+public import Lean.Elab.BuiltinDo.Forward
 meta import Lean.Parser.Do
 
 public section
@@ -24,6 +25,8 @@ open InternalSyntax in
 
 @[builtin_doElem_elab Lean.Parser.Term.doExpr] def elabDoExpr : DoElab := fun stx dec => do
   let `(doExpr| $e:term) := stx | throwUnsupportedSyntax
+  if let some r ← tryElabForwardApp? e dec then
+    return r
   let mα ← mkMonadApp dec.resultType
   let e ← Term.elabTermEnsuringType e mα
   dec.mkBindUnlessPure e

src/Lean/Elab/BuiltinDo/TryCatch.lean

@@ -17,7 +17,7 @@ namespace Lean.Elab.Do
 open Lean.Parser.Term
 open Lean.Meta
 
-private def elabDoCatch (lifter : ControlLifter) (body : Expr) (catch_ : TSyntax ``doCatch) : DoElabM Expr := do
+private def elabDoCatch (lifter : EffectForwarder) (body : Expr) (catch_ : TSyntax ``doCatch) : DoElabM Expr := do
   let mi := (← read).monadInfo
   let `(doCatch| catch $x $[: $eType?]? => $catchSeq) := catch_ | throwUnsupportedSyntax
   let x := Term.mkExplicitBinder x <| match eType? with
@@ -56,7 +56,8 @@ private def elabDoCatch (lifter : ControlLifter) (body : Expr) (catch_ : TSyntax
   -- So we need to pack up our effects and unpack them after the `try`.
   -- We could optimize for the terminal action case by omitting the state tuple ... in the future.
   let mi := (← read).monadInfo
-  let lifter ← ControlLifter.ofCont (← inferControlInfoElem stx) dec
+  let info ← inferControlInfoElem stx
+  let lifter ← EffectForwarder.ofCont info dec
   let body ← do
     let body ← lifter.lift (elabDoSeq trySeq)
     let body ← catches.foldlM (init := body) fun body catch_ => do
@@ -75,4 +76,6 @@ private def elabDoCatch (lifter : ControlLifter) (body : Expr) (catch_ : TSyntax
         let instFunctor ← Term.mkInstMVar <| mkApp (mkConst ``Functor [mi.u, mi.v]) mi.m
         pure <| mkApp7 (mkConst ``tryFinally [mi.u, mi.v])
           mi.m lifter.liftedDoBlockResultType β instMonadFinally instFunctor body fin
-  (← lifter.restoreCont).mkBindUnlessPure body
+  -- `info` already combines try and catch arms via `ControlInfo.alternative`,
+  -- so its `noFallthrough` is the right liveness signal for the continuation.
+  ((← lifter.restoreCont).withDeadCodeFromInfo info).mkBindUnlessPure body

src/Lean/Elab/Do/Basic.lean

@@ -542,6 +542,10 @@ def DoElemCont.elabAsSyntacticallyDeadCode (dec : DoElemCont) : DoElabM Unit :=
     s.restore
     Core.setMessageLog (log ++ warnings)
 
+/-- Wrap `dec.k` so it elaborates as dead iff `info.noFallthrough`. -/
+def DoElemCont.withDeadCodeFromInfo (dec : DoElemCont) (info : ControlInfo) : DoElemCont :=
+  { dec with k := withDeadCode (if info.noFallthrough then .deadSemantically else .alive) dec.k }
+
 /--
 Given a list of mut vars `vars` and an FVar `tupleVar` binding a tuple, bind the mut vars to the
 fields of the tuple and call `k` in the resulting local context.
@@ -646,8 +650,7 @@ def DoElemCont.withDuplicableCont (nondupDec : DoElemCont) (callerInfo : Control
     withLocalDeclD nondupDec.resultName nondupDec.resultType fun r => do
     withLocalDeclsDND (mutDecls.map fun (d : LocalDecl) => (d.userName, d.type)) fun muts => do
     for (x, newX) in mutVars.zip muts do Term.addTermInfo' x newX
-    withDeadCode (if callerInfo.noFallthrough then .deadSemantically else .alive) do
-    let e ← nondupDec.k
+    let e ← (nondupDec.withDeadCodeFromInfo callerInfo).k
     mkLambdaFVars (#[r] ++ muts) e
   unless ← joinRhsMVar.mvarId!.checkedAssign joinRhs do
     joinRhsMVar.mvarId!.withContext do

src/Lean/Elab/Do/Control.lean

@@ -186,20 +186,27 @@ def ControlStack.mkPure (base : ControlStack) (resultName : Name) : DoElabM Expr
   let r ← getFVarFromUserName resultName
   base.runInBase <| mkApp4 (mkConst ``Pure.pure [mi.u, mi.v]) mi.m instPure (← inferType r) r
 
-structure ControlLifter where
+/-- Plan for embedding a non-tail-resumptive body inside `origCont` via a monad-transformer stack. -/
+structure EffectForwarder where
+  /-- Continuation of the surrounding `do` block; restored after the body. -/
   origCont : DoElemCont
+  /-- Substack at which the early-return handler was installed, if any. -/
   returnBase? : Option ControlStack
+  /-- Substack at which the `break` handler was installed, if any. -/
   breakBase? : Option ControlStack
+  /-- Substack at which the `continue` handler was installed, if any. -/
   continueBase? : Option ControlStack
-  pureBase : ControlStack
-  pureDeadCode : CodeLiveness
+  /-- The full transformer stack over the base monad. -/
+  liftedStack : ControlStack
+  /-- The body's elaboration type, `stM dec.resultType`. -/
   liftedDoBlockResultType : Expr
 
 -- abbrev M := List
 -- #reduce (types := true) M (Except Nat (Option (Option Bool) × String))
 -- #reduce (types := true) OptionT (OptionT (StateT String (ExceptT Nat M))) Bool
 
-def ControlLifter.ofCont (info : ControlInfo) (dec : DoElemCont) : DoElabM ControlLifter := do
+/-- Build the lifter plan for a body whose effects are summarised by `info`. -/
+def EffectForwarder.ofCont (info : ControlInfo) (dec : DoElemCont) : DoElabM EffectForwarder := do
   let mi := (← read).monadInfo
   let reassignedMutVars := (← read).mutVars |>.filter (info.reassigns.contains ·.getId)
   let reassignedMutVarNames := reassignedMutVars.map (·.getId)
@@ -231,21 +238,20 @@ def ControlLifter.ofCont (info : ControlInfo) (dec : DoElemCont) : DoElabM Contr
     returnBase?,
     breakBase?,
     continueBase?,
-    pureBase := controlStack,
-    -- The success continuation `origCont` is dead code iff the `ControlInfo` says so semantically.
-    pureDeadCode := if info.noFallthrough then .deadSemantically else .alive,
+    liftedStack := controlStack,
     liftedDoBlockResultType := (← controlStack.stM dec.resultType),
   }
 
 /--
-This function is like `MonadControl.liftWith fun runInBase => elabElem (runInBase pure)`.
-All continuations should be thought of as wrapped in `runInBase`, so that their effects are embedded
-in the terminal `stM m (t m)` result type. This wrapping will be realized by
-`ControlStack.synthesizeConts`, after we know what the transformer stack `t` looks like.
-What `t` looks like depends on whether reassignments, early `return`, `break` and `continue` are
-used, considering *all* the use sites of `ControlLifter.lift`.
+Elaborate `elabElem` in the lifted stack: install `break`/`continue`/`return`
+handlers funnelling into `liftedStack` and set the doBlock result type to
+`liftedDoBlockResultType`. Semantically
+`MonadControl.liftWith fun runInBase => elabElem (runInBase pure)`. Conts
+passed to `elabElem` are implicitly wrapped in `runInBase`, realised later by
+`ControlStack.synthesizeConts` once the transformer stack `t` is fixed by
+aggregating effects across *all* `lift` sites for this lifter.
 -/
-def ControlLifter.lift (l : ControlLifter) (elabElem : DoElemCont → DoElabM Expr) : DoElabM Expr := do
+def EffectForwarder.lift (l : EffectForwarder) (elabElem : DoElemCont → DoElabM Expr) : DoElabM Expr := do
   let oldBreakCont ← getBreakCont
   let oldContinueCont ← getContinueCont
   let oldReturnCont ← getReturnCont
@@ -262,9 +268,10 @@ def ControlLifter.lift (l : ControlLifter) (elabElem : DoElemCont → DoElabM Ex
     | some returnBase => { oldReturnCont with k := returnBase.mkReturn }
     | _ => oldReturnCont
   let contInfo := ContInfo.toContInfoRef { breakCont, continueCont, returnCont }
-  let pureCont := { l.origCont with k := l.pureBase.mkPure l.origCont.resultName, kind := .duplicable }
+  let pureCont := { l.origCont with k := l.liftedStack.mkPure l.origCont.resultName, kind := .duplicable }
   withReader (fun ctx => { ctx with contInfo, doBlockResultType := l.liftedDoBlockResultType }) do
     elabElem pureCont
 
-def ControlLifter.restoreCont (l : ControlLifter) : DoElabM DoElemCont := do
-  l.pureBase.restoreCont { l.origCont with k := withDeadCode l.pureDeadCode l.origCont.k }
+/-- Build a `DoElemCont` that unpacks the lifted body's result and resumes `origCont.k`. -/
+def EffectForwarder.restoreCont (l : EffectForwarder) : DoElabM DoElemCont :=
+  l.liftedStack.restoreCont l.origCont

src/Lean/Elab/Do/ForwardSyntax.lean

@@ -0,0 +1,47 @@
+/-
+Copyright (c) 2026 Lean FRO LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Sebastian Graf
+-/
+module
+
+prelude
+public import Lean.Parser.Term
+import Lean.Elab.Term.TermElabM
+meta import Lean.Parser.Do
+
+namespace Lean.Elab.Do
+
+open Lean.Parser.Term
+
+/--
+A body that effect-forwards into the enclosing `do` block via a `do←` marker, captured from
+syntax of the form `do← body` or `fun b₁ … bₖ => do← body`. `binders` is empty in the bare
+form.
+-/
+public structure ForwardArg where
+  binders : Array (TSyntax ``funBinder)
+  body    : TSyntax ``doSeq
+
+/--
+If `e` is a function application whose last argument effect-forwards into the enclosing `do`
+block (its body bears a `do←` marker, possibly under `fun` binders), return the application
+head sans the forwarding argument together with the `ForwardArg` description.
+-/
+public def Forward.matchApp? (e : Term) : Option (Term × ForwardArg) := do
+  let `($head $args*) := (⟨dropParens e.raw⟩ : Term) | none
+  let last : Term := ⟨dropParens (← args.back?).raw⟩
+  let arg : ForwardArg ←
+    if last.raw.isOfKind ``doForward then
+      some ⟨#[], ⟨last.raw[2]⟩⟩
+    else if let `(fun $bs:funBinder* => $rest) := last then
+      let rest := dropParens rest.raw
+      if rest.isOfKind ``doForward then some ⟨bs, ⟨rest[2]⟩⟩ else none
+    else none
+  let preArgs := args.pop
+  let head' : Term :=
+    if preArgs.isEmpty then head
+    else ⟨mkNode ``Lean.Parser.Term.app #[head, mkNullNode preArgs]⟩
+  some (head', arg)
+
+end Lean.Elab.Do

src/Lean/Elab/Do/InferControlInfo.lean

@@ -7,6 +7,7 @@ module
 
 prelude
 public import Lean.Elab.Term
+public import Lean.Elab.Do.ForwardSyntax
 meta import Lean.Parser.Do
 import Lean.Elab.Do.PatternVar
 
@@ -56,6 +57,7 @@ structure ControlInfo where
   reassigns : NameSet := {}
   deriving Inhabited
 
+/-- A `ControlInfo` for an element that always falls through normally with a single exit. -/
 def ControlInfo.pure : ControlInfo := {}
 
 /--
@@ -64,6 +66,7 @@ at all (so no regular exits and the next element is trivially unreachable).
 -/
 def ControlInfo.empty : ControlInfo := { numRegularExits := 0, noFallthrough := true }
 
+/-- Combine info for `a; b`: union effect flags, exits/dead follow `b`/either. -/
 def ControlInfo.sequence (a b : ControlInfo) : ControlInfo := {
     -- Syntactic fields aggregate unconditionally; the elaborator keeps visiting `b` unless `a` is
     -- a syntactically-terminal element (only top-level `return`/`break`/`continue` are, via
@@ -77,6 +80,7 @@ def ControlInfo.sequence (a b : ControlInfo) : ControlInfo := {
     noFallthrough := a.noFallthrough || b.noFallthrough,
   }
 
+/-- Combine info for branches `a | b`: union flags, sum exits, dead iff both dead. -/
 def ControlInfo.alternative (a b : ControlInfo) : ControlInfo := {
     breaks := a.breaks || b.breaks,
     continues := a.continues || b.continues,
@@ -127,7 +131,14 @@ partial def ofElem (stx : TSyntax `doElem) : TermElabM ControlInfo := do
   | `(doElem| break) => return { breaks := true, numRegularExits := 0, noFallthrough := true }
   | `(doElem| continue) => return { continues := true, numRegularExits := 0, noFallthrough := true }
   | `(doElem| return $[$_]?) => return { returnsEarly := true, numRegularExits := 0, noFallthrough := true }
-  | `(doExpr| $_:term) => return { numRegularExits := 1 }
+  | `(doExpr| $e:term) =>
+    if let some (_, arg) := Forward.matchApp? e then
+      -- The last arg is a `do + resume $body` block. Depending on the function, it might or might
+      -- not be called. We may definitely fall through and the control lifting mechanism prevents
+      -- creation of a join point.
+      return { (← ofSeq arg.body) with noFallthrough := false, numRegularExits := 1 }
+    else
+      return { numRegularExits := 1 }
   | `(doElem| do $doSeq) => ofSeq doSeq
   -- Let
   | `(doElem| let $[mut]? $_:letConfig $_:letDecl) => return .pure
@@ -262,8 +273,10 @@ end
 
 end InferControlInfo
 
+/-- Infer the `ControlInfo` of `doSeq`. -/
 def inferControlInfoSeq (doSeq : TSyntax ``doSeq) : TermElabM ControlInfo :=
   InferControlInfo.ofSeq doSeq
 
+/-- Infer the `ControlInfo` of a single doElem. -/
 def inferControlInfoElem (doElem : TSyntax `doElem) : TermElabM ControlInfo :=
   InferControlInfo.ofElem doElem

src/Lean/Elab/Term/TermElabM.lean

@@ -1660,11 +1660,6 @@ private def isLambdaWithImplicit (stx : Syntax) : Bool :=
   | `(fun $binders* => $_) => binders.raw.any fun b => b.isOfKind ``Lean.Parser.Term.implicitBinder || b.isOfKind `Lean.Parser.Term.instBinder
   | _                      => false
 
-private partial def dropTermParens : Syntax → Syntax := fun stx =>
-  match stx with
-  | `(($stx)) => dropTermParens stx
-  | _         => stx
-
 private def isHole (stx : Syntax) : Bool :=
   stx.isOfKind ``Lean.Parser.Term.hole || stx.isOfKind ``Lean.Parser.Term.syntheticHole
 
@@ -1692,7 +1687,7 @@ def mkNoImplicitLambdaAnnotation (type : Expr) : Expr :=
 
 /-- Block usage of implicit lambdas if `stx` is `@f` or `@f arg1 ...` or `fun` with an implicit binder annotation. -/
 def blockImplicitLambda (stx : Syntax) : Bool :=
-  let stx := dropTermParens stx
+  let stx := Parser.Term.dropParens stx
   -- TODO: make it extensible
   isExplicit stx || isExplicitApp stx || isLambdaWithImplicit stx || isHole stx || isTacticBlock stx ||
   isNoImplicitLambda stx || isTypeAscription stx