feat: polymorphic ranges and slices

src/Init/Data.lean

@@ -46,3 +46,4 @@ import Init.Data.NeZero
 import Init.Data.Function
 import Init.Data.RArray
 import Init.Data.Vector
+import Init.Data.Iterators

src/Init/Data/Array/Lex/Basic.lean

@@ -8,7 +8,7 @@ module
 prelude
 import Init.Data.Array.Basic
 import Init.Data.Nat.Lemmas
-import Init.Data.Range
+import Init.Data.Range.New.Nat
 
 set_option linter.listVariables true -- Enforce naming conventions for `List`/`Array`/`Vector` variables.
 set_option linter.indexVariables true -- Enforce naming conventions for index variables.
@@ -24,9 +24,9 @@ Specifically, `Array.lex as bs lt` is true if
 * there is an index `i` such that `lt as[i] bs[i]`, and for all `j < i`, `as[j] == bs[j]`.
 -/
 def lex [BEq α] (as bs : Array α) (lt : α → α → Bool := by exact (· < ·)) : Bool := Id.run do
-  for h : i in [0 : min as.size bs.size] do
+  for h : i in 0,,<(min as.size bs.size) do
     -- TODO: `omega` should be able to find this itself.
-    have : i < min as.size bs.size := Membership.get_elem_helper h rfl
+    have : i < min as.size bs.size := Range.get_elem_helper h rfl
     if lt as[i] bs[i] then
       return true
     else if as[i] != bs[i] then

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

@@ -3,7 +3,10 @@ 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 Std.Data.Iterators.Consumers.Monadic.Collect
-import Std.Data.Iterators.Consumers.Monadic.Loop
-import Std.Data.Iterators.Consumers.Monadic.Partial
+import Init.Data.Iterators.Basic
+import Init.Data.Iterators.PostconditionMonad
+import Init.Data.Iterators.Consumers
+import Init.Data.Iterators.Lemmas

src/Init/Data/Iterators/Consumers.lean

@@ -0,0 +1,13 @@
+/-
+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.Consumers.Monadic
+import Init.Data.Iterators.Consumers.Access
+import Init.Data.Iterators.Consumers.Collect
+import Init.Data.Iterators.Consumers.Loop
+import Init.Data.Iterators.Consumers.Partial

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

@@ -3,8 +3,10 @@ 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 Std.Data.Iterators.Consumers.Partial
+import Init.Data.Iterators.Consumers.Partial
 
 namespace Std.Iterators
 

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

@@ -3,10 +3,12 @@ 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 Std.Data.Iterators.Basic
-import Std.Data.Iterators.Consumers.Partial
-import Std.Data.Iterators.Consumers.Monadic.Collect
+import Init.Data.Iterators.Basic
+import Init.Data.Iterators.Consumers.Partial
+import Init.Data.Iterators.Consumers.Monadic.Collect
 
 /-!
 # Collectors

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

@@ -3,9 +3,11 @@ 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 Std.Data.Iterators.Consumers.Monadic.Loop
-import Std.Data.Iterators.Consumers.Partial
+import Init.Data.Iterators.Consumers.Monadic.Loop
+import Init.Data.Iterators.Consumers.Partial
 
 /-!
 # Loop consumers
@@ -24,9 +26,11 @@ namespace Std.Iterators
 
 instance (α : Type w) (β : Type w) (n : Type w → Type w') [Monad n]
     [Iterator α Id β] [Finite α Id] [IteratorLoop α Id n] :
-    ForIn n (Iter (α := α) β) β where
-  forIn it init f :=
-    IteratorLoop.finiteForIn (fun δ (c : Id δ) => pure c.run) |>.forIn it.toIterM init f
+    ForIn' n (Iter (α := α) β) β ⟨fun it out => it.IsPlausibleIndirectOutput out⟩ where
+  forIn' it init f :=
+    IteratorLoop.finiteForIn' (fun δ (c : Id δ) => pure c.run) |>.forIn' it.toIterM init
+        fun out h acc =>
+          f out (Iter.isPlausibleIndirectOutput_iff_isPlausibleIndirectOutput_toIterM.mpr h) acc
 
 instance (α : Type w) (β : Type w) (n : Type w → Type w') [Monad n]
     [Iterator α Id β] [IteratorLoopPartial α Id n] :
@@ -111,4 +115,14 @@ def Iter.Partial.fold {α : Type w} {β : Type w} {γ : Type w} [Iterator α Id
     (init : γ) (it : Iter.Partial (α := α) β) : γ :=
   ForIn.forIn (m := Id) it init (fun x acc => ForInStep.yield (f acc x))
 
+@[always_inline, inline, inherit_doc IterM.size]
+def Iter.size {α : Type w} {β : Type w} [Iterator α Id β] [IteratorSize α Id]
+    (it : Iter (α := α) β) : Nat :=
+  (IteratorSize.size it.toIterM).run.down
+
+@[always_inline, inline, inherit_doc IterM.Partial.size]
+def Iter.Partial.size {α : Type w} {β : Type w} [Iterator α Id β] [IteratorSizePartial α Id]
+    (it : Iter (α := α) β) : Nat :=
+  (IteratorSizePartial.size it.toIterM).run.down
+
 end Std.Iterators

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

@@ -0,0 +1,11 @@
+/-
+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.Consumers.Monadic.Collect
+import Init.Data.Iterators.Consumers.Monadic.Loop
+import Init.Data.Iterators.Consumers.Monadic.Partial

src/Std/Data/Iterators/Consumers/Monadic/Collect.lean → src/Init/Data/Iterators/Consumers/Monadic/Collect.lean

@@ -3,9 +3,11 @@ 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 Std.Data.Iterators.Consumers.Monadic.Partial
-import Std.Data.Internal.LawfulMonadLiftFunction
+import Init.Data.Iterators.Consumers.Monadic.Partial
+import Init.Data.Iterators.Internal.LawfulMonadLiftFunction
 
 /-!
 # Collectors

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

@@ -3,10 +3,12 @@ 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.RCases
-import Std.Data.Iterators.Basic
-import Std.Data.Iterators.Consumers.Monadic.Partial
+import Init.Data.Iterators.Basic
+import Init.Data.Iterators.Consumers.Monadic.Partial
 
 /-!
 # Loop-based consumers
@@ -62,8 +64,9 @@ class IteratorLoop (α : Type w) (m : Type w → Type w') {β : Type w} [Iterato
   forIn : ∀ (_lift : (γ : Type w) → m γ → n γ) (γ : Type w),
       (plausible_forInStep : β → γ → ForInStep γ → Prop) →
       IteratorLoop.WellFounded α m plausible_forInStep →
-      IterM (α := α) m β → γ →
-      ((b : β) → (c : γ) → n (Subtype (plausible_forInStep b c))) → n γ
+      (it : IterM (α := α) m β) → γ →
+      ((b : β) → it.IsPlausibleIndirectOutput b → (c : γ) → n (Subtype (plausible_forInStep b c))) →
+      n γ
 
 /--
 `IteratorLoopPartial α m` provides efficient implementations of loop-based consumers for `α`-based
@@ -76,7 +79,14 @@ provided by the standard library.
 class IteratorLoopPartial (α : Type w) (m : Type w → Type w') {β : Type w} [Iterator α m β]
     (n : Type w → Type w'') where
   forInPartial : ∀ (_lift : (γ : Type w) → m γ → n γ) {γ : Type w},
-      IterM (α := α) m β → γ → ((b : β) → (c : γ) → n (ForInStep γ)) → n γ
+      (it : IterM (α := α) m β) → γ →
+      ((b : β) → it.IsPlausibleIndirectOutput b → (c : γ) → n (ForInStep γ)) → n γ
+
+class IteratorSize (α : Type w) (m : Type w → Type w') {β : Type w} [Iterator α m β] where
+  size : IterM (α := α) m β → m (ULift Nat)
+
+class IteratorSizePartial (α : Type w) (m : Type w → Type w') {β : Type w} [Iterator α m β] where
+  size : IterM (α := α) m β → m (ULift Nat)
 
 end Typeclasses
 
@@ -91,7 +101,7 @@ private def IteratorLoop.WFRel.mk {α : Type w} {m : Type w → Type w'} {β : T
     IteratorLoop.WFRel wf :=
   (it, c)
 
-instance {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
+private instance {α : Type w} {m : Type w → Type w'} {β : Type w} [Iterator α m β]
     {γ : Type x} {plausible_forInStep : β → γ → ForInStep γ → Prop}
     (wf : IteratorLoop.WellFounded α m plausible_forInStep) :
     WellFoundedRelation (IteratorLoop.WFRel wf) where
@@ -109,16 +119,20 @@ def IterM.DefaultConsumers.forIn {m : Type w → Type w'} {α : Type w} {β : Ty
     (plausible_forInStep : β → γ → ForInStep γ → Prop)
     (wf : IteratorLoop.WellFounded α m plausible_forInStep)
     (it : IterM (α := α) m β) (init : γ)
-    (f : (b : β) → (c : γ) → n (Subtype (plausible_forInStep b c))) : n γ :=
+    (f : (b : β) → it.IsPlausibleIndirectOutput 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 _ =>
-      match ← f out init with
-      | ⟨.yield c, _⟩ => IterM.DefaultConsumers.forIn lift _ plausible_forInStep wf it' c f
+    | .yield it' out h =>
+      match ← f out (.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)
       | ⟨.done c, _⟩ => return c
-    | .skip it' _ => IterM.DefaultConsumers.forIn lift _ plausible_forInStep wf it' init f
+    | .skip it' h =>
+      IterM.DefaultConsumers.forIn lift _ plausible_forInStep wf it' init
+        (fun out h' acc => f out (.indirect ⟨_, rfl, h⟩ h') acc)
     | .done _ => return init
 termination_by IteratorLoop.WFRel.mk wf it init
 decreasing_by
@@ -153,15 +167,19 @@ partial def IterM.DefaultConsumers.forInPartial {m : Type w → Type w'} {α : T
     {n : Type w → Type w''} [Monad n]
     (lift : ∀ γ, m γ → n γ) (γ : Type w)
     (it : IterM (α := α) m β) (init : γ)
-    (f : (b : β) → (c : γ) → n (ForInStep γ)) : n γ :=
+    (f : (b : β) → it.IsPlausibleIndirectOutput b → (c : γ) → n (ForInStep γ)) : n γ :=
   letI : MonadLift m n := ⟨fun {γ} => lift γ⟩
   do
     match ← it.step with
-    | .yield it' out _ =>
-      match ← f out init with
-      | .yield c => IterM.DefaultConsumers.forInPartial lift _ it' c f
+    | .yield it' out h =>
+      match ← f out (.direct ⟨_, h⟩) init with
+      | .yield c =>
+        IterM.DefaultConsumers.forInPartial lift _ it' c
+          fun out h' acc => f out (.indirect ⟨_, rfl, h⟩ h') acc
       | .done c => return c
-    | .skip it' _ => IterM.DefaultConsumers.forInPartial lift _ it' init f
+    | .skip it' h =>
+      IterM.DefaultConsumers.forInPartial lift _ it' init
+        fun out h' acc => f out (.indirect ⟨_, rfl, h⟩ h') acc
     | .done _ => return init
 
 /--
@@ -196,27 +214,28 @@ theorem IteratorLoop.wellFounded_of_finite {m : Type w → Type w'}
     exact WellFoundedRelation.wf
 
 /--
-This `ForIn`-style loop construct traverses a finite iterator using an `IteratorLoop` instance.
+This `ForIn'`-style loop construct traverses a finite iterator using an `IteratorLoop` instance.
 -/
 @[always_inline, inline]
-def IteratorLoop.finiteForIn {m : Type w → Type w'} {n : Type w → Type w''}
+def IteratorLoop.finiteForIn' {m : Type w → Type w'} {n : Type w → Type w''}
     {α : Type w} {β : Type w} [Iterator α m β] [Finite α m] [IteratorLoop α m n]
     (lift : ∀ γ, m γ → n γ) :
-    ForIn n (IterM (α := α) m β) β where
-  forIn {γ} [Monad n] it init f :=
+    ForIn' n (IterM (α := α) m β) β ⟨fun it out => it.IsPlausibleIndirectOutput out⟩ where
+  forIn' {γ} [Monad n] it init f :=
     IteratorLoop.forIn (α := α) (m := m) lift γ (fun _ _ _ => True)
       wellFounded_of_finite
-      it init ((⟨·, .intro⟩) <$> f · ·)
+      it init (fun out h acc => (⟨·, .intro⟩) <$> f out h acc)
 
 instance {m : Type w → Type w'} {n : Type w → Type w''}
     {α : Type w} {β : Type w} [Iterator α m β] [Finite α m] [IteratorLoop α m n]
     [MonadLiftT m n] :
-    ForIn n (IterM (α := α) m β) β := IteratorLoop.finiteForIn (fun _ => monadLift)
+    ForIn' n (IterM (α := α) m β) β ⟨fun it out => it.IsPlausibleIndirectOutput out⟩ :=
+  IteratorLoop.finiteForIn' (fun _ => monadLift)
 
 instance {m : Type w → Type w'} {n : Type w → Type w''}
     {α : Type w} {β : Type w} [Iterator α m β] [IteratorLoopPartial α m n] [MonadLiftT m n] :
-    ForIn n (IterM.Partial (α := α) m β) β where
-  forIn it init f :=
+    ForIn' n (IterM.Partial (α := α) m β) β ⟨fun it out => it.it.IsPlausibleIndirectOutput out⟩ where
+  forIn' it init f :=
     IteratorLoopPartial.forInPartial (α := α) (m := m) (fun _ => monadLift) it.it init f
 
 instance {m : Type w → Type w'} {n : Type w → Type w''}
@@ -327,4 +346,42 @@ def IterM.Partial.drain {α : Type w} {m : Type w → Type w'} [Monad m] {β : T
     m PUnit :=
   it.fold (γ := PUnit) (fun _ _ => .unit) .unit
 
+section Size
+
+@[always_inline, inline]
+def IterM.DefaultConsumers.size {α : Type w} {m : Type w → Type w'} [Monad m] {β : Type w}
+    [Iterator α m β] [IteratorLoop α m m] [Finite α m] (it : IterM (α := α) m β) :
+    m (ULift Nat) :=
+  it.fold (init := .up 0) fun acc _ => .up (acc.down + 1)
+
+@[always_inline, inline]
+def IterM.DefaultConsumers.sizePartial {α : Type w} {m : Type w → Type w'} [Monad m] {β : Type w}
+    [Iterator α m β] [IteratorLoopPartial α m m] (it : IterM (α := α) m β) :
+    m (ULift Nat) :=
+  it.allowNontermination.fold (init := .up 0) fun acc _ => .up (acc.down + 1)
+
+@[always_inline, inline]
+instance IteratorSize.defaultImplementation {α β : Type w} {m : Type w → Type w'} [Monad m]
+    [Iterator α m β] [Finite α m] [IteratorLoop α m m] :
+    IteratorSize α m where
+  size := IterM.DefaultConsumers.size
+
+@[always_inline, inline]
+instance IteratorSize.Partial.defaultImplementation {α β : Type w} {m : Type w → Type w'} [Monad m]
+    [Iterator α m β] [Finite α m] [IteratorLoopPartial α m m] :
+    IteratorSizePartial α m where
+  size := IterM.DefaultConsumers.sizePartial
+
+@[always_inline, inline]
+def IterM.size {α : Type} {m : Type → Type w'} {β : Type} [Iterator α m β] [Monad m]
+    (it : IterM (α := α) m β) [IteratorSize α m] : m Nat :=
+  ULift.down <$> IteratorSize.size it
+
+@[always_inline, inline]
+def IterM.Partial.size {α : Type} {m : Type → Type w'} {β : Type} [Iterator α m β] [Monad m]
+    (it : IterM.Partial (α := α) m β) [IteratorSizePartial α m] : m Nat :=
+  ULift.down <$> IteratorSizePartial.size it.it
+
+end Size
+
 end Std.Iterators

src/Std/Data/Iterators/Consumers/Monadic/Partial.lean → src/Init/Data/Iterators/Consumers/Monadic/Partial.lean

@@ -3,8 +3,10 @@ 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 Std.Data.Iterators.Basic
+import Init.Data.Iterators.Basic
 
 namespace Std.Iterators
 

src/Std/Data/Iterators/Consumers/Partial.lean → src/Init/Data/Iterators/Consumers/Partial.lean

@@ -3,8 +3,10 @@ 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 Std.Data.Iterators.Basic
+import Init.Data.Iterators.Basic
 
 namespace Std.Iterators
 

src/Std/Data/Iterators/Internal.lean → src/Init/Data/Iterators/Internal.lean

@@ -3,5 +3,8 @@ 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 Std.Data.Iterators.Internal.Termination
+import Init.Data.Iterators.Internal.LawfulMonadLiftFunction
+import Init.Data.Iterators.Internal.Termination

src/Std/Data/Internal/LawfulMonadLiftFunction.lean → src/Init/Data/Iterators/Internal/LawfulMonadLiftFunction.lean

@@ -3,6 +3,8 @@ 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.Control.Basic
 import Init.Control.Lawful.Basic

src/Std/Data/Iterators/Internal/Termination.lean → src/Init/Data/Iterators/Internal/Termination.lean

@@ -3,8 +3,10 @@ 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 Std.Data.Iterators.Basic
+import Init.Data.Iterators.Basic
 
 /-!
 This is an internal module used by iterator implementations.

src/Init/Data/Iterators/Lemmas.lean

@@ -0,0 +1,4 @@
+module
+
+prelude
+import Init.Data.Iterators.Lemmas.Consumers