feat: new, extensible do elaborator

src/Init/Control/Do.lean

@@ -0,0 +1,63 @@
+/-
+Copyright (c) 2025 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 Init.Control.Except
+public import Init.Control.Option
+
+public section
+
+/-!
+This module provides specialized wrappers around `ExceptT` to support the `do` elaborator.
+
+Specifically, the types here are used to tunnel early `return`, `break` and `continue` through
+non-algebraic higher-order effect combinators such as `tryCatch`.
+-/
+
+/-- A wrapper around `ExceptT` signifying early return. -/
+@[expose]
+abbrev EarlyReturnT (ρ m α) := ExceptT ρ m α
+
+/-- Exit a computation by returning a value `r : ρ` early. -/
+@[always_inline, inline, expose]
+abbrev EarlyReturnT.return {ρ m α} [Monad m] (r : ρ) : EarlyReturnT ρ m α :=
+  throw r
+
+/-- A specialization of `Except.casesOn`. -/
+@[always_inline, inline, expose]
+abbrev EarlyReturn.runK {ρ α : Type u} {β : Type v} (x : Except ρ α) (ret : ρ → β) (pure : α → β) : β :=
+  x.casesOn ret pure
+
+/-- A wrapper around `OptionT` signifying `break` in a loop. -/
+@[expose]
+abbrev BreakT := OptionT
+
+/-- Exit a loop body via `break`. -/
+@[always_inline, inline, expose]
+abbrev BreakT.break {m : Type w → Type x} [Monad m] : BreakT m α := failure
+
+/-- A specialization of `Option.casesOn`. -/
+@[always_inline, inline, expose]
+abbrev Break.runK {α : Type u} {β : Type v} (x : Option α) (breakK : Unit → β) (successK : α → β) : β :=
+  -- Note: The matcher below is used in the elaborator targeting `forIn` loops.
+  -- If you change the order of match arms here, you may need to adjust the elaborator.
+  match x with
+  | some a => successK a
+  | none => breakK ()
+
+/-- A wrapper around `OptionT` signifying `continue` in a loop. -/
+@[expose]
+abbrev ContinueT := OptionT
+
+/-- Exit a loop body via `continue`. -/
+@[always_inline, inline, expose]
+abbrev ContinueT.continue {m : Type w → Type x} [Monad m] : ContinueT m α := failure
+
+/-- A specialization of `Option.casesOn`. -/
+@[always_inline, inline, expose]
+abbrev Continue.runK {α : Type u} {β : Type v} (x : Option α) (continueK : Unit → β) (successK : α → β) : β :=
+  x.casesOn continueK (fun a _ => successK a) ()

src/Init/Control/Id.lean

@@ -79,3 +79,11 @@ instance : LawfulMonadAttach Id where
     exact x.run.2
 
 end Id
+
+/-- Turn a collection with a pure `ForIn` instance into an array. -/
+def ForIn.toArray {α : Type u} [inst : ForIn Id ρ α] (xs : ρ) : Array α :=
+  ForIn.forIn xs Array.empty (fun a acc => pure (.yield (acc.push a))) |> Id.run
+
+/-- Turn a collection with a pure `ForIn` instance into a list. -/
+def ForIn.toList {α : Type u} [ForIn Id ρ α] (xs : ρ) : List α :=
+  ForIn.toArray xs |>.toList

src/Init/Data/Array/Basic.lean

@@ -6,6 +6,7 @@ Authors: Leonardo de Moura
 module
 
 prelude
+public import Init.Control.Do
 public import Init.GetElem
 public import Init.Data.List.ToArrayImpl
 import all Init.Data.List.ToArrayImpl

src/Init/Meta/Defs.lean

@@ -784,8 +784,8 @@ def mkOptionalNode (arg : Option Syntax) : Syntax :=
 /--
 Creates a hole (`_`). The hole's position is copied from `ref`.
 -/
-def mkHole (ref : Syntax) (canonical := false) : Syntax :=
-  mkNode `Lean.Parser.Term.hole #[mkAtomFrom ref "_" canonical]
+def mkHole (ref : Syntax) (canonical := false) : Term :=
+  ⟨mkNode `Lean.Parser.Term.hole #[mkAtomFrom ref "_" canonical]⟩
 
 namespace Syntax
 

src/Init/Util.lean

@@ -30,6 +30,14 @@ def dbgTraceIfShared {α : Type u} (s : String) (a : α) : α := a
 @[never_extract, extern "lean_dbg_stack_trace"]
 def dbgStackTrace {α : Type u} (f : Unit → α) : α := f ()
 
+/--
+Print stack trace to stderr before evaluating given closure if `cond` is true.
+Currently supported on Linux only.
+-/
+@[never_extract]
+def dbgStackTraceIf {α : Type u} (cond : Bool) (f : Unit → α) : α :=
+  if cond then dbgStackTrace f else f ()
+
 @[extern "lean_dbg_sleep"]
 def dbgSleep {α : Type u} (ms : UInt32) (f : Unit → α) : α := f ()
 

src/Init/While.lean

@@ -40,7 +40,7 @@ macro_rules
 syntax "while " ident " : " termBeforeDo " do " doSeq : doElem
 
 macro_rules
-  | `(doElem| while $h : $cond do $seq) => `(doElem| repeat if $h : $cond then $seq else break)
+  | `(doElem| while $h : $cond do $seq) => `(doElem| repeat if $h:ident : $cond then $seq else break)
 
 syntax "while " termBeforeDo " do " doSeq : doElem
 

src/Lean/Elab.lean

@@ -65,3 +65,4 @@ public import Lean.Elab.ErrorExplanation
 public import Lean.Elab.DocString
 public import Lean.Elab.DocString.Builtin
 public import Lean.Elab.Parallel
+public import Lean.Elab.BuiltinDo

src/Lean/Elab/BindersUtil.lean

@@ -8,6 +8,7 @@ module
 prelude
 public import Lean.Parser.Term
 meta import Lean.Parser.Term
+meta import Lean.Parser.Do
 import Init.Syntax
 
 public section
@@ -65,12 +66,19 @@ def shouldExpandMatchAlt : TSyntax ``matchAlt → Bool
 def expandMatchAlts? (stx : Syntax) : MacroM (Option Syntax) := do
   match stx with
   | `(match $[$gen]? $[$motive]? $discrs,* with $alts:matchAlt*) =>
-    if alts.any shouldExpandMatchAlt then
-      let alts ← alts.foldlM (init := #[]) fun alts alt => return alts ++ (expandMatchAlt alt)
+     expand alts >>= fun alts? => alts?.mapM fun alts =>
       `(match $[$gen]? $[$motive]? $discrs,* with $alts:matchAlt*)
-    else
-      return none
+  | `(doElem| match $[$dep?]? $[$gen]? $[$motive]? $discrs,* with $alts:matchAlt*) =>
+     expand alts >>= fun alts? => alts?.mapM fun alts =>
+      `(doElem| match $[$dep?]? $[$gen]? $[$motive]? $discrs,* with $alts:matchAlt*)
   | _ => return none
+  where
+    expand (alts : Array (TSyntax ``matchAlt)) : MacroM (Option (Array (TSyntax ``matchAlt))) := do
+      if alts.any shouldExpandMatchAlt then
+        let alts ← alts.foldlM (init := #[]) fun alts alt => return alts ++ expandMatchAlt alt
+        return some alts
+      else
+        return none
 
 open TSyntax.Compat in
 def clearInMatchAlt (stx : TSyntax ``matchAlt) (vars : Array Ident) : TSyntax ``matchAlt :=

src/Lean/Elab/BuiltinDo.lean

@@ -0,0 +1,17 @@
+/-
+Copyright (c) 2025 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.BuiltinDo.Basic
+public import Lean.Elab.BuiltinDo.Let
+public import Lean.Elab.BuiltinDo.Match
+public import Lean.Elab.BuiltinDo.MatchExpr
+public import Lean.Elab.BuiltinDo.If
+public import Lean.Elab.BuiltinDo.Jump
+public import Lean.Elab.BuiltinDo.Misc
+public import Lean.Elab.BuiltinDo.For
+public import Lean.Elab.BuiltinDo.TryCatch

src/Lean/Elab/BuiltinDo/Basic.lean

@@ -0,0 +1,27 @@
+/-
+Copyright (c) 2025 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
+
+public section
+
+namespace Lean.Elab.Do
+
+open Lean.Parser.Term
+open Lean.Meta
+
+def elabDoIdDecl (x : Ident) (xType? : Option Term) (rhs : TSyntax `doElem) (k : DoElabM Expr)
+    (kind : DoElemContKind := .nonDuplicable) : DoElabM Expr := do
+  let xType ← Term.elabType (xType?.getD (mkHole x))
+  let lctx ← getLCtx
+  let ctx ← read
+  elabDoElem rhs <| .mk (kind := kind) x.getId xType do
+    withLCtxKeepingMutVarDefs lctx ctx x.getId do
+      Term.addLocalVarInfo x (← getFVarFromUserName x.getId)
+      k

src/Lean/Elab/BuiltinDo/For.lean

@@ -0,0 +1,190 @@
+/-
+Copyright (c) 2025 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.BuiltinDo.Basic
+meta import Lean.Parser.Do
+import Init.Control.Do
+import Lean.Meta.ProdN
+
+public section
+
+namespace Lean.Elab.Do
+
+open Lean.Parser.Term
+open Lean.Meta
+
+@[builtin_macro Lean.Parser.Term.doFor] def expandDoFor : Macro := fun stx => do
+  match stx with
+  | `(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) =>
+    let decls := decls.getElems
+    let `(doForDecl| $[$h? : ]? $pattern in $xs) := decls[0]! | Macro.throwUnsupported
+    let mut doElems := #[]
+    let mut body := body
+    -- Expand `pattern` into an `Ident` `x`:
+    let x ←
+      if pattern.raw.isIdent then
+        pure ⟨pattern⟩
+      else if pattern.raw.isOfKind ``Lean.Parser.Term.hole then
+        Term.mkFreshIdent pattern
+      else
+        -- This case is a last resort, because it introduces a `match` and that will cause eager
+        -- defaulting. In practice this means that `mut` vars default to `Nat` too often.
+        -- Hence we try to only generate a `match` if we absolutely must.
+        let x ← Term.mkFreshIdent pattern
+        body ← `(doSeq| match $x:term with | $pattern => $body)
+        pure x
+    -- Expand the remaining `doForDecl`s:
+    for doForDecl in decls[1...*] do
+      /-
+        Expand
+        ```
+        for x in xs, y in ys do
+          body
+        ```
+        into
+        ```
+        let mut s := Std.toStream ys
+        for x in xs do
+          match Std.Stream.next? s with
+          | none => break
+          | some (y, s') =>
+            s := s'
+            body
+        ```
+      -/
+      let `(doForDecl| $[$h? : ]? $y in $ys) := doForDecl | Macro.throwUnsupported
+      if let some h := h? then
+        Macro.throwErrorAt h "The proof annotation here has not been implemented yet."
+      /- Recall that `@` (explicit) disables `coeAtOutParam`.
+         We used `@` at `Stream` functions to make sure `resultIsOutParamSupport` is not used. -/
+      let toStreamApp ← withRef ys `(@Std.toStream _ _ _ $ys)
+      let s := mkIdentFrom ys (← withFreshMacroScope <| MonadQuotation.addMacroScope `__s)
+      doElems := doElems.push (← `(doSeqItem| let mut $s := $toStreamApp:term))
+      body ← `(doSeq|
+        match @Std.Stream.next? _ _ _ $s with
+          | none => break
+          | some ($y, s') =>
+            $s:ident := s'
+            do $body)
+    doElems := doElems.push (← `(doSeqItem| for $[$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
+  checkMutVarsForShadowing #[x]
+  let uα ← mkFreshLevelMVar
+  let uρ ← mkFreshLevelMVar
+  let α ← mkFreshExprMVar (mkSort (uα.succ)) (userName := `α) -- assigned by outParam
+  let ρ ← mkFreshExprMVar (mkSort (uρ.succ)) (userName := `ρ) -- assigned in the next line
+  let xs ← Term.elabTermEnsuringType xs ρ
+  let mi := (← read).monadInfo
+  let mutVars := (← read).mutVars
+
+  let info ← inferControlInfoSeq body
+  let oldReturnCont ← getReturnCont
+  let returnVarName ← mkFreshUserName `__r
+  let loopMutVars := mutVars.filter fun x => info.reassigns.contains x.getId
+  let loopMutVarNames :=
+    if info.returnsEarly then
+      returnVarName :: (loopMutVars.map (·.getId)).toList
+    else
+      (loopMutVars.map (·.getId)).toList
+  let useLoopMutVars (e : Option Expr) : TermElabM (Array Expr) := do
+    let mut defs := #[]
+    unless e.isNone || info.returnsEarly do
+      throwError "Early returning {e} but the info said there is no early return"
+    if info.returnsEarly then
+      let returnVar ←
+        match e with
+        | none => mkNone oldReturnCont.resultType
+        | some e => mkSome oldReturnCont.resultType e
+      defs := defs.push returnVar
+    for x in loopMutVars do
+      let defn ← getLocalDeclFromUserName x.getId
+      Term.addTermInfo' x defn.toExpr
+      -- ForIn forces all mut vars into the same universe: that of the do block result type.
+      discard <| Term.ensureHasType (mkSort (mi.u.succ)) defn.type
+      defs := defs.push defn.toExpr
+    if info.returnsEarly && loopMutVars.isEmpty then
+      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
+    | none =>
+      let instForIn ← Term.mkInstMVar <| mkApp3 (mkConst ``ForIn [uρ, uα, mi.u, mi.v]) mi.m ρ α
+      let app := mkConst ``ForIn.forIn [uρ, uα, mi.u, mi.v]
+      -- ForIn.forIn : {m ρ α : _} → [ForIn m ρ α] → {β : _} → ρ → β → (α → β → m (ForInStep β)) → m β
+      let app := mkApp7 app mi.m ρ α instForIn σ xs preS -- 1 arg remaining: loop body
+      pure (app, none)
+    | some _ =>
+      let d ← mkFreshExprMVar (mkApp2 (mkConst ``Membership [uα, uρ]) α ρ) (userName := `d) -- outParam
+      let instForIn ← Term.mkInstMVar <| mkApp4 (mkConst ``ForIn' [uρ, uα, mi.u, mi.v]) mi.m ρ α d
+      let app := mkConst ``ForIn'.forIn' [uρ, uα, mi.u, mi.v]
+      -- ForIn'.forIn' : {m ρ α : _} → [Membership α ρ] → [ForIn' m ρ α d] → {β : _} → ρ → β → ((a : α) → a ∈ x → β → m (ForInStep β)) → m β
+      let app := mkApp8 app mi.m ρ α d instForIn σ xs preS -- 1 arg remaining: loop body
+      pure (app, some d)
+  let s ← mkFreshUserName `__s
+  let xh : Array (Name × (Array Expr → DoElabM Expr)) := match h?, p? with
+    | some h, some d =>
+      #[(x.getId, fun _ => pure α),
+        (h.getId, fun x => pure (mkApp5 (mkConst ``Membership.mem [uα, uρ]) α ρ d xs x[0]!))]
+    | _, _ =>
+      #[(x.getId, fun _ => pure α)]
+
+  let body ←
+    withLocalDeclsD xh fun xh => do
+    withLocalDecl s .default σ (kind := .implDetail) fun loopS => do
+    mkLambdaFVars (xh.push loopS) <| ← do
+    bindMutVarsFromTuple loopMutVarNames loopS.fvarId! do
+    let newDoBlockResultType := mkApp (mkConst ``ForInStep [mi.u]) σ
+    withDoBlockResultType newDoBlockResultType do
+    let continueCont := do
+      let (tuple, _tupleTy) ← mkProdMkN (← useLoopMutVars none) mi.u
+      let yield := mkApp2 (mkConst ``ForInStep.yield [mi.u]) σ tuple
+      mkPureApp newDoBlockResultType yield
+    let breakCont := do
+      let (tuple, _tupleTy) ← mkProdMkN (← useLoopMutVars none) mi.u
+      let done := mkApp2 (mkConst ``ForInStep.done [mi.u]) σ tuple
+      mkPureApp newDoBlockResultType done
+    let returnCont := { oldReturnCont with k := fun e => do
+        let (tuple, _tupleTy) ← mkProdMkN (← useLoopMutVars (some e)) mi.u
+        let done := mkApp2 (mkConst ``ForInStep.done [mi.u]) σ tuple
+        mkPureApp newDoBlockResultType done
+      }
+    enterLoopBody breakCont continueCont returnCont do
+    -- Elaborate the loop body, which must have result type `PUnit`, just like the whole `for` loop.
+    elabDoSeq body { dec with k := continueCont, kind := .duplicable }
+
+  let forIn := mkApp app body
+
+  let γ := (← read).doBlockResultType
+  let rest ←
+    withLocalDeclD s σ fun postS => do mkLambdaFVars #[postS] <| ← do
+      bindMutVarsFromTuple loopMutVarNames postS.fvarId! do
+        if info.returnsEarly then
+          let ret ← getFVarFromUserName returnVarName
+          let ret ← if loopMutVars.isEmpty then mkAppM ``Prod.fst #[ret] else pure ret
+          let motive := mkLambda `_ .default (← inferType ret) (← mkMonadicType γ)
+          let app := mkApp3 (mkConst ``Break.runK.match_1 [mi.u, mi.v.succ]) oldReturnCont.resultType motive ret
+          let none := mkSimpleThunk (← dec.continueWithUnit)
+          let some ← withLocalDeclD (← mkFreshUserName `r) oldReturnCont.resultType fun r => do
+            mkLambdaFVars #[r] (← oldReturnCont.k r)
+          return mkApp2 app some none
+        else
+          dec.continueWithUnit
+
+  mkBindApp σ γ forIn rest