zip lemmas and getAtIdx? consumer

Author: datokrat

Iterator/Combinators/Monadic/Zip.lean

@@ -59,7 +59,7 @@ instance Zip.instIterator [Monad m] :
           pure <| .done (.doneRight hm hp)
 
 @[inline]
-def IterM.zip [Monad m]
+def IterM.zip
     (left : IterM (α := α₁) m β₁) (right : IterM (α := α₂) m β₂) :
     IterM (α := Zip α₁ m α₂ β₂) m (β₁ × β₂) :=
   toIter ⟨left, none, right⟩ m _

Iterator/Consumers.lean

@@ -5,6 +5,7 @@ Authors: Paul Reichert
 -/
 prelude
 import Iterator.Consumers.Monadic
+import Iterator.Consumers.Access
 import Iterator.Consumers.Collect
 import Iterator.Consumers.Loop
 import Iterator.Consumers.Partial

Iterator/Consumers/Access.lean

@@ -0,0 +1,31 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Paul Reichert
+-/
+prelude
+import Iterator.Consumers.Monadic.Access
+import Iterator.Consumers.Partial
+
+@[specialize]
+def Iter.getAtIdx? {α β} [Iterator α Id β] [Productive α Id]
+    (n : Nat) (it : Iter (α := α) β) : Option β :=
+  match it.step with
+  | .yield it' out _ =>
+    match n with
+    | 0 => some out
+    | k + 1 => it'.getAtIdx? k
+  | .skip it' _ => it'.getAtIdx? n
+  | .done _ => none
+termination_by (n, it.finitelyManySkips)
+
+@[specialize]
+partial def Iter.Partial.getAtIdx? {α β} [Iterator α Id β] [Monad Id]
+    (n : Nat) (it : Iter.Partial (α := α) β) : Option β := do
+  match it.it.step with
+  | .yield it' out _ =>
+    match n with
+    | 0 => some out
+    | k + 1 => (⟨it'⟩ : Iter.Partial (α := α) β).getAtIdx? k
+  | .skip it' _ => (⟨it'⟩ : Iter.Partial (α := α) β).getAtIdx? n
+  | .done _ => none

Iterator/Consumers/Monadic/Access.lean

@@ -0,0 +1,31 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Paul Reichert
+-/
+prelude
+import Iterator.Basic
+import Iterator.Consumers.Monadic.Partial
+
+@[specialize]
+def IterM.getAtIdx? {α m β} [Iterator α m β] [Productive α m] [Monad m]
+    (n : Nat) (it : IterM (α := α) m β) : m (Option β) := do
+  match ← it.step with
+  | .yield it' out _ =>
+    match n with
+    | 0 => return some out
+    | k + 1 => it'.getAtIdx? k
+  | .skip it' _ => it'.getAtIdx? n
+  | .done _ => return none
+termination_by (n, it.finitelyManySkips)
+
+@[specialize]
+partial def IterM.Partial.getAtIdx? {α m β} [Iterator α m β] [Monad m]
+    (n : Nat) (it : IterM.Partial (α := α) m β) : m (Option β) := do
+  match ← it.it.step with
+  | .yield it' out _ =>
+    match n with
+    | 0 => return some out
+    | k + 1 => (⟨it'⟩ : IterM.Partial (α := α) m β).getAtIdx? k
+  | .skip it' _ => (⟨it'⟩ : IterM.Partial (α := α) m β).getAtIdx? n
+  | .done _ => return none

Iterator/Lemmas/Combinators/Take.lean

@@ -5,6 +5,7 @@ Authors: Paul Reichert
 -/
 prelude
 import Iterator.Combinators.Take
+import Iterator.Consumers.Access
 import Iterator.Lemmas.Monadic.Take
 import Iterator.Lemmas.Consumers
 
@@ -52,3 +53,26 @@ theorem Iter.toList_take_of_finite {α β} [Iterator α Id β] {n : Nat}
     · simp [ihy h]
     · simp [ihs h]
     · simp
+
+theorem Iter.getAtIdx?_take {α β}
+    [Iterator α Id β] [Productive α Id] {k l : Nat}
+    {it : Iter (α := α) β} :
+    (it.take k).getAtIdx? l = if l < k then it.getAtIdx? l else none := by
+  revert k
+  fun_induction it.getAtIdx? l
+  case case1 it it' out h h' =>
+    intro k
+    simp [getAtIdx?.eq_def (it := it.take k), getAtIdx?.eq_def (it := it), step_take, h']
+    cases k <;> simp
+  case case2 it it' out h h' l ih =>
+    intro k
+    simp [getAtIdx?.eq_def (it := it.take k), getAtIdx?.eq_def (it := it), step_take, h']
+    cases k <;> cases l <;> simp [ih]
+  case case3 l it it' h h' ih =>
+    intro k
+    simp [getAtIdx?.eq_def (it := it.take k), getAtIdx?.eq_def (it := it), step_take, h']
+    cases k <;> cases l <;> simp [ih]
+  case case4 l it h h' =>
+    intro k
+    simp only [getAtIdx?.eq_def (it := it.take k), getAtIdx?.eq_def (it := it), step_take, h']
+    cases k <;> cases l <;> simp