fix: fixes #633

src/Lean/Elab/Do.lean

@@ -1605,15 +1605,23 @@ private def mkMonadAlias (m : Expr) : TermElabM Syntax := do
   assignExprMVar mvar.mvarId! m
   pure result
 
-private partial def ensureArrowNotUsed : Syntax → MacroM Unit
-  | stx@(Syntax.node k args) => do
-    if k == ``Parser.Term.liftMethod || k == ``Parser.Term.doLetArrow || k == ``Parser.Term.doReassignArrow || k == ``Parser.Term.doIfLetBind then
-      Macro.throwErrorAt stx "`←` and `<-` are not allowed in pure `do` blocks, i.e., blocks where Lean implicitly used the `Id` monad"
-    args.forM ensureArrowNotUsed
-  | _ => pure ()
+private partial def ensureArrowNotUsed (stx : Syntax) : MacroM Unit := do
+  /- We expand macros here because we stop the search at nested `do`s
+     So, we also want to stop the search at macros that expand `do ...`.
+     Hopefully this is not a performance bottleneck in practice. -/
+  let stx ← expandMacros stx
+  stx.getArgs.forM go
+where
+  go (stx : Syntax) : MacroM Unit :=
+    match stx with
+    | Syntax.node k args => do
+      if k == ``Parser.Term.liftMethod || k == ``Parser.Term.doLetArrow || k == ``Parser.Term.doReassignArrow || k == ``Parser.Term.doIfLetBind then
+        Macro.throwErrorAt stx "`←` and `<-` are not allowed in pure `do` blocks, i.e., blocks where Lean implicitly used the `Id` monad"
+      unless k == ``Parser.Term.do do
+        args.forM go
+    | _ => pure ()
 
-@[builtinTermElab «do»]
-def elabDo : TermElab := fun stx expectedType? => do
+@[builtinTermElab «do»] def elabDo : TermElab := fun stx expectedType? => do
   tryPostponeIfNoneOrMVar expectedType?
   let bindInfo ← extractBind expectedType?
   if bindInfo.isPure then

tests/lean/370.lean.expected.out

@@ -4,4 +4,4 @@
 370.lean:8:2-8:16: error: `←` and `<-` are not allowed in pure `do` blocks, i.e., blocks where Lean implicitly used the `Id` monad
 370.lean:12:16-12:24: error: `←` and `<-` are not allowed in pure `do` blocks, i.e., blocks where Lean implicitly used the `Id` monad
 370.lean:19:2-19:12: error: `←` and `<-` are not allowed in pure `do` blocks, i.e., blocks where Lean implicitly used the `Id` monad
-370.lean:23:2-23:26: error: `←` and `<-` are not allowed in pure `do` blocks, i.e., blocks where Lean implicitly used the `Id` monad
+370.lean:23:2-23:25: error: `←` and `<-` are not allowed in pure `do` blocks, i.e., blocks where Lean implicitly used the `Id` monad

tests/lean/run/633.lean

@@ -0,0 +1,37 @@
+abbrev semantics (α:Type) := StateM (List Nat) α
+
+inductive expression : Nat → Type
+| const : (n : Nat) → expression n
+
+def uext {w:Nat} (x: expression w) (o:Nat) : expression w := expression.const w
+def eval {n : Nat} (v:expression n) : semantics (expression n) := pure (expression.const n)
+def set_overflow {w : Nat} (e : expression w) : semantics Unit := pure ()
+
+structure instruction :=
+  (mnemonic:String)
+  (patterns:List Nat)
+
+def definst (mnem:String) (body: expression 8 -> semantics Unit) : instruction :=
+{ mnemonic := mnem
+, patterns := ((body (expression.const 8)).run []).snd.reverse
+}
+
+def mul : instruction := do -- this is a "pure" do block (as in it is the Id monad)
+ definst "mul" $ fun (src : expression 8) =>
+    let action : semantics Unit := do -- this is not "pure" do block
+      let tmp <- eval $ uext src 16
+      set_overflow $ tmp
+    action
+
+def mul' : instruction := do -- this is a "pure" do block (as in it is the Id monad)
+ definst "mul" $ fun (src : expression 8) =>
+    let rec action : semantics Unit := do -- this is not "pure" do block
+      let tmp <- eval $ uext src 16
+      set_overflow $ tmp
+    action
+
+def mul'' : instruction := do -- this is a "pure" do block (as in it is the Id monad)
+ definst "mul" $ fun (src : expression 8) =>
+    let action : semantics (expression 8) :=
+      return (<- eval $ uext src 16)
+    pure ()