From d378df47d72c19a8783a75e45ef37b7b1c819363 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Thu, 16 Sep 2021 14:11:34 -0700
Subject: [PATCH] fix: fixes #633

---
 src/Lean/Elab/Do.lean            | 24 ++++++++++++++-------
 tests/lean/370.lean.expected.out |  2 +-
 tests/lean/run/633.lean          | 37 ++++++++++++++++++++++++++++++++
 3 files changed, 54 insertions(+), 9 deletions(-)
 create mode 100644 tests/lean/run/633.lean

diff --git a/src/Lean/Elab/Do.lean b/src/Lean/Elab/Do.lean
index 28196eb1d1..3aac240c53 100644
--- a/src/Lean/Elab/Do.lean
+++ b/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
diff --git a/tests/lean/370.lean.expected.out b/tests/lean/370.lean.expected.out
index 4308dfea2a..d06c872e53 100644
--- a/tests/lean/370.lean.expected.out
+++ b/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
diff --git a/tests/lean/run/633.lean b/tests/lean/run/633.lean
new file mode 100644
index 0000000000..d293d660df
--- /dev/null
+++ b/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 ()