refactor: replace some Subarray functions with generic slice functions (#9017)

This PR removes the `Subarray`-specific `toArray`, `foldlM` and `foldl`
methods and instead provides these operations on `Std.Slice`, which are
implemented with the `ToIterator` instance of the slice. Calling
`subarray.toArray` etc. still works, since `Subarray` is an abbreviation
for `Slice _`.

Because the benchmarks are not so clear, to be safe, I will merge this
only after the release. In contrast to the ranges, the iteration over
slices is not quite as efficient as the old `Subarray`-specific
implementation, which would require either more optimizations in the
iterator library (special `IteratorLoop` and `IteratorCollect`
implementations) or better unboxing support by the compiler.

src/Init/Data/Array/Subarray.lean

@@ -6,7 +6,8 @@ Authors: Leonardo de Moura
 module
 
 prelude
-public import Init.Data.Array.Basic
+public import Init.GetElem
+import Init.Data.Array.GetLit
 public import Init.Data.Slice.Basic
 
 public section
@@ -159,64 +160,9 @@ instance : EmptyCollection (Subarray α) :=
 instance : Inhabited (Subarray α) :=
   ⟨{}⟩
 
-/--
-The run-time implementation of `ForIn.forIn` for `Subarray`, which allows it to be used with `for`
-loops in `do`-notation.
-
-This definition replaces `Subarray.forIn`.
+/-!
+`ForIn` and `foldlM` are implemented in `Init.Data.Slice.Array.Iterator` using the slice iterator.
 -/
-@[inline] unsafe def forInUnsafe {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (s : Subarray α) (b : β) (f : α → β → m (ForInStep β)) : m β :=
-  let sz := USize.ofNat s.stop
-  let rec @[specialize] loop (i : USize) (b : β) : m β := do
-    if i < sz then
-      let a := s.array.uget i lcProof
-      match (← f a b) with
-      | ForInStep.done  b => pure b
-      | ForInStep.yield b => loop (i+1) b
-    else
-      pure b
-  loop (USize.ofNat s.start) b
-
-/--
-The implementation of `ForIn.forIn` for `Subarray`, which allows it to be used with `for` loops in
-`do`-notation.
--/
--- TODO: provide reference implementation
-@[implemented_by Subarray.forInUnsafe]
-protected opaque forIn {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (s : Subarray α) (b : β) (f : α → β → m (ForInStep β)) : m β :=
-  pure b
-
-instance : ForIn m (Subarray α) α where
-  forIn := Subarray.forIn
-
-/--
-Folds a monadic operation from left to right over the elements in a subarray.
-
-An accumulator of type `β` is constructed by starting with `init` and monadically combining each
-element of the subarray with the current accumulator value in turn. The monad in question may permit
-early termination or repetition.
-
-Examples:
-```lean example
-#eval #["red", "green", "blue"].toSubarray.foldlM (init := "") fun acc x => do
-  let l ← Option.guard (· ≠ 0) x.length
-  return s!"{acc}({l}){x} "
-```
-```output
-some "(3)red (5)green (4)blue "
-```
-```lean example
-#eval #["red", "green", "blue"].toSubarray.foldlM (init := 0) fun acc x => do
-  let l ← Option.guard (· ≠ 5) x.length
-  return s!"{acc}({l}){x} "
-```
-```output
-none
-```
--/
-@[inline]
-def foldlM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : β → α → m β) (init : β) (as : Subarray α) : m β :=
-  as.array.foldlM f (init := init) (start := as.start) (stop := as.stop)
 
 /--
 Folds a monadic operation from right to left over the elements in a subarray.
@@ -313,20 +259,6 @@ The elements are processed starting at the highest index and moving down.
 def forRevM {α : Type u} {m : Type v → Type w} [Monad m] (f : α → m PUnit) (as : Subarray α) : m PUnit :=
   as.array.forRevM f (start := as.stop) (stop := as.start)
 
-/--
-Folds an operation from left to right over the elements in a subarray.
-
-An accumulator of type `β` is constructed by starting with `init` and combining each
-element of the subarray with the current accumulator value in turn.
-
-Examples:
- * `#["red", "green", "blue"].toSubarray.foldl (· + ·.length) 0 = 12`
- * `#["red", "green", "blue"].toSubarray.popFront.foldl (· + ·.length) 0 = 9`
--/
-@[inline]
-def foldl {α : Type u} {β : Type v} (f : β → α → β) (init : β) (as : Subarray α) : β :=
-  Id.run <| as.foldlM (pure <| f · ·) (init := init)
-
 /--
 Folds an operation from right to left over the elements in a subarray.
 
@@ -464,18 +396,6 @@ def toSubarray (as : Array α) (start : Nat := 0) (stop : Nat := as.size) : Suba
          start_le_stop := Nat.le_refl _,
          stop_le_array_size := Nat.le_refl _ }⟩
 
-/--
-Allocates a new array that contains the contents of the subarray.
--/
-@[coe]
-def ofSubarray (s : Subarray α) : Array α := Id.run do
-  let mut as := mkEmpty (s.stop - s.start)
-  for a in s do
-    as := as.push a
-  return as
-
-instance : Coe (Subarray α) (Array α) := ⟨ofSubarray⟩
-
 /-- A subarray with the provided bounds.-/
 syntax:max term noWs "[" withoutPosition(term ":" term) "]" : term
 /-- A subarray with the provided lower bound that extends to the rest of the array. -/
@@ -489,22 +409,3 @@ macro_rules
   | `($a[$start : ])      => `(let a := $a; Array.toSubarray a $start a.size)
 
 end Array
-
-@[inherit_doc Array.ofSubarray]
-def Subarray.toArray (s : Subarray α) : Array α :=
-  Array.ofSubarray s
-
-instance : Append (Subarray α) where
-  append x y :=
-   let a := x.toArray ++ y.toArray
-   a.toSubarray 0 a.size
-
-/-- `Subarray` representation. -/
-protected def Subarray.repr [Repr α] (s : Subarray α) : Std.Format :=
-  repr s.toArray ++ ".toSubarray"
-
-instance [Repr α] : Repr (Subarray α) where
-  reprPrec s  _ := Subarray.repr s
-
-instance [ToString α] : ToString (Subarray α) where
-  toString s := toString s.toArray

src/Init/Data/Iterators/Consumers.lean

@@ -12,4 +12,6 @@ public import Init.Data.Iterators.Consumers.Collect
 public import Init.Data.Iterators.Consumers.Loop
 public import Init.Data.Iterators.Consumers.Partial
 
+public import Init.Data.Iterators.Consumers.Stream
+
 public section

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

@@ -6,7 +6,6 @@ Authors: Paul Reichert
 module
 
 prelude
-public import Init.Data.Stream
 public import Init.Data.Iterators.Consumers.Partial
 public import Init.Data.Iterators.Consumers.Loop
 public import Init.Data.Iterators.Consumers.Monadic.Access
@@ -64,15 +63,4 @@ def Iter.atIdx? {α β} [Iterator α Id β] [Productive α Id] [IteratorAccess
   | .skip _ => none
   | .done => none
 
-instance {α β} [Iterator α Id β] [Productive α Id] [IteratorAccess α Id] :
-    Stream (Iter (α := α) β) β where
-  next? it := match (it.toIterM.nextAtIdx? 0).run with
-    | .yield it' out _ => some (out, it'.toIter)
-    | .skip _ h => False.elim ?noskip
-    | .done _ => none
-  where finally
-    case noskip =>
-      revert h
-      exact IterM.not_isPlausibleNthOutputStep_yield
-
 end Std.Iterators

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

@@ -0,0 +1,27 @@
+/-
+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
+public import Init.Data.Stream
+public import Init.Data.Iterators.Consumers.Access
+
+public section
+
+namespace Std.Iterators
+
+instance {α β} [Iterator α Id β] [Productive α Id] [IteratorAccess α Id] :
+    Stream (Iter (α := α) β) β where
+  next? it := match (it.toIterM.nextAtIdx? 0).run with
+    | .yield it' out _ => some (out, it'.toIter)
+    | .skip _ h => False.elim ?noskip
+    | .done _ => none
+  where finally
+    case noskip =>
+      revert h
+      exact IterM.not_isPlausibleNthOutputStep_yield
+
+end Std.Iterators

src/Init/Data/Iterators/ToIterator.lean

@@ -6,7 +6,8 @@ Authors: Paul Reichert
 module
 
 prelude
-public import Init.Data.Iterators.Consumers
+public import Init.Data.Iterators.Consumers.Collect
+public import Init.Data.Iterators.Consumers.Loop
 
 public section
 

src/Init/Data/Range/Polymorphic.lean

@@ -8,6 +8,7 @@ module
 prelude
 public import Init.Data.Range.Polymorphic.Basic
 public import Init.Data.Range.Polymorphic.Iterators
+public import Init.Data.Range.Polymorphic.Stream
 public import Init.Data.Range.Polymorphic.Lemmas
 public import Init.Data.Range.Polymorphic.Nat
 public import Init.Data.Range.Polymorphic.NatLemmas

src/Init/Data/Range/Polymorphic/Iterators.lean

@@ -9,7 +9,6 @@ prelude
 public import Init.Data.Range.Polymorphic.RangeIterator
 public import Init.Data.Range.Polymorphic.Basic
 public import Init.Data.Iterators.Combinators.Attach
-public import Init.Data.Stream
 
 public section
 
@@ -26,10 +25,6 @@ def Internal.iter {sl su α} [UpwardEnumerable α] [BoundedUpwardEnumerable sl
     (r : PRange ⟨sl, su⟩ α) : Iter (α := RangeIterator su α) α :=
   ⟨⟨BoundedUpwardEnumerable.init? r.lower, r.upper⟩⟩
 
-instance {sl su α} [UpwardEnumerable α] [BoundedUpwardEnumerable sl α] :
-    ToStream (PRange ⟨sl, su⟩ α) (Iter (α := RangeIterator su α) α) where
-  toStream r := Internal.iter r
-
 /--
 Returns the elements of the given range as a list in ascending order, given that ranges of the given
 type and shape support this function and the range is finite.

src/Init/Data/Range/Polymorphic/Stream.lean

@@ -0,0 +1,22 @@
+/-
+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
+public import Init.Data.Range.Polymorphic.Iterators
+public import Init.Data.Stream
+
+public section
+
+open Std.Iterators
+
+namespace Std.PRange
+
+instance {sl su α} [UpwardEnumerable α] [BoundedUpwardEnumerable sl α] :
+    ToStream (PRange ⟨sl, su⟩ α) (Iter (α := RangeIterator su α) α) where
+  toStream r := Internal.iter r
+
+end Std.PRange

src/Init/Data/Slice/Array/Iterator.lean

@@ -8,8 +8,11 @@ module
 prelude
 public import Init.Core
 public import Init.Data.Slice.Array.Basic
-public import Init.Data.Iterators.Combinators.Attach
-public import Init.Data.Iterators.Combinators.FilterMap
+import Init.Data.Iterators.Combinators.Attach
+import Init.Data.Iterators.Combinators.FilterMap
+import Init.Data.Iterators.Combinators.ULift
+public import Init.Data.Iterators.Consumers.Collect
+public import Init.Data.Iterators.Consumers.Loop
 public import all Init.Data.Range.Polymorphic.Basic
 public import Init.Data.Range.Polymorphic.Nat
 public import Init.Data.Range.Polymorphic.Iterators
@@ -26,6 +29,7 @@ open Std Slice PRange Iterators
 
 variable {shape : RangeShape} {α : Type u}
 
+@[no_expose]
 instance {s : Subarray α} : ToIterator s Id α :=
   .of _
     (PRange.Internal.iter (s.internalRepresentation.start...<s.internalRepresentation.stop)
@@ -39,3 +43,89 @@ where finally
     intro out _ h
     have := s.internalRepresentation.stop_le_array_size
     omega
+
+universe v w
+
+@[no_expose] instance {s : Subarray α} : Iterator (ToIterator.State s Id) Id α := inferInstance
+@[no_expose] instance {s : Subarray α} : Finite (ToIterator.State s Id) Id := inferInstance
+@[no_expose] instance {s : Subarray α} : IteratorCollect (ToIterator.State s Id) Id Id := inferInstance
+@[no_expose] instance {s : Subarray α} : IteratorCollectPartial (ToIterator.State s Id) Id Id := inferInstance
+@[no_expose] instance {s : Subarray α} {m : Type v → Type w} [Monad m] :
+    IteratorLoop (ToIterator.State s Id) Id m := inferInstance
+@[no_expose] instance {s : Subarray α} {m : Type v → Type w} [Monad m] :
+    IteratorLoopPartial (ToIterator.State s Id) Id m := inferInstance
+@[no_expose] instance {s : Subarray α} :
+    IteratorSize (ToIterator.State s Id) Id := inferInstance
+@[no_expose] instance {s : Subarray α} :
+    IteratorSizePartial (ToIterator.State s Id) Id := inferInstance
+
+@[no_expose]
+instance {α : Type u} {m : Type v → Type w} :
+    ForIn m (Subarray α) α where
+  forIn xs init f := forIn (Std.Slice.Internal.iter xs) init f
+
+/-!
+Without defining the following function `Subarray.foldlM`, it is still possible to call
+`subarray.foldlM`, which would be elaborated to `Slice.foldlM (s := subarray)`. However, in order to
+maximize backward compatibility and avoid confusion in the manual entry for `Subarray`, we
+explicitly provide the wrapper function `Subarray.foldlM` for `Slice.foldlM`, providing a more
+specific docstring.
+-/
+
+/--
+Folds a monadic operation from left to right over the elements in a subarray.
+An accumulator of type `β` is constructed by starting with `init` and monadically combining each
+element of the subarray with the current accumulator value in turn. The monad in question may permit
+early termination or repetition.
+Examples:
+```lean example
+#eval #["red", "green", "blue"].toSubarray.foldlM (init := "") fun acc x => do
+  let l ← Option.guard (· ≠ 0) x.length
+  return s!"{acc}({l}){x} "
+```
+```output
+some "(3)red (5)green (4)blue "
+```
+```lean example
+#eval #["red", "green", "blue"].toSubarray.foldlM (init := 0) fun acc x => do
+  let l ← Option.guard (· ≠ 5) x.length
+  return s!"{acc}({l}){x} "
+```
+```output
+none
+```
+-/
+@[inline]
+def Subarray.foldlM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : β → α → m β) (init : β) (as : Subarray α) : m β :=
+  Slice.foldlM f (init := init) as
+
+namespace Array
+
+/--
+Allocates a new array that contains the contents of the subarray.
+-/
+@[coe]
+def ofSubarray (s : Subarray α) : Array α := Id.run do
+  let mut as := mkEmpty (s.stop - s.start)
+  for a in s do
+    as := as.push a
+  return as
+
+instance : Coe (Subarray α) (Array α) := ⟨ofSubarray⟩
+
+instance : Append (Subarray α) where
+  append x y :=
+   let a := x.toArray ++ y.toArray
+   a.toSubarray 0 a.size
+
+/-- `Subarray` representation. -/
+protected def Subarray.repr [Repr α] (s : Subarray α) : Std.Format :=
+  repr s.toArray ++ ".toSubarray"
+
+instance [Repr α] : Repr (Subarray α) where
+  reprPrec s  _ := Subarray.repr s
+
+instance [ToString α] : ToString (Subarray α) where
+  toString s := toString s.toArray
+
+end Array

src/Init/Data/Slice/Array/Lemmas.lean

@@ -21,7 +21,7 @@ open Std.Iterators Std.PRange
 
 namespace Std.Slice.Array
 
-theorem internalIter_eq {α : Type u} {s : Subarray α} :
+private theorem internalIter_eq {α : Type u} {s : Subarray α} :
     Internal.iter s = (PRange.Internal.iter (s.start...<s.stop)
       |>.attachWith (· < s.array.size)
         (fun out h => h
@@ -32,7 +32,7 @@ theorem internalIter_eq {α : Type u} {s : Subarray α} :
       |>.map fun | .up i => s.array[i.1]) := by
   simp [Internal.iter, ToIterator.iter_eq, Subarray.start, Subarray.stop, Subarray.array]
 
-theorem toList_internalIter {α : Type u} {s : Subarray α} :
+private theorem toList_internalIter {α : Type u} {s : Subarray α} :
     (Internal.iter s).toList =
       ((s.start...s.stop).toList
         |>.attachWith (· < s.array.size)

src/Init/Data/Slice/Lemmas.lean

@@ -7,6 +7,7 @@ module
 
 prelude
 public import all Init.Data.Slice.Operations
+import Init.Data.Iterators.Lemmas.Consumers
 
 public section
 

src/Init/Data/Slice/Operations.lean

@@ -8,7 +8,7 @@ module
 prelude
 public import Init.Data.Slice.Basic
 public import Init.Data.Slice.Notation
-public import Init.Data.Iterators
+public import Init.Data.Iterators.ToIterator
 
 public section
 
@@ -34,26 +34,72 @@ Returns the number of elements with distinct indices in the given slice.
 Example: `#[1, 1, 1][0...2].size = 2`.
 -/
 @[always_inline, inline]
-def size (s : Slice g) [ToIterator s Id β] [Iterator (ToIterator.State s Id) Id β]
+def size (s : Slice γ) [ToIterator s Id β] [Iterator (ToIterator.State s Id) Id β]
     [IteratorSize (ToIterator.State s Id) Id] :=
   Internal.iter s |>.size
 
 /-- Allocates a new array that contains the elements of the slice. -/
 @[always_inline, inline]
-def toArray (s : Slice g) [ToIterator s Id β] [Iterator (ToIterator.State s Id) Id β]
+def toArray (s : Slice γ) [ToIterator s Id β] [Iterator (ToIterator.State s Id) Id β]
     [IteratorCollect (ToIterator.State s Id) Id Id] [Finite (ToIterator.State s Id) Id] : Array β :=
   Internal.iter s |>.toArray
 
 /-- Allocates a new list that contains the elements of the slice. -/
 @[always_inline, inline]
-def toList (s : Slice g) [ToIterator s Id β] [Iterator (ToIterator.State s Id) Id β]
+def toList (s : Slice γ) [ToIterator s Id β] [Iterator (ToIterator.State s Id) Id β]
     [IteratorCollect (ToIterator.State s Id) Id Id] [Finite (ToIterator.State s Id) Id] : List β :=
   Internal.iter s |>.toList
 
 /-- Allocates a new list that contains the elements of the slice in reverse order. -/
 @[always_inline, inline]
-def toListRev (s : Slice g) [ToIterator s Id β] [Iterator (ToIterator.State s Id) Id β]
+def toListRev (s : Slice γ) [ToIterator s Id β] [Iterator (ToIterator.State s Id) Id β]
     [Finite (ToIterator.State s Id) Id] : List β :=
   Internal.iter s |>.toListRev
 
+/--
+Folds a monadic operation from left to right over the elements in a slice.
+An accumulator of type `β` is constructed by starting with `init` and monadically combining each
+element of the slice with the current accumulator value in turn. The monad in question may permit
+early termination or repetition.
+
+Examples for the special case of subarrays:
+```lean example
+#eval #["red", "green", "blue"].toSubarray.foldlM (init := "") fun acc x => do
+  let l ← Option.guard (· ≠ 0) x.length
+  return s!"{acc}({l}){x} "
+```
+```output
+some "(3)red (5)green (4)blue "
+```
+```lean example
+#eval #["red", "green", "blue"].toSubarray.foldlM (init := 0) fun acc x => do
+  let l ← Option.guard (· ≠ 5) x.length
+  return s!"{acc}({l}){x} "
+```
+```output
+none
+```
+-/
+@[always_inline, inline]
+def foldlM {γ : Type u} {β : Type v}
+    {δ : Type w} {m : Type w → Type w'} [Monad m] (f : δ → β → m δ) (init : δ)
+    (s : Slice γ) [ToIterator s Id β] [Iterator (ToIterator.State s Id) Id β]
+    [IteratorLoop (ToIterator.State s Id) Id m] [Finite (ToIterator.State s Id) Id] : m δ :=
+  Internal.iter s |>.foldM f init
+
+/--
+Folds an operation from left to right over the elements in a slice.
+An accumulator of type `β` is constructed by starting with `init` and combining each
+element of the slice with the current accumulator value in turn.
+Examples for the special case of subarrays:
+ * `#["red", "green", "blue"].toSubarray.foldl (· + ·.length) 0 = 12`
+ * `#["red", "green", "blue"].toSubarray.popFront.foldl (· + ·.length) 0 = 9`
+-/
+@[always_inline, inline]
+def foldl {γ : Type u} {β : Type v}
+    {δ : Type w} (f : δ → β → δ) (init : δ)
+    (s : Slice γ) [ToIterator s Id β] [Iterator (ToIterator.State s Id) Id β]
+    [IteratorLoop (ToIterator.State s Id) Id Id] [Finite (ToIterator.State s Id) Id] : δ :=
+  Internal.iter s |>.fold f init
+
 end Std.Slice

src/Lean/Message.lean

@@ -6,6 +6,7 @@ Author: Sebastian Ullrich, Leonardo de Moura
 Message type used by the Lean frontend
 -/
 prelude
+import Init.Data.Slice.Array
 import Lean.Data.Position
 import Lean.Data.OpenDecl
 import Lean.MetavarContext
@@ -713,7 +714,7 @@ instance : ToMessageData MVarId        := ⟨MessageData.ofGoal⟩
 instance : ToMessageData MessageData   := ⟨id⟩
 instance [ToMessageData α] : ToMessageData (List α)  := ⟨fun as => MessageData.ofList <| as.map toMessageData⟩
 instance [ToMessageData α] : ToMessageData (Array α) := ⟨fun as => toMessageData as.toList⟩
-instance [ToMessageData α] : ToMessageData (Subarray α) := ⟨fun as => toMessageData as.toArray.toList⟩
+instance [ToMessageData α] : ToMessageData (Subarray α) := ⟨fun as => toMessageData as.toList⟩
 instance [ToMessageData α] : ToMessageData (Option α) := ⟨fun | none => "none" | some e => "some (" ++ toMessageData e ++ ")"⟩
 instance [ToMessageData α] [ToMessageData β] : ToMessageData (α × β) :=
   ⟨fun (a, b) => .paren <| toMessageData a ++ "," ++ Format.line ++ toMessageData b⟩

src/Lean/Util/Diff.lean

@@ -5,6 +5,7 @@ Authors: David Thrane Christiansen
 -/
 prelude
 import Init.Data.Array.Subarray.Split
+import Init.Data.Slice.Array.Iterator
 import Init.Data.Range
 import Std.Data.HashMap.Basic
 import Init.Omega