From 61232e788ac3a02c8d0e7a124830fe7356a01506 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 9 Jun 2022 18:47:53 +0200
Subject: [PATCH] fix: check types at do pattern reassignment

---
 src/Lean/Elab/Do.lean                   | 21 ++++-----------------
 tests/lean/1206.lean                    | 12 ++++++++++++
 tests/lean/1206.lean.expected.out       |  0
 tests/lean/doErrorMsg.lean              |  5 +++++
 tests/lean/doErrorMsg.lean.expected.out |  8 ++++++++
 5 files changed, 29 insertions(+), 17 deletions(-)
 create mode 100644 tests/lean/1206.lean
 create mode 100644 tests/lean/1206.lean.expected.out

diff --git a/src/Lean/Elab/Do.lean b/src/Lean/Elab/Do.lean
index 9687028dc8..67c633b572 100644
--- a/src/Lean/Elab/Do.lean
+++ b/src/Lean/Elab/Do.lean
@@ -961,23 +961,10 @@ def declToTerm (decl : Syntax) (k : Syntax) : M Syntax := withRef decl <| withFr
     Macro.throwErrorAt decl "unexpected kind of 'do' declaration"
 
 def reassignToTerm (reassign : Syntax) (k : Syntax) : MacroM Syntax := withRef reassign <| withFreshMacroScope do
-  let kind := reassign.getKind
-  if kind == ``Lean.Parser.Term.doReassign then
-    -- doReassign := leading_parser (letIdDecl <|> letPatDecl)
-    let arg := reassign[0]
-    if arg.getKind == ``Lean.Parser.Term.letIdDecl then
-      -- letIdDecl := leading_parser ident >> many (ppSpace >> bracketedBinder) >> optType >>  " := " >> termParser
-      let x   := arg[0]
-      let val := arg[4]
-      let newVal ← `(ensure_type_of% $x $(quote "invalid reassignment, value") $val)
-      let arg := arg.setArg 4 newVal
-      let letDecl := mkNode `Lean.Parser.Term.letDecl #[arg]
-      `(let $letDecl:letDecl; $k)
-    else
-      -- TODO: ensure the types did not change
-      let letDecl := mkNode `Lean.Parser.Term.letDecl #[arg]
-      `(let $letDecl:letDecl; $k)
-  else
+  match reassign with
+  | `(doElem| $x:ident := $rhs) => `(let $x:ident := ensure_type_of% $x $(quote "invalid reassignment, value") $rhs; $k)
+  | `(doElem| $e:term  := $rhs) => `(let $e:term  := ensure_type_of% $e $(quote "invalid reassignment, value") $rhs; $k)
+  | _ =>
     -- Note that `doReassignArrow` is expanded by `doReassignArrowToCode
     Macro.throwErrorAt reassign "unexpected kind of 'do' reassignment"
 
diff --git a/tests/lean/1206.lean b/tests/lean/1206.lean
new file mode 100644
index 0000000000..0553f06739
--- /dev/null
+++ b/tests/lean/1206.lean
@@ -0,0 +1,12 @@
+import Std
+
+set_option linter.unusedVariables true
+
+open Std
+
+def boo : Id (HashSet Nat) := do
+  let mut vs : HashSet Nat := HashSet.empty
+  for i in List.range 10 do
+    /- unused variable `vs` -/
+    (_, vs) := (i, vs.insert i)
+  return vs
diff --git a/tests/lean/1206.lean.expected.out b/tests/lean/1206.lean.expected.out
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/lean/doErrorMsg.lean b/tests/lean/doErrorMsg.lean
index d6e9cc13a7..b8aa7ddc59 100644
--- a/tests/lean/doErrorMsg.lean
+++ b/tests/lean/doErrorMsg.lean
@@ -22,3 +22,8 @@ def g2 : ExceptT String (StateT Nat Id) Unit := do
 def g3 : ExceptT String (StateT Nat Id) String := do
   let x ← f2
   f1
+
+example : Nat := Id.run do
+  let mut n : Nat := 0
+  (n, _) := (false, false)
+  n
diff --git a/tests/lean/doErrorMsg.lean.expected.out b/tests/lean/doErrorMsg.lean.expected.out
index 6ad2922f43..9851f3df3e 100644
--- a/tests/lean/doErrorMsg.lean.expected.out
+++ b/tests/lean/doErrorMsg.lean.expected.out
@@ -28,3 +28,11 @@ has type
   Nat → ExceptT String (StateT Nat Id) Nat : Type
 but is expected to have type
   ExceptT String (StateT Nat Id) ?m : Type
+doErrorMsg.lean:28:13-28:18: error: application type mismatch
+  Prod.mk false
+argument
+  false
+has type
+  Bool : Type
+but is expected to have type
+  Nat : Type