From 791388400b4a4f68d9ea4d4462302eeee07615d5 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Thu, 14 Jan 2021 14:18:56 -0800
Subject: [PATCH] feat: improve `do` error messages

cc @Kha @Vtec234
---
 src/Init/Meta.lean                      | 24 ++++++++++++++++++++
 src/Lean/Elab/Do.lean                   |  9 +++++---
 tests/lean/doErrorMsg.lean              | 24 ++++++++++++++++++++
 tests/lean/doErrorMsg.lean.expected.out | 30 +++++++++++++++++++++++++
 4 files changed, 84 insertions(+), 3 deletions(-)
 create mode 100644 tests/lean/doErrorMsg.lean
 create mode 100644 tests/lean/doErrorMsg.lean.expected.out

diff --git a/src/Init/Meta.lean b/src/Init/Meta.lean
index 88290d500b..e94912d27d 100644
--- a/src/Init/Meta.lean
+++ b/src/Init/Meta.lean
@@ -229,6 +229,30 @@ def copyInfo (s : Syntax) (source : Syntax) : Syntax :=
  let s := s.copyHeadInfo source
  s.copyTailInfo source
 
+/--
+  Copy head and tail position information from `source` to `s`.
+  `leading` and `trailing` information is not preserved. -/
+def copyRangePos (s : Syntax) (source : Syntax) : Syntax :=
+  match source.getPos with
+  | none     => s
+  | some pos =>
+    let s := s.setHeadInfo { pos := pos }
+    match source.getTailInfo with
+    | some { pos := some pos, .. } =>
+      let s := s.setTailInfo { pos := pos }
+      /- The trailing token at `s` may be different from `source`.
+         So, we retrieve the tail positions and adjust `pos` to make sure the `s.getTailPos == source.getTailPos`. -/
+      match source.getTailPos, s.getTailPos with
+      | some pos₁, some pos₂ =>
+        if pos₁ < pos₂ then
+          s.setTailInfo { pos := some ((pos : Nat) - (pos₂ - pos₁) : Nat) }
+        else if pos₁ > pos₂ then
+          s.setTailInfo { pos := some ((pos : Nat) + (pos₁ - pos₂) : Nat) }
+        else
+          s
+      | _, _ => s
+    | _ => s
+
 end Syntax
 
 def mkAtom (val : String) : Syntax :=
diff --git a/src/Lean/Elab/Do.lean b/src/Lean/Elab/Do.lean
index 710114f9a8..a205d57807 100644
--- a/src/Lean/Elab/Do.lean
+++ b/src/Lean/Elab/Do.lean
@@ -862,7 +862,7 @@ def actionTerminalToTerm (action : Syntax) : M Syntax := do
   let r ← actionTerminalToTermCore action
   pure $ r.copyInfo ref
 
-def seqToTermCore (action : Syntax) (k : Syntax) : MacroM Syntax := withFreshMacroScope do
+def seqToTermCore (action : Syntax) (k : Syntax) : M Syntax := withFreshMacroScope do
   if action.getKind == `Lean.Parser.Term.doDbgTrace then
     let msg := action[1]
     `(dbgTrace! $msg; $k)
@@ -870,9 +870,10 @@ def seqToTermCore (action : Syntax) (k : Syntax) : MacroM Syntax := withFreshMac
     let cond := action[1]
     `(assert! $cond; $k)
   else
+    let action := Syntax.copyRangePos (← `(($action : $((←read).m) PUnit))) action
     `(Bind.bind $action (fun (_ : PUnit) => $k))
 
-def seqToTerm (action : Syntax) (k : Syntax) : MacroM Syntax := do
+def seqToTerm (action : Syntax) (k : Syntax) : M Syntax := do
   let r ← seqToTermCore action k
   return r.copyInfo action
 
@@ -894,7 +895,9 @@ def declToTermCore (decl : Syntax) (k : Syntax) : M Syntax := withFreshMacroScop
       let doElem := arg[3]
       -- `doElem` must be a `doExpr action`. See `doLetArrowToCode`
       match isDoExpr? doElem with
-      | some action => `(Bind.bind $action (fun ($id:ident : $type) => $k))
+      | some action =>
+        let action := Syntax.copyRangePos (← `(($action : $((← read).m) $type))) action
+        `(Bind.bind $action (fun ($id:ident : $type) => $k))
       | none        => Macro.throwErrorAt decl "unexpected kind of 'do' declaration"
     else
       Macro.throwErrorAt decl "unexpected kind of 'do' declaration"
diff --git a/tests/lean/doErrorMsg.lean b/tests/lean/doErrorMsg.lean
new file mode 100644
index 0000000000..d6e9cc13a7
--- /dev/null
+++ b/tests/lean/doErrorMsg.lean
@@ -0,0 +1,24 @@
+def f : IO Nat := do
+  IO.println "hello"
+  IO.getStdin
+  return 10
+
+def f1 : ExceptT String (StateT Nat Id) Nat := do
+  modify (· + 1)
+  get
+
+def f2 (x : Nat) : ExceptT String (StateT Nat Id) Nat := do
+  modify (· + x)
+  get
+
+def g1 : ExceptT String (StateT Nat Id) Unit := do
+  let x : String ← f1
+  return ()
+
+def g2 : ExceptT String (StateT Nat Id) Unit := do
+  let x : String ← f2 10
+  return ()
+
+def g3 : ExceptT String (StateT Nat Id) String := do
+  let x ← f2
+  f1
diff --git a/tests/lean/doErrorMsg.lean.expected.out b/tests/lean/doErrorMsg.lean.expected.out
new file mode 100644
index 0000000000..e9b90f9bc1
--- /dev/null
+++ b/tests/lean/doErrorMsg.lean.expected.out
@@ -0,0 +1,30 @@
+doErrorMsg.lean:3:2: error: type mismatch
+  IO.getStdin
+has type
+  EIO IO.Error IO.FS.Stream
+but is expected to have type
+  EIO IO.Error PUnit
+doErrorMsg.lean:15:19: error: type mismatch
+  f1
+has type
+  ExceptT String (StateT Nat Id) Nat
+but is expected to have type
+  ExceptT String (StateT Nat Id) String
+doErrorMsg.lean:19:19: error: type mismatch
+  f2 10
+has type
+  ExceptT String (StateT Nat Id) Nat
+but is expected to have type
+  ExceptT String (StateT Nat Id) String
+doErrorMsg.lean:24:2: error: type mismatch
+  f1
+has type
+  ExceptT String (StateT Nat Id) Nat
+but is expected to have type
+  ExceptT String (StateT Nat Id) String
+doErrorMsg.lean:23:10: error: type mismatch
+  f2
+has type
+  Nat → ExceptT String (StateT Nat Id) Nat
+but is expected to have type
+  ExceptT String (StateT Nat Id) ?m