From 87bf97bdc19d7fa09ffc78877f30936724f4d342 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sat, 31 Oct 2020 08:13:09 -0700
Subject: [PATCH] feat: expand term `try`, `for`, `unless`, and `return`

---
 src/Lean/Elab/Do.lean                | 18 ++++++++++++++++++
 tests/lean/doElemAsTermNotation.lean | 27 +++++++++++++++++++++++++++
 2 files changed, 45 insertions(+)
 create mode 100644 tests/lean/doElemAsTermNotation.lean

diff --git a/src/Lean/Elab/Do.lean b/src/Lean/Elab/Do.lean
index d3bca52f8a..3258c3842f 100644
--- a/src/Lean/Elab/Do.lean
+++ b/src/Lean/Elab/Do.lean
@@ -1514,6 +1514,24 @@ end Do
 
 builtin_initialize registerTraceClass `Elab.do
 
+private def toDoElem (newKind : SyntaxNodeKind) : Macro := fun stx => do
+  let stx := stx.setKind newKind
+  let stxNew ← `(do $stx:doElem)
+  return stxNew.copyInfo stx
+
+@[builtinMacro Lean.Parser.Term.termFor]
+def expandTermFor : Macro := toDoElem `Lean.Parser.Term.doFor
+
+@[builtinMacro Lean.Parser.Term.termTry]
+def expandTermTry : Macro := toDoElem `Lean.Parser.Term.doTry
+
+@[builtinMacro Lean.Parser.Term.termUnless]
+def expandTermUnless : Macro := toDoElem `Lean.Parser.Term.doUnless
+
+@[builtinMacro Lean.Parser.Term.termReturn]
+def expandTermReturn : Macro := toDoElem `Lean.Parser.Term.doReturn
+
+
 end Term
 end Elab
 end Lean
diff --git a/tests/lean/doElemAsTermNotation.lean b/tests/lean/doElemAsTermNotation.lean
new file mode 100644
index 0000000000..92875e4ff1
--- /dev/null
+++ b/tests/lean/doElemAsTermNotation.lean
@@ -0,0 +1,27 @@
+def f1 (x : Nat) : IO Unit :=
+  unless x > 10 do
+    IO.println s!"x: {x}"
+
+#eval f1 0
+#eval f1 100
+
+def f2 (x : Nat) (ref : IO.Ref Nat) : IO Nat :=
+  return x + (←ref.get)
+
+#eval id (α := IO Nat) do f2 5 (← IO.mkRef 10)
+
+def f3 (x : Nat) : IO Nat :=
+  try
+    IO.println x
+    throw $ IO.userError "failed"
+  catch
+    | IO.Error.userError msg => IO.println s!"at catch: {msg}"; pure 0
+    | ex => throw ex
+
+#eval f3 5
+
+def f4 (xs : List Nat) : IO Unit :=
+  for x in xs do
+    IO.println x
+
+#eval f4 [1,2,3]