use pure iterators in examples

Author: datokrat

Iterator.lean

@@ -1,6 +1,7 @@
 -- This module serves as the root of the `Iterator` library.
 -- Import modules here that should be built as part of the library.
 import Iterator.Basic
+import Iterator.Pure
 import Iterator.Producers
 import Iterator.Combinators
 import Iterator.Consumers

Iterator/Combinators/Drop.lean

@@ -143,6 +143,6 @@ instance Drop.instFinitenessRelation [Iterator α m β] [Monad m] [Finite α m]
 instance Drop.instIteratorToArray [Monad m] [Iterator α m β] [Finite α m] : IteratorToArray (Drop α m β) m :=
   .defaultImplementation
 
-instance Drop.instIteratorFor [Monad m] [Monad n] [MonadLiftT m n] [Iterator α m β] [Finite α m] :
+instance Drop.instIteratorFor [Monad m] [Monad n] [Iterator α m β] [Finite α m] :
     IteratorFor (Drop α m β) m n :=
   .defaultImplementation

Iterator/Combinators/FilterMap.lean

@@ -176,7 +176,7 @@ instance {f : β → HetT m (Option γ)} [Finite α m] :
     IteratorToArray (FilterMapMH α f) m :=
   .defaultImplementation
 
-instance FilterMapMH.instIteratorFor [Monad m] [Monad n] [MonadLiftT m n]
+instance FilterMapMH.instIteratorFor [Monad m] [Monad n]
     [Iterator α m β] [Finite α m] :
     IteratorFor (FilterMapMH α f) m n :=
   .defaultImplementation
@@ -192,7 +192,7 @@ instance {f : β → HetT m γ} [IteratorToArray α m] :
       (fun x => do g ((← (f x).computation).inflate (small := _)))
       it.inner.inner
 
-instance MapMH.instIteratorFor [Monad m] [Monad n] [MonadLiftT m n]
+instance MapMH.instIteratorFor [Monad m] [Monad n]
     [Iterator α m β] [Finite α m] :
     IteratorFor (MapMH α f) m n :=
   .defaultImplementation

Iterator/Combinators/FlatMap.lean

@@ -340,7 +340,7 @@ instance FlatMap.instIteratorToArray [Monad m] [Iterator α₂ m γ] [Finite α
     IteratorToArray (FlatMap α f) m :=
   .defaultImplementation
 
-instance FlatMap.instIteratorFor [Monad m] [Monad n] [MonadLiftT m n] [Iterator α₂ m γ]
+instance FlatMap.instIteratorFor [Monad m] [Monad n] [Iterator α₂ m γ]
     [Finite α m] [Finite α₂ m] :
     IteratorFor (FlatMap α f) m n :=
   .defaultImplementation
@@ -453,7 +453,7 @@ instance SigmaIterator.instIteratorToArray [Monad m] [∀ b, Iterator (α b) m 
     IteratorToArray (SigmaIterator α m γ) m :=
   .defaultImplementation
 
-instance SigmaIterator.instIteratorFor [Monad m] [Monad n] [MonadLiftT m n]
+instance SigmaIterator.instIteratorFor [Monad m] [Monad n]
     [∀ b, Iterator (α b) m γ] [∀ b, Finite (α b) m] :
     IteratorFor (SigmaIterator α m γ) m n :=
   .defaultImplementation

Iterator/Combinators/MonadLift.lean

@@ -52,7 +52,7 @@ instance {_ : Iterator α m β} [Finite α m] {_ : MonadLiftT m n} : FinitenessR
     case done h =>
       cases h
 
-instance MonadLiftIterator.instIteratorFor [Monad n] [Monad n'] [MonadLiftT n n']
+instance MonadLiftIterator.instIteratorFor [Monad n] [Monad n']
     [Iterator α m β] [Finite α m] {_ : MonadLiftT m n} :
     IteratorFor (MonadLiftIterator α m n) n n' :=
   .defaultImplementation

Iterator/Combinators/Take.lean

@@ -168,7 +168,7 @@ instance Take.instFinitenessRelation [Monad m] [Iterator α m β] [Productive α
 instance Take.instIteratorToArray [Monad m] [Iterator α m β] [Productive α m] : IteratorToArray (Take α m β) m :=
   .defaultImplementation
 
-instance Take.instIteratorFor [Monad m] [Monad n] [MonadLiftT m n] [Iterator α m β] [Finite α m] :
+instance Take.instIteratorFor [Monad m] [Monad n] [Iterator α m β] [Finite α m] :
     IteratorFor (Take α m β) m n :=
   .defaultImplementation
   -- TODO: use [IteratorFor α m n]

Iterator/Combinators/Zip.lean

@@ -304,7 +304,7 @@ instance [Monad m] [Productive α₁ m] [Finite α₂ m] :
 instance ZipH.instIteratorToArray [Monad m] [Finite (ZipH α₁ m α₂ β₂) m] : IteratorToArray (ZipH α₁ m α₂ β₂) m :=
   .defaultImplementation
 
-instance ZipH.instIteratorFor [Monad m] [Monad n] [MonadLiftT m n] [Finite (ZipH α₁ m α₂ β₂) m] :
+instance ZipH.instIteratorFor [Monad m] [Monad n] [Finite (ZipH α₁ m α₂ β₂) m] :
     IteratorFor (ZipH α₁ m α₂ β₂) m n :=
   .defaultImplementation
 

Iterator/Consumers/Loop.lean

@@ -1,50 +1,56 @@
 import Iterator.Basic
 
 @[specialize]
-def IterM.DefaultConsumers.forIn {m : Type w → Type w'} {n : Type w → Type w''} [Monad n] [MonadLiftT m n]
+def IterM.DefaultConsumers.forIn {m : Type w → Type w'}
+    {n : Type w → Type w''} [Monad n] (lift : ∀ γ, m γ → n γ)
     {α : Type w} {β : Type v} {γ : Type w}
     [Iterator α m β] [Finite α m]
     (it : IterM (α := α) m β) (init : γ)
-    (f : (b : β) → (c : γ) → n (ForInStep γ)) : n γ := do
-  match (← it.stepH).inflate with
-  | .yield it' out _ =>
-    match ← f out init with
-    | .yield c => IterM.DefaultConsumers.forIn it' c f
-    | .done c => return c
-  | .skip it' _ => IterM.DefaultConsumers.forIn it' init f
-  | .done _ => return init
+    (f : (b : β) → (c : γ) → n (ForInStep γ)) : n γ :=
+  letI : MonadLift m n := ⟨fun {γ} => lift γ⟩
+  do
+    match (← it.stepH).inflate with
+    | .yield it' out _ =>
+      match ← f out init with
+      | .yield c => IterM.DefaultConsumers.forIn lift it' c f
+      | .done c => return c
+    | .skip it' _ => IterM.DefaultConsumers.forIn lift it' init f
+    | .done _ => return init
 termination_by it.finitelyManySteps
 
-class IteratorFor (α : Type w) (m : Type w → Type w') {β : Type v} [Iterator α m β] (n : Type w → Type w'') where
-  forIn : ∀ {γ : Type w}, IterM (α := α) m β → γ →
-      ((b : β) → (c : γ) → n (ForInStep γ)) →
-      n γ
+class IteratorFor (α : Type w) (m : Type w → Type w') {β : Type v} [Iterator α m β]
+    (n : Type w → Type w'') where
+  forIn : ∀ (_lift : (γ : Type w) → m γ → n γ) {γ : Type w},
+      IterM (α := α) m β → γ →
+      ((b : β) → (c : γ) → n (ForInStep γ)) → n γ
 
 class LawfulIteratorFor (α : Type w) (m : Type w → Type w') (n : Type w → Type w'')
-    [Monad n] [Iterator α m β] [Finite α m] [IteratorFor α m n] [MonadLiftT m n] where
-  lawful : ∀ γ, IteratorFor.forIn (α := α) (m := m) (n := n) (γ := γ) =
-    IterM.DefaultConsumers.forIn (α := α) (m := m) (n := n) (γ := γ)
+    [Monad n] [Iterator α m β] [Finite α m] [IteratorFor α m n] where
+  lawful : ∀ lift γ, IteratorFor.forIn lift (α := α) (m := m) (γ := γ) =
+    IterM.DefaultConsumers.forIn (α := α) (m := m) (n := n) (lift := lift) (γ := γ)
 
 def IteratorFor.defaultImplementation {α : Type w} {m : Type w → Type w'} {n : Type w → Type w'}
-    [Monad m] [Monad n] [MonadLiftT m n] [Iterator α m β] [Finite α m] :
+    [Monad m] [Monad n] [Iterator α m β] [Finite α m] :
     IteratorFor α m n where
-  forIn := IterM.DefaultConsumers.forIn
+  forIn lift := IterM.DefaultConsumers.forIn lift
 
 instance (α : Type w) (m : Type w → Type w') (n : Type w → Type w')
-    [Monad m] [Monad n] [MonadLiftT m n] [Iterator α m β] [Finite α m] :
-    haveI : IteratorFor α m n := .defaultImplementation
+    [Monad m] [Monad n] [Iterator α m β] [Finite α m] :
+    letI : IteratorFor α m n := .defaultImplementation
     LawfulIteratorFor α m n :=
   letI : IteratorFor α m n := .defaultImplementation
-  ⟨fun _ => rfl⟩
+  ⟨fun _ _ => rfl⟩
 
 instance {m : Type w → Type w'} {n : Type w → Type w''}
-    {α : Type w} {β : Type v} [Iterator α m β] [IteratorFor α m n] :
+    {α : Type w} {β : Type v} [Iterator α m β] [IteratorFor α m n]
+    [MonadLiftT m n] :
     ForIn n (IterM (α := α) m β) β where
-  forIn := IteratorFor.forIn
+  forIn := IteratorFor.forIn (fun _ => MonadLiftT.monadLift)
 
 @[specialize]
 def IterM.foldM {m : Type w → Type w'} {n : Type w → Type w'} [Monad n]
     {α : Type w} {β : Type v} {γ : Type w} [Iterator α m β] [IteratorFor α m n]
+    [MonadLiftT m n]
     (f : γ → β → n γ) (init : γ) (it : IterM (α := α) m β) : n γ :=
   ForIn.forIn it init (fun x acc => ForInStep.yield <$> f acc x)
 

Iterator/Examples.lean

@@ -22,19 +22,18 @@ end Primes
 
 def hideEggs : IO Unit := do
   -- Repeat colors and locations indefinitely
-  let colors := IterM.unfold Id (0 : Nat) (· + 1) |>.map fun n =>
+  let colors := Iter.unfold (0 : Nat) (· + 1) |>.map fun n =>
     #["green", "red", "yellow"][n % 3]!
-  let locations := IterM.unfold Id (0 : Nat) (· + 1) |>.map fun n =>
+  let locations := Iter.unfold (0 : Nat) (· + 1) |>.map fun n =>
     #["in a boot", "underneath a lampshade", "under the cushion", "in the lawn"][n % 4]!
   -- Summon the chickens
-  let chickens := ["Clucky", "Patches", "Fluffy"].iter Id
+  let chickens := ["Clucky", "Patches", "Fluffy"].iter
   -- Gather, color and hide the eggs
-  let eggs := chickens.flatMap (fun x : String => IterM.unfold Id x id |>.take 3)
+  let eggs := chickens.flatMap (fun x : String => Iter.unfold x id |>.take 3)
     |>.zip colors
     |>.zip locations
   -- Report the results (top secret)
-  let eggsIO := eggs.switchMonad (pure : ∀ {γ}, γ → IO γ) -- Obtain an IO-monadic iterator
-  for x in eggsIO do
+  for x in eggs do
     IO.println s! "{x.1.1} hides a {x.1.2} egg {x.2}."
   -- Alternative : eggsIO.mapM (fun x => IO.println s!"{x.1.1} hides a {x.1.2} egg {x.2}.") |>.drain
 
@@ -60,7 +59,7 @@ def firstOfEach (l : List (List Nat)) : List Nat :=
 #eval! firstOfEach [[1, 2], [], [3, 4]]
 
 def deepSum (l : List (List Nat)) : Nat :=
-  go (l.iter Id |>.flatMap fun l' => l'.iter Id) 0
+  go (l.iter |>.flatMap fun l' => l'.iter) 0
 where
   @[specialize]
   go it acc :=
@@ -107,15 +106,15 @@ Additional diagnostic information may be available using the `set_option diagnos
 -- as Rust
 
 def dependentFlatMap (l : List (List Nat)) : List Nat :=
-  let it := IterM.unfold Id 2 (· + 1) |>.zip l.iter
-  it.flatMapD (fun (x : Nat × List Nat) => (x.2.iter Id).filter fun y => y % x.1 = 0) |>.toList
+  let it := Iter.unfold 2 (· + 1) |>.zip l.iter
+  it.flatMapD (fun (x : Nat × List Nat) => x.2.iter.filter fun y => y % x.1 = 0) |>.toList
 
 /-- info: [2, 6, 9] -/
 #guard_msgs in
 #eval! dependentFlatMap [[1, 2, 3], [4, 5, 6, 9]]
 
 def addIndices (l : List Nat) : List (Nat × Nat) :=
-  IterM.unfold Id 0 (· + 1) |>.zip (l.iter Id) |>.toList
+  Iter.unfold 0 (· + 1) |>.zip l.iter |>.toList
 
 /-- info: [(0, 3), (1, 1), (2, 4), (3, 1), (4, 5), (5, 9)] -/
 #guard_msgs in
@@ -126,7 +125,7 @@ def addIndices (l : List Nat) : List (Nat × Nat) :=
 #eval! [3, 1, 4, 1, 5, 9].iter.filter (· % 2 = 0) |>.toList
 
 def printNums (l : List Nat) : IO Unit :=
-  l.iter IO |>.mapM (fun n => do IO.println (toString n); pure n) |>.drain
+  l.iterM IO |>.mapM (fun n => do IO.println (toString n); pure n) |>.drain
 
 
 /--
@@ -139,7 +138,7 @@ info: 1
 
 -- Same result as `printNums` but showcasing that we can use iterators in `for-in` constructs
 def printNumsForIn (l : List Nat) : IO Unit := do
-  for n in l.iter IO do
+  for n in l.iter do
     IO.println n
 
 def timestamps (n : Nat) : IO (List Std.Time.PlainTime) := do
@@ -162,7 +161,7 @@ def timestamps (n : Nat) : IO (List Std.Time.PlainTime) := do
 
 -- This example demonstrates that chained `mapM` calls are executed in a different order than with `List.mapM`.
 def chainedMapM (l : List Nat) : IO Unit :=
-  l.iter IO |>.mapM (IO.println <| s!"1st {.}") |>.mapM (IO.println <| s!"2nd {·}") |>.drain
+  l.iterM IO |>.mapM (IO.println <| s!"1st {.}") |>.mapM (IO.println <| s!"2nd {·}") |>.drain
 
 /--
 info: 1st 1

Iterator/Producers.lean

@@ -7,6 +7,7 @@ prelude
 import Iterator.Basic
 import Iterator.Consumers.Collect
 import Iterator.Consumers.Loop
+import Iterator.Pure.Basic
 import Init.Data.Nat.Lemmas
 
 section ListIterator
@@ -21,10 +22,15 @@ Returns a finite iterator for the given list.
 The iterator yields the elements of the list in order and then terminates.
 -/
 @[always_inline, inline]
-def List.iter {α : Type w} (l : List α) (m : Type w → Type w' := by exact Id) [Pure m] :
+def List.iterM {α : Type w} (l : List α) (m : Type w → Type w') [Pure m] :
     IterM (α := ListIterator α) m α :=
   toIter { list := l } m α
 
+@[always_inline, inline]
+def List.iter {α : Type w} (l : List α) :
+    Iter (α := ListIterator α) α :=
+  ((l.iterM Id).toPureIter : Iter α)
+
 instance {α : Type w} [Pure m] : Iterator (ListIterator α) m α where
   plausible_step it
     | .yield it' out => it.inner.list = out :: it'.inner.list
@@ -46,7 +52,7 @@ instance [Pure m] : FinitenessRelation (ListIterator α) m where
 instance {α : Type w} [Monad m] : IteratorToArray (ListIterator α) m :=
   .defaultImplementation
 
-instance {α : Type w} [Monad m] [Monad n] [MonadLiftT m n] : IteratorFor (ListIterator α) m n :=
+instance {α : Type w} [Monad m] [Monad n] : IteratorFor (ListIterator α) m n :=
   .defaultImplementation
 
 end ListIterator
@@ -72,10 +78,14 @@ instance [Pure m] : Iterator (UnfoldIterator α f) m α where
 Creates an infinite, productive iterator. First it yields `init`.
 If the last step yielded `a`, the next will yield `f a`.
 -/
-@[inline]
+@[always_inline, inline]
 def IterM.unfold (m : Type w → Type w') [Pure m] {α : Type w} (init : α) (f : α → α) :=
   toIter (UnfoldIterator.mk (f := f) init) m α
 
+@[always_inline, inline]
+def Iter.unfold {α : Type w} (init : α) (f : α → α) :=
+  ((IterM.unfold Id init f).toPureIter : Iter α)
+
 instance [Pure m] : ProductivenessRelation (UnfoldIterator α f) m where
   rel := emptyWf.rel
   wf := emptyWf.wf