From eacdb5ff83135077d57239d2e1ba2ceeb785a59c Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 12 Oct 2020 11:50:13 -0700
Subject: [PATCH] feat: add `Range` notation

---
 src/Init/Data.lean        |  1 +
 src/Init/Data/Range.lean  | 44 +++++++++++++++++++++++++++++++++++++++
 tests/lean/run/range.lean | 29 ++++++++++++++++++++++++++
 3 files changed, 74 insertions(+)
 create mode 100644 src/Init/Data/Range.lean
 create mode 100644 tests/lean/run/range.lean

diff --git a/src/Init/Data.lean b/src/Init/Data.lean
index a1baaba130..450482db3d 100644
--- a/src/Init/Data.lean
+++ b/src/Init/Data.lean
@@ -19,3 +19,4 @@ import Init.Data.Float
 import Init.Data.Option
 import Init.Data.Random
 import Init.Data.ToString
+import Init.Data.Range
diff --git a/src/Init/Data/Range.lean b/src/Init/Data/Range.lean
new file mode 100644
index 0000000000..b3985be7a1
--- /dev/null
+++ b/src/Init/Data/Range.lean
@@ -0,0 +1,44 @@
+/-
+Copyright (c) 2020 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+prelude
+import Init.LeanInit
+new_frontend
+
+namespace Std
+-- We put `Range` in `Init` because we want the notation `[i:j]`  without importing `Std`
+-- We don't put `Range` in the top-level namespace to avoid collisions with user defined types
+structure Range :=
+(start : Nat := 0)
+(stop  : Nat)
+(step  : Nat := 1)
+
+namespace Range
+universes u v
+
+@[inline] def forIn {β : Type u} {m : Type u → Type v} [Monad m] (range : Range) (init : β) (f : Nat → β → m (ForInStep β)) : m β :=
+let rec loop (i : Nat) (j : Nat) (b : β) : m β := do
+  if j ≥ range.stop then
+    pure b
+  else match i with
+   | 0   => pure b
+   | i+1 => match ← f j b with
+      | ForInStep.done b  => pure b
+      | ForInStep.yield b => loop i (j + range.step) b
+loop range.stop range.start init
+
+syntax:max "[" ":" term "]" : term
+syntax:max "[" term ":" term "]" : term
+syntax:max "[" ":" term ":" term "]" : term
+syntax:max "[" term ":" term ":" term "]" : term
+
+macro_rules
+| `([ : $stop]) => `({ stop := $stop : Range })
+| `([ $start : $stop ]) => `({ start := $start, stop := $stop : Range })
+| `([ $start : $stop : $step ]) => `({ start := $start, stop := $stop, step := $step : Range })
+| `([ : $stop : $step ]) => `({ stop := $stop, step := $step : Range })
+
+end Range
+end Std
diff --git a/tests/lean/run/range.lean b/tests/lean/run/range.lean
new file mode 100644
index 0000000000..2538f32dbe
--- /dev/null
+++ b/tests/lean/run/range.lean
@@ -0,0 +1,29 @@
+new_frontend
+
+def ex1 : IO Unit := do
+IO.println "example 1"
+for x in [:100:10] do
+  IO.println s!"x: {x}"
+
+#eval ex1
+
+def ex2 : IO Unit := do
+IO.println "example 2"
+for x in [:10] do
+  IO.println s!"x: {x}"
+
+#eval ex2
+
+def ex3 : IO Unit := do
+IO.println "example 3"
+for x in [1:10] do
+  IO.println s!"x: {x}"
+
+#eval ex3
+
+def ex4 : IO Unit := do
+IO.println "example 4"
+for x in [1:10:3] do
+  IO.println s!"x: {x}"
+
+#eval ex4