feat: polymorphic ranges

src/Init/Data.lean

@@ -47,3 +47,4 @@ import Init.Data.Function
 import Init.Data.RArray
 import Init.Data.Vector
 import Init.Data.Iterators
+import Init.Data.Range.Polymorphic

src/Init/Data/Iterators/Basic.lean

@@ -354,7 +354,7 @@ Makes a single step with the given iterator `it`, potentially emitting a value a
 succeeding iterator. If this function is used recursively, termination can sometimes be proved with
 the termination measures `it.finitelyManySteps` and `it.finitelyManySkips`.
 -/
-@[always_inline, inline]
+@[always_inline, inline, expose]
 def IterM.step {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
     (it : IterM (α := α) m β) : m it.Step :=
   Iterator.step it
@@ -383,6 +383,24 @@ inductive IterM.IsPlausibleIndirectOutput {α β : Type w} {m : Type w → Type
   | indirect {it it' : IterM (α := α) m β} {out : β} : it'.IsPlausibleSuccessorOf it →
       it'.IsPlausibleIndirectOutput out → it.IsPlausibleIndirectOutput out
 
+/--
+Asserts that an iterator `it'` could plausibly produce `it'` as a successor iterator after
+finitely many steps. This relation is reflexive.
+-/
+inductive IterM.IsPlausibleIndirectSuccessorOf {α β : Type w} {m : Type w → Type w'}
+    [Iterator α m β] : IterM (α := α) m β → IterM (α := α) m β → Prop where
+  | refl (it : IterM (α := α) m β) : it.IsPlausibleIndirectSuccessorOf it
+  | cons_right {it'' it' it : IterM (α := α) m β} (h' : it''.IsPlausibleIndirectSuccessorOf it')
+      (h : it'.IsPlausibleSuccessorOf it) : it''.IsPlausibleIndirectSuccessorOf it
+
+theorem IterM.IsPlausibleIndirectSuccessorOf.trans {α β : Type w} {m : Type w → Type w'}
+    [Iterator α m β] {it'' it' it : IterM (α := α) m β}
+    (h' : it''.IsPlausibleIndirectSuccessorOf it') (h : it'.IsPlausibleIndirectSuccessorOf it) :
+    it''.IsPlausibleIndirectSuccessorOf it := by
+  induction h
+  case refl => exact h'
+  case cons_right ih => exact IsPlausibleIndirectSuccessorOf.cons_right ih ‹_›
+
 /--
 The type of the step object returned by `Iter.step`, containing an `IterStep`
 and a proof that this is a plausible step for the given iterator.
@@ -431,6 +449,16 @@ def Iter.IsPlausibleOutput {α : Type w} {β : Type w} [Iterator α Id β]
     (it : Iter (α := α) β) (out : β) : Prop :=
   it.toIterM.IsPlausibleOutput out
 
+theorem Iter.isPlausibleOutput_iff_exists {α : Type w} {β : Type w} [Iterator α Id β]
+    {it : Iter (α := α) β} {out : β} :
+    it.IsPlausibleOutput out ↔ ∃ it', it.IsPlausibleStep (.yield it' out) := by
+  simp only [IsPlausibleOutput, IterM.IsPlausibleOutput]
+  constructor
+  · rintro ⟨it', h⟩
+    exact ⟨it'.toIter, h⟩
+  · rintro ⟨it', h⟩
+    exact ⟨it'.toIterM, h⟩
+
 /--
 Asserts that a certain iterator `it'` could plausibly be the directly succeeding iterator of another
 given iterator `it`.
@@ -440,6 +468,18 @@ def Iter.IsPlausibleSuccessorOf {α : Type w} {β : Type w} [Iterator α Id β]
     (it' it : Iter (α := α) β) : Prop :=
   it'.toIterM.IsPlausibleSuccessorOf it.toIterM
 
+theorem Iter.isPlausibleSuccessorOf_iff_exists {α : Type w} {β : Type w} [Iterator α Id β]
+    {it' it : Iter (α := α) β} :
+    it'.IsPlausibleSuccessorOf it ↔ ∃ step, step.successor = some it' ∧ it.IsPlausibleStep step := by
+  simp only [IsPlausibleSuccessorOf, IterM.IsPlausibleSuccessorOf]
+  constructor
+  · rintro ⟨step, h₁, h₂⟩
+    exact ⟨step.mapIterator IterM.toIter,
+      by cases step <;> simp_all [IterStep.successor, Iter.IsPlausibleStep]⟩
+  · rintro ⟨step, h₁, h₂⟩
+    exact ⟨step.mapIterator Iter.toIterM,
+      by cases step <;> simp_all [IterStep.successor, Iter.IsPlausibleStep]⟩
+
 /--
 Asserts that a certain iterator `it` could plausibly yield the value `out` after an arbitrary
 number of steps.
@@ -472,6 +512,45 @@ theorem Iter.isPlausibleIndirectOutput_iff_isPlausibleIndirectOutput_toIterM {α
       replace h : it'.toIter.IsPlausibleSuccessorOf it.toIter := h
       exact .indirect (α := α) h ih
 
+/--
+Asserts that an iterator `it'` could plausibly produce `it'` as a successor iterator after
+finitely many steps. This relation is reflexive.
+-/
+inductive Iter.IsPlausibleIndirectSuccessorOf {α : Type w} {β : Type w} [Iterator α Id β] :
+    Iter (α := α) β → Iter (α := α) β → Prop where
+  | refl (it : Iter (α := α) β) : IsPlausibleIndirectSuccessorOf it it
+  | cons_right {it'' it' it : Iter (α := α) β} (h' : it''.IsPlausibleIndirectSuccessorOf it')
+      (h : it'.IsPlausibleSuccessorOf it) : it''.IsPlausibleIndirectSuccessorOf it
+
+theorem Iter.isPlausibleIndirectSuccessor_iff_isPlausibleIndirectSuccessor_toIterM {α β : Type w}
+    [Iterator α Id β] {it' it : Iter (α := α) β} :
+    it'.IsPlausibleIndirectSuccessorOf it ↔ it'.toIterM.IsPlausibleIndirectSuccessorOf it.toIterM := by
+  constructor
+  · intro h
+    induction h with
+    | refl => exact .refl _
+    | cons_right _ h ih => exact .cons_right ih h
+  · intro h
+    rw [← Iter.toIter_toIterM (it := it), ← Iter.toIter_toIterM (it := it')]
+    generalize it.toIterM = it at ⊢ h
+    induction h with
+    | refl => exact .refl _
+    | cons_right _ h ih => exact .cons_right ih h
+
+theorem Iter.IsPlausibleIndirectSuccessorOf.trans {α : Type w} {β : Type w} [Iterator α Id β]
+    {it'' it' it : Iter (α := α) β} (h' : it''.IsPlausibleIndirectSuccessorOf it')
+    (h : it'.IsPlausibleIndirectSuccessorOf it) : it''.IsPlausibleIndirectSuccessorOf it := by
+  induction h
+  case refl => exact h'
+  case cons_right ih => exact IsPlausibleIndirectSuccessorOf.cons_right ih ‹_›
+
+theorem Iter.IsPlausibleIndirectOutput.trans {α : Type w} {β : Type w} [Iterator α Id β]
+    {it' it : Iter (α := α) β} {out : β} (h : it'.IsPlausibleIndirectSuccessorOf it)
+    (h' : it'.IsPlausibleIndirectOutput out) : it.IsPlausibleIndirectOutput out := by
+  induction h
+  case refl => exact h'
+  case cons_right ih => exact IsPlausibleIndirectOutput.indirect ‹_› ih
+
 /--
 Asserts that a certain iterator `it'` could plausibly be the directly succeeding iterator of another
 given iterator `it` while no value is emitted (see `IterStep.skip`).
@@ -687,6 +766,21 @@ instance [Iterator α m β] [Finite α m] : Productive α m where
 
 end Productive
 
+/--
+This typeclass characterizes iterators that have deterministic return values. This typeclass does
+*not* guarantee that there are no monadic side effects such as exceptions.
+
+General monadic iterators can be nondeterministic, so that `it.IsPlausibleStep step` will be true
+for no or more than one choice of `step`. This typeclass ensures that there is exactly one such
+choice.
+
+This is an experimental instance and it should not be explicitly used downstream of the standard
+library.
+-/
+class LawfulDeterministicIterator (α : Type w) (m : Type w → Type w') [Iterator α m β]
+    where
+  isPlausibleStep_eq_eq : ∀ it : IterM (α := α) m β, ∃ step, it.IsPlausibleStep = (· = step)
+
 end Iterators
 
 export Iterators (Iter IterM)

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

@@ -7,6 +7,8 @@ module
 
 prelude
 import Init.Data.Iterators.Consumers.Partial
+import Init.Data.Iterators.Consumers.Loop
+import Init.Data.Iterators.Consumers.Monadic.Access
 
 namespace Std.Iterators
 
@@ -51,4 +53,12 @@ partial def Iter.Partial.atIdxSlow? {α β} [Iterator α Id β] [Monad Id]
   | .skip it' _ => (⟨it'⟩ : Iter.Partial (α := α) β).atIdxSlow? n
   | .done _ => none
 
+@[always_inline, inline, inherit_doc IterM.atIdx?]
+def Iter.atIdx? {α β} [Iterator α Id β] [Productive α Id] [IteratorAccess α Id]
+    (n : Nat) (it : Iter (α := α) β) : Option β :=
+  match (IteratorAccess.nextAtIdx? it.toIterM n).run.val with
+  | .yield _ out => some out
+  | .skip _ => none
+  | .done => none
+
 end Std.Iterators

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

@@ -56,18 +56,4 @@ def Iter.Partial.toList {α : Type w} {β : Type w}
     [Iterator α Id β] [IteratorCollectPartial α Id Id] (it : Iter.Partial (α := α) β) : List β :=
   it.it.toIterM.allowNontermination.toList.run
 
-/--
-This class charaterizes how the plausibility behavior (`IsPlausibleStep`) and the actual iteration
-behavior (`it.step`) should relate to each other for pure iterators. Intuitively, a step should
-only be plausible if it is possible. For simplicity's sake, the actual definition is weaker but
-presupposes that the iterator is finite.
-
-This is an experimental instance and it should not be explicitly used downstream of the standard
-library.
--/
-class LawfulPureIterator (α : Type w) [Iterator α Id β]
-    [Finite α Id] [IteratorCollect α Id Id] where
-  mem_toList_iff_isPlausibleIndirectOutput {it : Iter (α := α) β} {out : β} :
-    out ∈ it.toList ↔ it.IsPlausibleIndirectOutput out
-
 end Std.Iterators

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

@@ -6,6 +6,7 @@ Authors: Paul Reichert
 module
 
 prelude
+import Init.Data.Iterators.Consumers.Collect
 import Init.Data.Iterators.Consumers.Monadic.Loop
 import Init.Data.Iterators.Consumers.Partial
 
@@ -29,6 +30,7 @@ A `ForIn'` instance for iterators. Its generic membership relation is not easy t
 so this is not marked as `instance`. This way, more convenient instances can be built on top of it
 or future library improvements will make it more comfortable.
 -/
+@[always_inline, inline]
 def Iter.instForIn' {α : Type w} {β : Type w} {n : Type w → Type w'} [Monad n]
     [Iterator α Id β] [Finite α Id] [IteratorLoop α Id n] :
     ForIn' n (Iter (α := α) β) β ⟨fun it out => it.IsPlausibleIndirectOutput out⟩ where
@@ -136,4 +138,18 @@ def Iter.Partial.size {α : Type w} {β : Type w} [Iterator α Id β] [IteratorS
     (it : Iter (α := α) β) : Nat :=
   (IteratorSizePartial.size it.toIterM).run.down
 
+/--
+`LawfulIteratorSize α m` ensures that the `size` function of an iterator behaves as if it
+iterated over the whole iterator, counting its elements and causing all the monadic side effects
+of the iterations. This is a fairly strong condition for monadic iterators and it will be false
+for many efficient implementations of `size` that compute the size without actually iterating.
+
+This class is experimental and users of the iterator API should not explicitly depend on it.
+-/
+class LawfulIteratorSize (α : Type w) {β : Type w} [Iterator α Id β] [Finite α Id]
+    [IteratorSize α Id] where
+    size_eq_size_toArray {it : Iter (α := α) β} : it.size =
+      haveI : IteratorCollect α Id Id := .defaultImplementation
+      it.toArray.size
+
 end Std.Iterators

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

@@ -6,6 +6,7 @@ Authors: Paul Reichert
 module
 
 prelude
+import Init.Data.Iterators.Consumers.Monadic.Access
 import Init.Data.Iterators.Consumers.Monadic.Collect
 import Init.Data.Iterators.Consumers.Monadic.Loop
 import Init.Data.Iterators.Consumers.Monadic.Partial

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

@@ -0,0 +1,95 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Paul Reichert
+-/
+module
+
+prelude
+import Init.Data.Iterators.Basic
+
+namespace Std.Iterators
+
+/--
+`it.IsPlausibleNthOutputStep n step` is the proposition that according to the
+`IsPlausibleStep` relation, it is plausible that `step` returns the step in which the `n`-th value
+of `it` is emitted, or `.done` if `it` can plausibly terminate before emitting `n` values.
+-/
+inductive IterM.IsPlausibleNthOutputStep {α β : Type w} {m : Type w → Type w'} [Iterator α m β] :
+    Nat → IterM (α := α) m β → IterStep (IterM (α := α) m β) β → Prop where
+  /-- If `it` plausibly yields in its immediate next step, this step is a plausible `0`-th output step. -/
+  | zero_yield {it : IterM (α := α) m β} : it.IsPlausibleStep (.yield it' out) →
+      it.IsPlausibleNthOutputStep 0 (.yield it' out)
+  /--
+  If `it` plausibly terminates in its immediate next step (`.done`), then `.done` is a plausible
+  `n`-th output step for arbitrary `n`.
+  -/
+  | done {it : IterM (α := α) m β} : it.IsPlausibleStep .done →
+      it.IsPlausibleNthOutputStep n .done
+  /--
+  If `it` plausibly yields in its immediate next step, the successor iterator being `it'`, and
+  if `step` is a plausible `n`-th output step of `it'`, then `step` is a plausible `n + 1`-th
+  output step of `it`.
+  -/
+  | yield {it it' : IterM (α := α) m β} {out step} : it.IsPlausibleStep (.yield it' out) →
+      it'.IsPlausibleNthOutputStep n step → it.IsPlausibleNthOutputStep (n + 1) step
+  /--
+  If `it` plausibly skips in its immediate next step, the successor iterator being `it'`, and
+  if `step` is a plausible `n`-th output step of `it'`, then `step` is also a plausible `n`-th
+  output step of `it`.
+  -/
+  | skip {it it' : IterM (α := α) m β} {step} : it.IsPlausibleStep (.skip it') →
+      it'.IsPlausibleNthOutputStep n step → it.IsPlausibleNthOutputStep n step
+
+/--
+`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
+(necessarily of the form `.yield _ _`) or `.done` if `it` terminates before emitting the `n`-th
+value.
+
+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 class is experimental and users of the iterator API should not explicitly depend on it.
+-/
+class IteratorAccess (α : Type w) (m : Type w → Type w') {β : Type w} [Iterator α m β] where
+  nextAtIdx? (it : IterM (α := α) m β) (n : Nat) :
+    m (PlausibleIterStep (it.IsPlausibleNthOutputStep n))
+
+/--
+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.
+-/
+@[always_inline, inline]
+def IterM.nextAtIdx? [Iterator α m β] [IteratorAccess α m] (it : IterM (α := α) m β)
+    (n : Nat) : m (PlausibleIterStep (it.IsPlausibleNthOutputStep n)) :=
+  IteratorAccess.nextAtIdx? it n
+
+/--
+Returns the `n`-th value emitted by `it`, or `none` if `it` terminates earlier.
+
+For monadic iterators, the monadic effects of this operation may differ from manually iterating
+to the `n`-th value because `atIdx?` 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.
+-/
+@[always_inline, inline]
+def IterM.atIdx? [Iterator α m β] [IteratorAccess α m] [Monad m] (it : IterM (α := α) m β)
+    (n : Nat) : m (Option β) := do
+  match (← IteratorAccess.nextAtIdx? it n).val with
+  | .yield _ out => return some out
+  | .skip _ => return none
+  | .done => return none
+
+end Std.Iterators

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

@@ -105,12 +105,14 @@ class IteratorSizePartial (α : Type w) (m : Type w → Type w') {β : Type w} [
 
 end Typeclasses
 
-private def IteratorLoop.WFRel {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
+/-- Internal implementation detail of the iterator library. -/
+def IteratorLoop.WFRel {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
     {γ : Type x} {plausible_forInStep : β → γ → ForInStep γ → Prop}
     (_wf : WellFounded α m plausible_forInStep) :=
   IterM (α := α) m β × γ
 
-private def IteratorLoop.WFRel.mk {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
+/-- Internal implementation detail of the iterator library. -/
+def IteratorLoop.WFRel.mk {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
     {γ : Type x} {plausible_forInStep : β → γ → ForInStep γ → Prop}
     (wf : WellFounded α m plausible_forInStep) (it : IterM (α := α) m β) (c : γ) :
     IteratorLoop.WFRel wf :=
@@ -134,20 +136,21 @@ def IterM.DefaultConsumers.forIn' {m : Type w → Type w'} {α : Type w} {β : T
     (plausible_forInStep : β → γ → ForInStep γ → Prop)
     (wf : IteratorLoop.WellFounded α m plausible_forInStep)
     (it : IterM (α := α) m β) (init : γ)
-    (f : (b : β) → it.IsPlausibleIndirectOutput b → (c : γ) → n (Subtype (plausible_forInStep b c))) : n γ :=
+    (P : β → Prop) (hP : ∀ b, it.IsPlausibleIndirectOutput b → P b)
+    (f : (b : β) → P b → (c : γ) → n (Subtype (plausible_forInStep b c))) : n γ :=
   haveI : WellFounded _ := wf
   letI : MonadLift m n := ⟨fun {γ} => lift γ⟩
   do
     match ← it.step with
     | .yield it' out h =>
-      match ← f out (.direct ⟨_, h⟩) init with
+      match ← f out (hP _ <| .direct ⟨_, h⟩) init with
       | ⟨.yield c, _⟩ =>
-        IterM.DefaultConsumers.forIn' lift _ plausible_forInStep wf it' c
-          (fun out h' acc => f out (.indirect ⟨_, rfl, h⟩ h') acc)
+        IterM.DefaultConsumers.forIn' lift _ plausible_forInStep wf it' c P
+          (fun _ h' => hP _ <| .indirect ⟨_, rfl, h⟩ h') f
       | ⟨.done c, _⟩ => return c
     | .skip it' h =>
-      IterM.DefaultConsumers.forIn' lift _ plausible_forInStep wf it' init
-        (fun out h' acc => f out (.indirect ⟨_, rfl, h⟩ h') acc)
+      IterM.DefaultConsumers.forIn' lift _ plausible_forInStep wf it' init P
+          (fun _ h' => hP _ <| .indirect ⟨_, rfl, h⟩ h') f
     | .done _ => return init
 termination_by IteratorLoop.WFRel.mk wf it init
 decreasing_by
@@ -163,7 +166,7 @@ implementations are possible and should be used instead.
 def IteratorLoop.defaultImplementation {α : Type w} {m : Type w → Type w'} {n : Type w → Type w''}
     [Monad n] [Iterator α m β] :
     IteratorLoop α m n where
-  forIn lift := IterM.DefaultConsumers.forIn' lift
+  forIn lift γ Pl wf it init := IterM.DefaultConsumers.forIn' lift γ Pl wf it init _ (fun _ => id)
 
 /--
 Asserts that a given `IteratorLoop` instance is equal to `IteratorLoop.defaultImplementation`.
@@ -246,6 +249,7 @@ A `ForIn'` instance for iterators. Its generic membership relation is not easy t
 so this is not marked as `instance`. This way, more convenient instances can be built on top of it
 or future library improvements will make it more comfortable.
 -/
+@[always_inline, inline]
 def IterM.instForIn' {m : Type w → Type w'} {n : Type w → Type w''}
     {α : Type w} {β : Type w} [Iterator α m β] [Finite α m] [IteratorLoop α m n]
     [MonadLiftT m n] :