From 21d90afa432ad1a699cf9befe3b3912bb72f5192 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sat, 3 Oct 2020 15:16:45 -0700
Subject: [PATCH] feat: add `ForInStep` type

---
 src/Init/Core.lean              |  5 +++++
 src/Init/Data/List/Control.lean | 24 ++++++++++++++++++++++++
 tests/playground/forIn2.lean    | 30 ++++++++----------------------
 3 files changed, 37 insertions(+), 22 deletions(-)

diff --git a/src/Init/Core.lean b/src/Init/Core.lean
index 665c17f6ba..a12ead9280 100644
--- a/src/Init/Core.lean
+++ b/src/Init/Core.lean
@@ -273,6 +273,11 @@ structure Subtype {α : Sort u} (p : α → Prop) :=
 inductive Exists {α : Sort u} (p : α → Prop) : Prop
 | intro (w : α) (h : p w) : Exists
 
+/- Auxiliary type used to compiler `for x in xs` notation. -/
+inductive ForInStep (α : Type u)
+| done  : α → ForInStep
+| yield : α → ForInStep
+
 class inductive Decidable (p : Prop)
 | isFalse (h : ¬p) : Decidable
 | isTrue  (h : p) : Decidable
diff --git a/src/Init/Data/List/Control.lean b/src/Init/Data/List/Control.lean
index baf14b032c..386b2258c5 100644
--- a/src/Init/Data/List/Control.lean
+++ b/src/Init/Data/List/Control.lean
@@ -154,4 +154,28 @@ def findSomeM? {m : Type u → Type v} [Monad m] {α : Type w} {β : Type u} (f
   | some b => pure b
   | none   => findSomeM? as
 
+@[specialize]
+def forInAux {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : α → β → m (ForInStep β)) : List α → β → m β
+| [], b    => pure b
+| a::as, b => do
+  s ← f a b;
+  match s with
+  | ForInStep.done b  => pure b
+  | ForInStep.yield b => forInAux as b
+
+@[inline] def forIn {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : List α) (init : β) (f : α → β → m (ForInStep β)) : m β :=
+forInAux f as init
+
+@[specialize]
+def forInMapAux {α β : Type u} {m : Type u → Type v} [Monad m] (f : α → β → m (ForInStep (α × β))) : List α → List α → β →  m (List α × β)
+| [],    rs, b => pure (rs.reverse, b)
+| a::as, rs, b => do
+  s ← f a b;
+  match s with
+  | ForInStep.done (a, b)  => pure ((a :: rs).reverse ++ as, b)
+  | ForInStep.yield (a, b) => forInMapAux as (a::rs) b
+
+@[inline] def forInMap {α β : Type u} {m : Type u → Type v} [Monad m] (as : List α) (init : β) (f : α → β → m (ForInStep (α × β))) : m (List α × β) :=
+forInMapAux f as [] init
+
 end List
diff --git a/tests/playground/forIn2.lean b/tests/playground/forIn2.lean
index 3ef2138ac3..3850e047fa 100644
--- a/tests/playground/forIn2.lean
+++ b/tests/playground/forIn2.lean
@@ -1,45 +1,31 @@
 new_frontend
 
-inductive ForIn.Step.{u} (α : Type u)
-| done  : α → Step α
-| yield : α → Step α
-
-@[inline] def List.forIn {α β m} [Monad m] (as : List α) (init : β) (f : α → β → m (ForIn.Step β)) : m β :=
-let rec @[specialize] loop
-  | [], b    => pure b
-  | a::as, b => do
-    let s ← f a b
-    (match s with
-    | ForIn.Step.done b     => pure b
-    | ForIn.Step.yield b => loop as b)
-loop as init
-
 def tst1 : IO Nat :=
 [1, 2, 3, 4, 5, 7, 8, 9, 10, 11, 12, 14].forIn 0 fun a b =>
   if a % 2 == 0 then do
     IO.println (">> " ++ toString a ++ " " ++ toString b)
-    (if b > 20 then return ForIn.Step.done b
-     else return ForIn.Step.yield (a+b))
+    (if b > 20 then pure $ ForInStep.done b
+     else pure $ ForInStep.yield (a+b))
   else
-    return ForIn.Step.yield b
+    pure $ ForInStep.yield b
 
 #eval tst1
 
 structure Range :=
 (lower upper : Nat)
 
-def Range.forIn {β m} [Monad m] (range : Range) (init : β) (f : Nat → β → m (ForIn.Step β)) : m β :=
+def Range.forIn {β m} [Monad m] (range : Range) (init : β) (f : Nat → β → m (ForInStep β)) : m β :=
 let base := range.lower + range.upper - 2
 let rec @[specialize] loop
   | 0,   b => pure b
   | i+1, b =>
     let j := base - i
-    if j >= range.upper then return b
+    if j >= range.upper then pure b
     else do
       let s ← f j b
       (match s with
-      | ForIn.Step.done b     => return b
-      | ForIn.Step.yield b => loop i b)
+      | ForInStep.done b  => pure b
+      | ForInStep.yield b => loop i b)
 loop (range.upper - 1) init
 
 @[inline] def range (a : Nat) (b : Option Nat := none) : Range :=
@@ -53,6 +39,6 @@ instance : HasOfNat (Option Nat) :=
 def tst3 : IO Nat :=
 (range 5 10).forIn 0 fun i s => do
   IO.println (">> " ++ toString i)
-  return ForIn.Step.yield (s+i)
+  pure $ ForInStep.yield (s+i)
 
 #eval tst3