provide combinators and some consumers

Author: datokrat

Iterator/Combinators/Drop.lean

@@ -44,8 +44,7 @@ _TODO_: prove `Productive`
 Currently, this combinator incurs an additional O(1) cost with each output of `it`, even when the iterator
 does not drop any elements anymore.
 -/
-def IterM.drop [Iterator α m β]
-    (n : Nat) (it : IterM (α := α) m β) :=
+def IterM.drop (n : Nat) (it : IterM (α := α) m β) :=
   toIter (Drop.mk n it) m β
 
 inductive Drop.PlausibleStep [Iterator α m β] (it : IterM (α := Drop α m β) m β) :

Iterator/Combinators/FlatMap.lean

@@ -462,7 +462,6 @@ end Dependent
 
 section IterM
 
-set_option pp.universes true in
 @[always_inline, inline]
 def IterM.flatMapD {α : Type w} {β : Type v} {α₂ : β → Type w}
     {γ : Type x} {m : Type w → Type w'}

Iterator/Consumers/Loop.lean

@@ -4,20 +4,20 @@ import Iterator.Basic
 def IterM.DefaultConsumers.forIn {m : Type w → Type w'} {n : Type w → Type w''} [Monad n] [MonadLiftT m n]
     {α : Type w} {β : Type v} {γ : Type w}
     [Iterator α m β] [Finite α m]
-    (it : IterM (α := α) m β) (init : γ) {successor_of : β → γ → γ → Prop}
-    (f : (b : β) → (c : γ) → n (Subtype fun s : ForInStep γ => successor_of b c s.value)) : n γ := do
+    (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
+    | .yield c => IterM.DefaultConsumers.forIn it' c f
+    | .done c => return c
   | .skip it' _ => IterM.DefaultConsumers.forIn 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 β → γ → {successor_of : β → γ → γ → Prop} →
-      ((b : β) → (c : γ) → n (Subtype fun s : ForInStep γ => successor_of b c s.value)) →
+  forIn : ∀ {γ : Type w}, IterM (α := α) m β → γ →
+      ((b : β) → (c : γ) → n (ForInStep γ)) →
       n γ
 
 class LawfulIteratorFor (α : Type w) (m : Type w → Type w') (n : Type w → Type w'')
@@ -38,15 +38,13 @@ instance (α : Type w) (m : Type w → Type w') (n : Type w → Type w')
   ⟨fun _ => rfl⟩
 
 instance {m : Type w → Type w'} {n : Type w → Type w''}
-    {α : Type w} {β : Type v} [Iterator α m β] [Finite α m] [IteratorFor α m n] :
+    {α : Type w} {β : Type v} [Iterator α m β] [IteratorFor α m n] :
     ForIn n (IterM (α := α) m β) β where
-  forIn it init stepper := IteratorFor.forIn it init
-      (successor_of := fun _ _ _ => True)
-      (fun b c => (⟨·, True.intro⟩) <$> stepper b c)
+  forIn := IteratorFor.forIn
 
 @[specialize]
 def IterM.foldM {m : Type w → Type w'} {n : Type w → Type w'} [Monad n]
-    {α : Type w} {β : Type v} {γ : Type w} [Iterator α m β] [Finite α m] [IteratorFor α m n]
+    {α : Type w} {β : Type v} {γ : Type w} [Iterator α m β] [IteratorFor α m n]
     (f : γ → β → n γ) (init : γ) (it : IterM (α := α) m β) : n γ :=
   ForIn.forIn it init (fun x acc => ForInStep.yield <$> f acc x)
 

Iterator/Pure/Basic.lean

@@ -1,3 +1,4 @@
+prelude
 import Iterator.Basic
 
 structure Iter {α : Type w} (β : Type v) : Type w where
@@ -39,7 +40,7 @@ def Iter.plausible_skip_successor_of {α : Type w} {β : Type v} [Iterator α Id
   it.plausible_step (.skip it')
 
 @[always_inline, inline]
-def Iter.step {α : Type w} {β : Type w} [Iterator α Id β] [Functor Id]
+def Iter.step {α : Type w} {β : Type v} [Iterator α Id β]
     (it : Iter (α := α) β) : it.Step := by
   refine it.toIterM.stepH.run.inflate.map IterM.toPureIter id _ ?_
   intro step hp

Iterator/Pure/Combinators/Drop.lean

@@ -0,0 +1,8 @@
+prelude
+import Iterator.Combinators.Drop
+import Iterator.Pure.Basic
+
+@[always_inline, inline]
+def Iter.drop {α : Type w} {β : Type v} (n : Nat) (it : Iter (α := α) β) :
+    Iter (α := Drop α Id β) β :=
+  it.toIterM.drop n |>.toPureIter

Iterator/Pure/Combinators/FilterMap.lean

@@ -0,0 +1,10 @@
+prelude
+import Iterator.Combinators.FilterMap
+import Iterator.Pure.Basic
+
+-- TODO: consistently require or do not require `Iterator` instances
+-- in combinators
+@[always_inline, inline]
+def Iter.filterMap {α : Type w} {β : Type v} [Iterator α Id β]
+    (f : β → Option γ) (it : Iter (α := α) β) :=
+  ((it.toIterM.filterMap f).toPureIter : Iter γ)

Iterator/Pure/Combinators/FlatMap.lean

@@ -0,0 +1,17 @@
+prelude
+import Iterator.Combinators.FlatMap
+import Iterator.Pure.Basic
+
+@[always_inline, inline]
+def Iter.flatMap {α : Type w} {β : Type v} {α₂ : Type w} {γ : Type x}
+    [Iterator α Id β] [Iterator α₂ Id γ] (f : β → Iter (α := α₂) γ)
+    (it : Iter (α := α) β) :=
+  ((it.toIterM.flatMap (Iter.toIterM ∘ f)).toPureIter : Iter γ)
+
+
+@[always_inline, inline]
+def Iter.flatMapD {α : Type w} {β : Type v} {α₂ : β → Type w} {γ : Type x}
+    [Iterator α Id β] [∀ b, Iterator (α₂ b) Id γ] (f : (b : β) → Iter (α := α₂ b) γ)
+    (it : Iter (α := α) β) :=
+  -- TODO: confusing error message when writing (toIterM ∘ (f b)) (it's fine that it fails)
+  ((it.toIterM.flatMapD (fun b => (f b).toIterM)).toPureIter : Iter γ)

Iterator/Pure/Combinators/Take.lean

@@ -0,0 +1,8 @@
+prelude
+import Iterator.Combinators.Take
+import Iterator.Pure.Basic
+
+@[always_inline, inline]
+def Iter.take {α : Type w} {β : Type v} (n : Nat) (it : Iter (α := α) β) :
+    Iter (α := Take α Id β) β :=
+  it.toIterM.take n |>.toPureIter

Iterator/Pure/Combinators/Zip.lean

@@ -0,0 +1,9 @@
+prelude
+import Iterator.Combinators.Zip
+import Iterator.Pure.Basic
+
+@[always_inline, inline]
+def Iter.zip {α₁ : Type w} {β₁: Type v₁} {α₂ : Type w} {β₂ : Type v₂}
+    [Iterator α₁ Id β₁] [Iterator α₂ Id β₂]
+    (left : Iter (α := α₁) β₁) (right : Iter (α := α₂) β₂) :=
+  ((left.toIterM.zipH right.toIterM).toPureIter : Iter (β₁ × β₂))

Iterator/Pure/Consumers/Collect.lean

@@ -0,0 +1,17 @@
+import Iterator.Pure.Basic
+import Iterator.Consumers.Collect
+
+@[always_inline, inline]
+def Iter.toArray {α : Type w} {β : Type w}
+    [Iterator α Id β] [IteratorToArray α Id] (it : Iter (α := α) β) : Array β :=
+  it.toIterM.toArray
+
+@[always_inline, inline]
+def Iter.reverseToList {α : Type w} {β : Type w}
+    [Iterator α Id β] [Finite α Id] (it : Iter (α := α) β) : List β :=
+  it.toIterM.reverseToList
+
+@[always_inline, inline]
+def Iter.toList {α : Type w} {β : Type w}
+    [Iterator α Id β] [IteratorToArray α Id] (it : Iter (α := α) β) : List β :=
+  it.toIterM.toList