From cd445dce7630b936e727450f6bf780ff2cddc7c8 Mon Sep 17 00:00:00 2001 From: Paul Reichert <6992158+datokrat@users.noreply.github.com> Date: Thu, 3 Jul 2025 21:33:19 +0200 Subject: [PATCH] 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 | 107 +----------------- src/Init/Data/Iterators/Consumers.lean | 2 + src/Init/Data/Iterators/Consumers/Access.lean | 12 -- src/Init/Data/Iterators/Consumers/Stream.lean | 27 +++++ src/Init/Data/Iterators/ToIterator.lean | 3 +- src/Init/Data/Range/Polymorphic.lean | 1 + .../Data/Range/Polymorphic/Iterators.lean | 5 - src/Init/Data/Range/Polymorphic/Stream.lean | 22 ++++ src/Init/Data/Slice/Array/Iterator.lean | 94 ++++++++++++++- src/Init/Data/Slice/Array/Lemmas.lean | 4 +- src/Init/Data/Slice/Lemmas.lean | 1 + src/Init/Data/Slice/Operations.lean | 56 ++++++++- src/Lean/Message.lean | 3 +- src/Lean/Util/Diff.lean | 1 + 14 files changed, 207 insertions(+), 131 deletions(-) create mode 100644 src/Init/Data/Iterators/Consumers/Stream.lean create mode 100644 src/Init/Data/Range/Polymorphic/Stream.lean diff --git a/src/Init/Data/Array/Subarray.lean b/src/Init/Data/Array/Subarray.lean index c9162ce856..21fd17c343 100644 --- a/src/Init/Data/Array/Subarray.lean +++ b/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 diff --git a/src/Init/Data/Iterators/Consumers.lean b/src/Init/Data/Iterators/Consumers.lean index e4c4aedc49..e04ead2e4a 100644 --- a/src/Init/Data/Iterators/Consumers.lean +++ b/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 diff --git a/src/Init/Data/Iterators/Consumers/Access.lean b/src/Init/Data/Iterators/Consumers/Access.lean index d845a787c6..37ebe1befe 100644 --- a/src/Init/Data/Iterators/Consumers/Access.lean +++ b/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 diff --git a/src/Init/Data/Iterators/Consumers/Stream.lean b/src/Init/Data/Iterators/Consumers/Stream.lean new file mode 100644 index 0000000000..c6453b6058 --- /dev/null +++ b/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 diff --git a/src/Init/Data/Iterators/ToIterator.lean b/src/Init/Data/Iterators/ToIterator.lean index eaa1d04455..d1029521f3 100644 --- a/src/Init/Data/Iterators/ToIterator.lean +++ b/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 diff --git a/src/Init/Data/Range/Polymorphic.lean b/src/Init/Data/Range/Polymorphic.lean index 9dc63d5eb4..29b7a9c65f 100644 --- a/src/Init/Data/Range/Polymorphic.lean +++ b/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 diff --git a/src/Init/Data/Range/Polymorphic/Iterators.lean b/src/Init/Data/Range/Polymorphic/Iterators.lean index 3247f2ff28..ad4ba7def1 100644 --- a/src/Init/Data/Range/Polymorphic/Iterators.lean +++ b/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. diff --git a/src/Init/Data/Range/Polymorphic/Stream.lean b/src/Init/Data/Range/Polymorphic/Stream.lean new file mode 100644 index 0000000000..2575792090 --- /dev/null +++ b/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 diff --git a/src/Init/Data/Slice/Array/Iterator.lean b/src/Init/Data/Slice/Array/Iterator.lean index f7f808187f..2de48257bf 100644 --- a/src/Init/Data/Slice/Array/Iterator.lean +++ b/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... 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 diff --git a/src/Init/Data/Slice/Array/Lemmas.lean b/src/Init/Data/Slice/Array/Lemmas.lean index f1991e5328..f01942fcc9 100644 --- a/src/Init/Data/Slice/Array/Lemmas.lean +++ b/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....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) diff --git a/src/Init/Data/Slice/Lemmas.lean b/src/Init/Data/Slice/Lemmas.lean index ba07e47764..00cf83d696 100644 --- a/src/Init/Data/Slice/Lemmas.lean +++ b/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 diff --git a/src/Init/Data/Slice/Operations.lean b/src/Init/Data/Slice/Operations.lean index f6e90ffded..8ea73d9251 100644 --- a/src/Init/Data/Slice/Operations.lean +++ b/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 diff --git a/src/Lean/Message.lean b/src/Lean/Message.lean index efe4e83db5..d3ce0d7eee 100644 --- a/src/Lean/Message.lean +++ b/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⟩ diff --git a/src/Lean/Util/Diff.lean b/src/Lean/Util/Diff.lean index 41cd4b6122..63d9839f67 100644 --- a/src/Lean/Util/Diff.lean +++ 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