feat: lemmas about step size iterators

src/Init/Data/Iterators/Basic.lean

@@ -214,6 +214,12 @@ theorem Iter.toIterM_comp_toIter {α : Type w} {β : Type w} :
     Iter.toIterM ∘ IterM.toIter (α := α) (β := β) = id :=
   rfl
 
+theorem Iter.toIterM_inj :
+    Iter.toIterM it = Iter.toIterM it' ↔ it = it' := by
+  apply Iff.intro
+  · cases it; cases it'; simp [Iter.toIterM]
+  · simp +contextual
+
 section IterStep
 
 variable {α : Type u} {β : Type w}
@@ -297,6 +303,12 @@ theorem IterStep.mapIterator_comp {α' : Type u'} {α'' : Type u''}
   apply funext
   exact fun _ => mapIterator_mapIterator.symm
 
+theorem IterStep.mapIterator_inj {α' : Type u'} {f : α → α'} {s s'} (hf : ∀ a b, f a = f b → a = b) :
+    IterStep.mapIterator (β := β) f s = IterStep.mapIterator (β := β) f s' ↔ s = s' := by
+  replace hf (a b : α) : f a = f b ↔ a = b := ⟨hf a b, by simp +contextual⟩
+  simp only [mapIterator]
+  split <;> split <;> simp [hf]
+
 @[simp]
 theorem IterStep.mapIterator_id {step : IterStep α β} :
     step.mapIterator id = step := by

src/Init/Data/Iterators/Consumers/Access.lean

@@ -19,6 +19,45 @@ set_option linter.missingDocs true
 namespace Std
 open Std.Iterators
 
+/-- Characterizes the plausible step when advancing to the `n`-th output of a pure iterator. -/
+def Iter.IsPlausibleNthOutputStep {α β : Type w} [Iterator α Id β]
+    (n : Nat) (it : Iter (α := α) β) (step : IterStep (Iter (α := α) β) β) : Prop :=
+  it.toIterM.IsPlausibleNthOutputStep n (step.mapIterator Iter.toIterM)
+
+/--
+Returns the step in which `it` yields its `n`-th element, or `.done` if it terminates earlier.
+In contrast to `step`, this function will always return either `.yield` or `.done` but never a
+`.skip` step.
+
+For monadic iterators, the monadic effects of this operation may differ from manually iterating
+to the `n`-th value because `nextAtIdx?` can take shortcuts. By the signature, the return value
+is guaranteed to plausible in the sense of `IterM.IsPlausibleNthOutputStep`.
+
+This function is only available for iterators that explicitly support it by implementing
+the `IteratorAccess` typeclass.
+-/
+@[inline]
+def Iter.nextAtIdx? [Iterator α Id β] [IteratorAccess α Id] (it : Iter (α := α) β)
+    (n : Nat) : PlausibleIterStep (it.IsPlausibleNthOutputStep n) :=
+  let step := it.toIterM.nextAtIdx? n |>.run
+  ⟨step.val.mapIterator IterM.toIter, by simp [IsPlausibleNthOutputStep, step.property]⟩
+
+/--
+Slow version of `IterM.nextAtIdx?` that does not require an `IteratorAccess α m` instance.
+
+Returns the step in which `it` yields its `n`-th element, or `.done` if it terminates earlier.
+In contrast to `step`, this function will always return either `.yield` or `.done` but never a
+`.skip` step.
+
+This function terminates after finitely many steps.
+-/
+@[inline]
+def Iter.nextAtIdxSlow? [Iterator α Id β] [Productive α Id]
+    (it : Iter (α := α) β)
+    (n : Nat) : PlausibleIterStep (it.IsPlausibleNthOutputStep n) :=
+  let step := it.toIterM.nextAtIdxSlow? n |>.run
+  ⟨step.val.mapIterator IterM.toIter, by simp [IsPlausibleNthOutputStep, step.property]⟩
+
 /--
 If possible, takes `n` steps with the iterator `it` and
 returns the `n`-th emitted value, or `none` if `it` finished

src/Init/Data/Iterators/Consumers/Monadic/Access.lean

@@ -7,6 +7,9 @@ module
 
 prelude
 public import Init.Data.Iterators.Basic
+public import Init.WFExtrinsicFix
+import Init.RCases
+import Init.Data.Iterators.Lemmas.Monadic.Basic
 
 set_option linter.missingDocs true
 
@@ -46,7 +49,7 @@ inductive IterM.IsPlausibleNthOutputStep {α β : Type w} {m : Type w → Type w
   | skip {it it' : IterM (α := α) m β} {step} : it.IsPlausibleStep (.skip it') →
       it'.IsPlausibleNthOutputStep n step → it.IsPlausibleNthOutputStep n step
 
-theorem IterM.not_isPlausibleNthOutputStep_yield {α β : Type w} {m : Type w → Type w'} [Iterator α m β]
+theorem IterM.not_isPlausibleNthOutputStep_skip {α β : Type w} {m : Type w → Type w'} [Iterator α m β]
     {n : Nat} {it it' : IterM (α := α) m β} :
     ¬ it.IsPlausibleNthOutputStep n (.skip it') := by
   intro h
@@ -57,6 +60,89 @@ theorem IterM.not_isPlausibleNthOutputStep_yield {α β : Type w} {m : Type w 
   · simp_all
   · simp_all
 
+theorem IterM.isPlausibleNthOutputStep_trans_of_yield {α β : Type w} {m : Type w → Type w'}
+    [Iterator α m β] {k n} {it it' : IterM (α := α) m β} {out step}
+    (h : it.IsPlausibleNthOutputStep k (.yield it' out))
+    (h' : it'.IsPlausibleNthOutputStep n step) :
+    it.IsPlausibleNthOutputStep (n + k + 1) step := by
+  generalize hs : (IterStep.yield it' out) = s at h
+  induction h generalizing h' it' out
+  case zero_yield =>
+    cases hs
+    exact .yield ‹_› h'
+  case done => cases hs
+  case yield ih =>
+    cases hs
+    refine .yield ‹_› ?_
+    simp only [Nat.add_assoc] at ih
+    exact ih h' rfl
+  case skip ih =>
+    cases hs
+    refine .skip ‹_› ?_
+    apply ih h' rfl
+
+theorem IterM.isPlausibleNthOutputStep_trans_of_done {α β : Type w} {m : Type w → Type w'}
+    [Iterator α m β] {k n} {it : IterM (α := α) m β}
+    (h : it.IsPlausibleNthOutputStep k .done) (hle : k ≤ n) :
+    it.IsPlausibleNthOutputStep n .done := by
+  generalize hs : IterStep.done = s at h
+  induction h generalizing n
+  case zero_yield => cases hs
+  case yield ih =>
+    cases hs
+    obtain ⟨n, rfl⟩ := Nat.exists_eq_add_one_of_ne_zero (n := n) (Nat.ne_zero_of_lt (Nat.lt_of_add_one_le hle))
+    exact .yield ‹_› (ih (Nat.le_of_add_le_add_right hle) rfl)
+  case skip ih =>
+    cases hs
+    exact .skip ‹_› (ih hle rfl)
+  case done =>
+    cases hs
+    exact .done ‹_›
+
+theorem IterM.IsPlausibleNthOutputStep.unique [Iterator α Id β]
+    [LawfulDeterministicIterator α Id] {it : IterM (α := α) Id β} {s s'}
+    (hs : it.IsPlausibleNthOutputStep n s) (hs' : it.IsPlausibleNthOutputStep n s') :
+    s = s' := by
+  induction hs
+  case zero_yield h =>
+    match hs' with
+    | .zero_yield h' ..
+    | .skip h' ..
+    | .done h' =>
+      replace h' := h'.eq_step
+      rw [← h.eq_step] at h'
+      cases h'
+      all_goals simp
+  case done h =>
+    match hs' with
+    | .zero_yield h' ..
+    | .yield h' ..
+    | .skip h' .. =>
+      replace h := h.eq_step
+      replace h' := h'.eq_step
+      rw [← h] at h'
+      cases h'
+    | .done h' => simp
+  case yield h _ ih =>
+    match hs' with
+    | .yield h' ..
+    | .skip h' ..
+    | .done h' =>
+      replace h' := h'.eq_step
+      rw [← h.eq_step] at h'
+      cases h'
+      all_goals apply ih ‹_›
+  case skip h _ ih =>
+    match hs' with
+    | .zero_yield h' ..
+    | .yield h' ..
+    | .skip h' ..
+    | .done h' =>
+      replace h' := h'.eq_step
+      rw [← h.eq_step] at h'
+      cases h'
+      all_goals apply ih ‹_›
+
 /--
 `IteratorAccess α m` provides efficient implementations for random access or iterators that support
 it. `it.nextAtIdx? n` either returns the step in which the `n`th value of `it` is emitted
@@ -95,6 +181,53 @@ def IterM.nextAtIdx? [Iterator α m β] [IteratorAccess α m] (it : IterM (α :=
     (n : Nat) : m (PlausibleIterStep (it.IsPlausibleNthOutputStep n)) :=
   IteratorAccess.nextAtIdx? it n
 
+/--
+Slow version of `IterM.nextAtIdx?` that does not require an `IteratorAccess α m` instance.
+
+Returns the step in which `it` yields its `n`-th element, or `.done` if it terminates earlier.
+In contrast to `step`, this function will always return either `.yield` or `.done` but never a
+`.skip` step.
+
+This function terminates after finitely many steps.
+-/
+@[inline]
+def IterM.atIdxSlow? [Monad m] [Iterator α m β] [Productive α m]
+    (it' : IterM (α := α) m β)
+    (n' : Nat) : m (Option β) := do
+    match (← it'.step).inflate with
+    | .yield it'' out _ =>
+      match n' with
+      | 0 => return some out
+      | k + 1 => atIdxSlow? it'' k
+    | .skip it'' _ => atIdxSlow? it'' n'
+    | .done _ => return none
+  termination_by (n', it'.finitelyManySkips)
+
+/--
+Slow version of `IterM.nextAtIdx?` that does not require an `IteratorAccess α m` instance.
+
+Returns the step in which `it` yields its `n`-th element, or `.done` if it terminates earlier.
+In contrast to `step`, this function will always return either `.yield` or `.done` but never a
+`.skip` step.
+
+This function terminates after finitely many steps.
+-/
+@[inline]
+def IterM.nextAtIdxSlow? [Monad m] [Iterator α m β] [Productive α m]
+    (it : IterM (α := α) m β)
+    (n : Nat) : m (PlausibleIterStep (it.IsPlausibleNthOutputStep n)) :=
+  go it n (fun s => id)
+where
+  go [Productive α m] it' n' (h : ∀ s, it'.IsPlausibleNthOutputStep n' s → it.IsPlausibleNthOutputStep n s) := do
+    match (← it'.step).inflate with
+    | .yield it'' out hp =>
+      match n' with
+      | 0 => return .yield it'' out (h _ (.zero_yield hp))
+      | k + 1 => go it'' k (fun s hp' => h s (.yield hp hp'))
+    | .skip it'' hp => go it'' n' (fun s hp' => h s (.skip hp hp'))
+    | .done hp => return .done (h _ (.done hp))
+  termination_by (n', it'.finitelyManySkips)
+
 /--
 Returns the `n`-th value emitted by `it`, or `none` if `it` terminates earlier.
 

src/Init/Data/Iterators/Consumers/Stream.lean

@@ -25,6 +25,6 @@ instance {α β} [Iterator α Id β] [Productive α Id] [IteratorAccess α Id] :
   where finally
     case noskip =>
       revert h
-      exact IterM.not_isPlausibleNthOutputStep_yield
+      exact IterM.not_isPlausibleNthOutputStep_skip
 
 end Std

src/Init/Data/Iterators/Lemmas/Basic.lean

@@ -7,10 +7,11 @@ module
 
 prelude
 public import Init.Data.Iterators.Basic
+public import Init.RCases
 
 public section
 
-namespace Std
+namespace Std.Iter
 open Std.Iterators
 
 /--
@@ -19,7 +20,7 @@ iterator `it` to an element of `motive it` by defining `f it` in terms of the va
 the plausible successors of `it'.
 -/
 @[specialize]
-def Iter.inductSteps {α β} [Iterator α Id β] [Finite α Id]
+def inductSteps {α β} [Iterator α Id β] [Finite α Id]
   (motive : Iter (α := α) β → Sort x)
   (step : (it : Iter (α := α) β) →
     (ih_yield : ∀ {it' : Iter (α := α) β} {out : β},
@@ -38,7 +39,7 @@ iterator `it` to an element of `motive it` by defining `f it` in terms of the va
 the plausible skip successors of `it'.
 -/
 @[specialize]
-def Iter.inductSkips {α β} [Iterator α Id β] [Productive α Id]
+def inductSkips {α β} [Iterator α Id β] [Productive α Id]
   (motive : Iter (α := α) β → Sort x)
   (step : (it : Iter (α := α) β) →
     (ih_skip : ∀ {it' : Iter (α := α) β}, it.IsPlausibleStep (.skip it') → motive it') →
@@ -47,4 +48,22 @@ def Iter.inductSkips {α β} [Iterator α Id β] [Productive α Id]
   step it (fun {it'} _ => inductSkips motive step it')
 termination_by it.finitelyManySkips
 
-end Std
+-- Not a real instance because the discrimination key would be to indiscriminate.
+local instance [Iterator α Id β] [LawfulDeterministicIterator α Id] {it : Iter (α := α) β} :
+    Subsingleton it.Step where
+  allEq s s' := by
+    obtain ⟨s'', hs''⟩ := LawfulDeterministicIterator.isPlausibleStep_eq_eq it.toIterM
+    obtain ⟨s, hs⟩ := s
+    obtain ⟨s', hs'⟩ := s'
+    simp only [Iter.IsPlausibleStep, hs''] at hs hs'
+    have := hs.trans hs'.symm
+    rw [IterStep.mapIterator_inj (fun _ _ => toIterM_inj.mp) ] at this
+    rwa [Subtype.mk.injEq]
+
+theorem IsPlausibleStep.eq_step [Iterator α Id β] [LawfulDeterministicIterator α Id]
+    {it : Iter (α := α) β} {step} (h : it.IsPlausibleStep step) :
+    step = it.step.val := by
+  have : ⟨step, h⟩ = it.step := Subsingleton.allEq _ _
+  simp [← this]
+
+end Std.Iter