From 09dcf718c138fc4c65c83b1bbac43f97db5b82a3 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Thu, 8 Oct 2020 11:56:03 -0700
Subject: [PATCH] feat: expand `doHave`

---
 src/Lean/Elab/Do.lean           | 25 ++++++++++++++++++++++---
 tests/lean/run/doNotation3.lean | 23 +++++++++++++++++++++++
 2 files changed, 45 insertions(+), 3 deletions(-)
 create mode 100644 tests/lean/run/doNotation3.lean

diff --git a/src/Lean/Elab/Do.lean b/src/Lean/Elab/Do.lean
index 0597d88a28..5ddc14cbc1 100644
--- a/src/Lean/Elab/Do.lean
+++ b/src/Lean/Elab/Do.lean
@@ -540,6 +540,21 @@ def getDoLetVars (doLet : Syntax) : TermElabM (Array Name) :=
 -- parser! "let " >> letDecl
 getLetDeclVars (doLet.getArg 1)
 
+def getDoHaveVar (doHave : Syntax) : Name :=
+/-
+  `parser! "have " >> Term.haveDecl`
+  where
+  ```
+  haveDecl := optIdent >> termParser >> (haveAssign <|> fromTerm <|> byTactic)
+  optIdent := optional (try (ident >> " : "))
+
+  ``` -/
+let optIdent := doHave.getArg 1;
+if optIdent.isNone then
+  `this
+else
+  optIdent.getIdAt 0
+
 def getDoLetRecVars (doLetRec : Syntax) : TermElabM (Array Name) := do
 -- letRecDecls is an array of `(group (optional attributes >> letDecl))`
 let letRecDecls := (doLetRec.getArg 1).getArgs.getSepElems;
@@ -778,7 +793,10 @@ else if kind == `Lean.Parser.Term.doLetArrow then
   else
     liftM $ Macro.throwError decl "unexpected kind of 'do' declaration"
 else if kind == `Lean.Parser.Term.doHave then
-  liftM $ Macro.throwError decl ("WIP " ++ toString decl)
+  -- The `have` term is of the form  `"have " >> haveDecl >> optSemicolon termParser`
+  let args := decl.getArgs;
+  let args := args ++ #[mkNullNode /- optional ';' -/, k];
+  pure $ mkNode `Lean.Parser.Term.«have» args
 else
   liftM $ Macro.throwError decl "unexpected kind of 'do' declaration"
 
@@ -1123,6 +1141,9 @@ partial def doSeqToCode : List Syntax → M CodeBlock
     if k == `Lean.Parser.Term.doLet then do
       vars ← liftM $ getDoLetVars doElem;
       mkVarDeclCore vars doElem <$> withNewVars vars (doSeqToCode doElems)
+    else if k == `Lean.Parser.Term.doHave then
+      let var := getDoHaveVar doElem;
+      mkVarDeclCore #[var] doElem <$> withNewVars #[var] (doSeqToCode doElems)
     else if k == `Lean.Parser.Term.doLetRec then do
       vars ← liftM $ getDoLetRecVars doElem;
       mkVarDeclCore vars doElem <$> withNewVars vars (doSeqToCode doElems)
@@ -1135,8 +1156,6 @@ partial def doSeqToCode : List Syntax → M CodeBlock
       doLetArrowToCode doSeqToCode doElem doElems
     else if k == `Lean.Parser.Term.doReassignArrow then
       throwError "WIP"
-    else if k == `Lean.Parser.Term.doHave then
-      throwError "WIP"
     else if k == `Lean.Parser.Term.doIf then
       doIfToCode doSeqToCode doElem doElems
     else if k == `Lean.Parser.Term.doUnless then do
diff --git a/tests/lean/run/doNotation3.lean b/tests/lean/run/doNotation3.lean
new file mode 100644
index 0000000000..97943a1e80
--- /dev/null
+++ b/tests/lean/run/doNotation3.lean
@@ -0,0 +1,23 @@
+new_frontend
+
+theorem zeroLtOfLt : {a b : Nat} → a < b → 0 < b
+| 0,   _, h => h
+| a+1, b, h =>
+  have a < b from Nat.ltTrans (Nat.ltSuccSelf _) h
+  zeroLtOfLt this
+
+def fold {m α β} [Monad m] (as : Array α) (b : β) (f : α → β → m β) : m β := do
+let rec loop : (i : Nat) → i ≤ as.size → β → m β
+  | 0,   h, b => b
+  | i+1, h, b => do
+    have h' : i < as.size          from Nat.ltOfLtOfLe (Nat.ltSuccSelf i) h
+    have as.size - 1 < as.size     from Nat.subLt (zeroLtOfLt h') (decide! (0 < 1))
+    have as.size - 1 - i < as.size from Nat.ltOfLeOfLt (Nat.subLe (as.size - 1) i) this
+    let b ← f (as.get ⟨as.size - 1 - i, this⟩) b
+    loop i (Nat.leOfLt h') b
+loop as.size (Nat.leRefl _) b
+
+#eval Id.run $ fold #[1, 2, 3, 4] 0 (pure $ · + ·)
+
+theorem ex : (Id.run $ fold #[1, 2, 3, 4] 0 (pure $ · + ·)) = 10 :=
+rfl