From 8eb02930985d2d13c94790359a8e99fa9068f1eb Mon Sep 17 00:00:00 2001 From: Paul Reichert <6992158+datokrat@users.noreply.github.com> Date: Mon, 17 Nov 2025 16:58:29 +0100 Subject: [PATCH] feat: add MPL specs for slice `for ... in` (#11141) This PR provides a polymorphic `ForIn` instance for slices and an MPL `spec` lemma for the iteration over slices using `for ... in`. It also provides a version specialized to `Subarray`. --- src/Init/Data/Array/Lemmas.lean | 4 +-- src/Init/Data/Iterators/ToIterator.lean | 12 ++++++++ src/Init/Data/Slice/Array/Iterator.lean | 26 +++++++++------- src/Init/Data/Slice/Array/Lemmas.lean | 29 ++++++++++++++--- src/Init/Data/Slice/Lemmas.lean | 41 +++++++++++++++++++++++++ src/Init/Data/Slice/Operations.lean | 8 +++++ src/Std/Do/Triple/SpecLemmas.lean | 29 +++++++++++++++++ tests/lean/run/doLogicTests.lean | 21 ++++++++++--- 8 files changed, 148 insertions(+), 22 deletions(-) diff --git a/src/Init/Data/Array/Lemmas.lean b/src/Init/Data/Array/Lemmas.lean index 71f56156e6..e391e8e389 100644 --- a/src/Init/Data/Array/Lemmas.lean +++ b/src/Init/Data/Array/Lemmas.lean @@ -2231,8 +2231,8 @@ theorem push_eq_flatten_iff {xss : Array (Array α)} {ys : Array α} {y : α} : -- zs = cs ++ ds.flatten := by sorry -/-- Two arrays of subarrays are equal iff their flattens coincide, as well as the sizes of the -subarrays. -/ +/-- Two arrays of arrays are equal iff their flattens coincide, as well as the sizes of the +arrays. -/ theorem eq_iff_flatten_eq {xss₁ xss₂ : Array (Array α)} : xss₁ = xss₂ ↔ xss₁.flatten = xss₂.flatten ∧ map size xss₁ = map size xss₂ := by cases xss₁ using array₂_induction with diff --git a/src/Init/Data/Iterators/ToIterator.lean b/src/Init/Data/Iterators/ToIterator.lean index c4338df0ac..3514474b4e 100644 --- a/src/Init/Data/Iterators/ToIterator.lean +++ b/src/Init/Data/Iterators/ToIterator.lean @@ -89,6 +89,12 @@ instance {x : γ} {State : Type w} {iter} IteratorCollect (α := i.State) m n := inferInstanceAs <| IteratorCollect (α := State) m n +instance {x : γ} {State : Type w} {iter} [Monad m] [Monad n] + [Iterator (α := State) m β] [IteratorCollect State m n] [LawfulIteratorCollect State m n] : + letI i : ToIterator x m β := .ofM State iter + LawfulIteratorCollect (α := i.State) m n := + inferInstanceAs <| LawfulIteratorCollect (α := State) m n + instance {x : γ} {State : Type w} {iter} [Iterator (α := State) m β] [IteratorCollectPartial State m n] : letI i : ToIterator x m β := .ofM State iter @@ -101,6 +107,12 @@ instance {x : γ} {State : Type w} {iter} IteratorLoop (α := i.State) m n := inferInstanceAs <| IteratorLoop (α := State) m n +instance {x : γ} {State : Type w} {iter} [Monad m] [Monad n] + [Iterator (α := State) m β] [IteratorLoop State m n] [LawfulIteratorLoop State m n]: + letI i : ToIterator x m β := .ofM State iter + LawfulIteratorLoop (α := i.State) m n := + inferInstanceAs <| LawfulIteratorLoop (α := State) m n + instance {x : γ} {State : Type w} {iter} [Iterator (α := State) m β] [IteratorLoopPartial State m n] : letI i : ToIterator x m β := .ofM State iter diff --git a/src/Init/Data/Slice/Array/Iterator.lean b/src/Init/Data/Slice/Array/Iterator.lean index dcd372346f..b9cef642c6 100644 --- a/src/Init/Data/Slice/Array/Iterator.lean +++ b/src/Init/Data/Slice/Array/Iterator.lean @@ -7,12 +7,14 @@ module prelude public import Init.Data.Slice.Array.Basic +public import Init.Data.Slice.Operations import Init.Data.Iterators.Combinators.Attach public import Init.Data.Iterators.Combinators.ULift import all Init.Data.Range.Polymorphic.Basic public import Init.Data.Range.Polymorphic.Iterators public import Init.Data.Slice.Operations import Init.Omega +import Init.Data.Iterators.Lemmas.Combinators.Monadic.FilterMap public section @@ -24,7 +26,7 @@ open Std Slice PRange Iterators variable {shape : RangeShape} {α : Type u} -instance {s : Subarray α} : ToIterator s Id α := +instance {s : Slice (Internal.SubarrayData α)} : ToIterator s Id α := .of _ (Rco.Internal.iter (s.internalRepresentation.start....attachWith (· < s.internalRepresentation.array.size) ?h @@ -39,22 +41,24 @@ where finally 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] : +@[no_expose] instance {s : Slice (Internal.SubarrayData α)} : Iterator (ToIterator.State s Id) Id α := inferInstance +@[no_expose] instance {s : Slice (Internal.SubarrayData α)} : Finite (ToIterator.State s Id) Id := inferInstance +@[no_expose] instance {s : Slice (Internal.SubarrayData α)} : IteratorCollect (ToIterator.State s Id) Id Id := inferInstance +@[no_expose] instance {s : Slice (Internal.SubarrayData α)} : LawfulIteratorCollect (ToIterator.State s Id) Id Id := inferInstance +@[no_expose] instance {s : Slice (Internal.SubarrayData α)} : IteratorCollectPartial (ToIterator.State s Id) Id Id := inferInstance +@[no_expose] instance {s : Slice (Internal.SubarrayData α)} {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] : +@[no_expose] instance {s : Slice (Internal.SubarrayData α)} {m : Type v → Type w} [Monad m] : + LawfulIteratorLoop (ToIterator.State s Id) Id m := inferInstance +@[no_expose] instance {s : Slice (Internal.SubarrayData α)} {m : Type v → Type w} [Monad m] : IteratorLoopPartial (ToIterator.State s Id) Id m := inferInstance instance : SliceSize (Internal.SubarrayData α) where size s := s.internalRepresentation.stop - s.internalRepresentation.start -@[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 +instance {α : Type u} {m : Type v → Type w} [Monad m] : + ForIn m (Subarray α) α := + inferInstance /-! Without defining the following function `Subarray.foldlM`, it is still possible to call diff --git a/src/Init/Data/Slice/Array/Lemmas.lean b/src/Init/Data/Slice/Array/Lemmas.lean index 88a3b00ee6..63e1348877 100644 --- a/src/Init/Data/Slice/Array/Lemmas.lean +++ b/src/Init/Data/Slice/Array/Lemmas.lean @@ -8,6 +8,7 @@ module prelude import all Init.Data.Array.Subarray import all Init.Data.Slice.Array.Basic +import Init.Data.Slice.Lemmas public import Init.Data.Slice.Array.Iterator import all Init.Data.Slice.Array.Iterator import all Init.Data.Slice.Operations @@ -16,11 +17,11 @@ import all Init.Data.Range.Polymorphic.Lemmas public import Init.Data.Slice.Lemmas public import Init.Data.Iterators.Lemmas -open Std.Iterators Std.PRange +open Std Std.Iterators Std.PRange Std.Slice -namespace Std.Slice.Array +namespace Subarray -theorem internalIter_rco_eq {α : Type u} {s : Subarray α} : +theorem internalIter_eq {α : Type u} {s : Subarray α} : Internal.iter s = (Rco.Internal.iter (s.start....attachWith (· < s.array.size) (fun out h => h @@ -40,7 +41,7 @@ theorem toList_internalIter {α : Type u} {s : Subarray α} : |> Rco.lt_upper_of_mem |> (Nat.lt_of_lt_of_le · s.stop_le_array_size)) |>.map fun i => s.array[i.1]) := by - rw [internalIter_rco_eq, Iter.toList_map, Iter.toList_uLift, Iter.toList_attachWith] + rw [internalIter_eq, Iter.toList_map, Iter.toList_uLift, Iter.toList_attachWith] simp [Rco.toList] public instance : LawfulSliceSize (Internal.SubarrayData α) where @@ -50,4 +51,22 @@ public instance : LawfulSliceSize (Internal.SubarrayData α) where Rco.size_toArray, Rco.size, Rxo.HasSize.size, Rxc.HasSize.size] omega -end Std.Slice.Array +public theorem toArray_eq_sliceToArray {α : Type u} {s : Subarray α} : + s.toArray = Slice.toArray s := by + simp [Subarray.toArray, Array.ofSubarray] + +@[simp] +public theorem forIn_toList {α : Type u} {s : Subarray α} + {m : Type v → Type w} [Monad m] [LawfulMonad m] {γ : Type v} {init : γ} + {f : α → γ → m (ForInStep γ)} : + ForIn.forIn s.toList init f = ForIn.forIn s init f := + Slice.forIn_toList + +@[simp] +public theorem forIn_toArray {α : Type u} {s : Subarray α} + {m : Type v → Type w} [Monad m] [LawfulMonad m] {γ : Type v} {init : γ} + {f : α → γ → m (ForInStep γ)} : + ForIn.forIn s.toArray init f = ForIn.forIn s init f := + Slice.forIn_toArray + +end Subarray diff --git a/src/Init/Data/Slice/Lemmas.lean b/src/Init/Data/Slice/Lemmas.lean index 82a395b3c6..681e186ae0 100644 --- a/src/Init/Data/Slice/Lemmas.lean +++ b/src/Init/Data/Slice/Lemmas.lean @@ -8,7 +8,9 @@ module prelude public import Init.Data.Slice.Operations import all Init.Data.Slice.Operations +import Init.Data.Iterators.Consumers import Init.Data.Iterators.Lemmas.Consumers +public import Init.Data.List.Control public section @@ -23,6 +25,45 @@ theorem Internal.iter_eq_toIteratorIter {γ : Type u} {s : Slice γ} Internal.iter s = ToIterator.iter s := (rfl) +theorem forIn_internalIter {γ : Type u} {β : Type v} + {m : Type w → Type x} [Monad m] {δ : Type w} + [∀ s : Slice γ, ToIterator s Id β] + [∀ s : Slice γ, Iterator (ToIterator.State s Id) Id β] + [∀ s : Slice γ, IteratorLoop (ToIterator.State s Id) Id m] + [∀ s : Slice γ, LawfulIteratorLoop (ToIterator.State s Id) Id m] + [∀ s : Slice γ, Finite (ToIterator.State s Id) Id] {s : Slice γ} + {init : δ} {f : β → δ → m (ForInStep δ)} : + ForIn.forIn (Internal.iter s) init f = ForIn.forIn s init f := + (rfl) + +@[simp] +public theorem forIn_toList {γ : Type u} {β : Type v} + {m : Type w → Type x} [Monad m] [LawfulMonad m] {δ : Type w} + [∀ s : Slice γ, ToIterator s Id β] + [∀ s : Slice γ, Iterator (ToIterator.State s Id) Id β] + [∀ s : Slice γ, IteratorLoop (ToIterator.State s Id) Id m] + [∀ s : Slice γ, LawfulIteratorLoop (ToIterator.State s Id) Id m] + [∀ s : Slice γ, IteratorCollect (ToIterator.State s Id) Id Id] + [∀ s : Slice γ, LawfulIteratorCollect (ToIterator.State s Id) Id Id] + [∀ s : Slice γ, Finite (ToIterator.State s Id) Id] {s : Slice γ} + {init : δ} {f : β → δ → m (ForInStep δ)} : + ForIn.forIn s.toList init f = ForIn.forIn s init f := by + rw [← forIn_internalIter, ← Iter.forIn_toList, Slice.toList] + +@[simp] +public theorem forIn_toArray {γ : Type u} {β : Type v} + {m : Type w → Type x} [Monad m] [LawfulMonad m] {δ : Type w} + [∀ s : Slice γ, ToIterator s Id β] + [∀ s : Slice γ, Iterator (ToIterator.State s Id) Id β] + [∀ s : Slice γ, IteratorLoop (ToIterator.State s Id) Id m] + [∀ s : Slice γ, LawfulIteratorLoop (ToIterator.State s Id) Id m] + [∀ s : Slice γ, IteratorCollect (ToIterator.State s Id) Id Id] + [∀ s : Slice γ, LawfulIteratorCollect (ToIterator.State s Id) Id Id] + [∀ s : Slice γ, Finite (ToIterator.State s Id) Id] {s : Slice γ} + {init : δ} {f : β → δ → m (ForInStep δ)} : + ForIn.forIn s.toArray init f = ForIn.forIn s init f := by + rw [← forIn_internalIter, ← Iter.forIn_toArray, Slice.toArray] + theorem Internal.size_eq_count_iter [∀ s : Slice γ, ToIterator s Id β] [∀ s : Slice γ, Iterator (ToIterator.State s Id) Id β] {s : Slice γ} [Finite (ToIterator.State s Id) Id] diff --git a/src/Init/Data/Slice/Operations.lean b/src/Init/Data/Slice/Operations.lean index cb4b6ac1e6..c31f0c7e4c 100644 --- a/src/Init/Data/Slice/Operations.lean +++ b/src/Init/Data/Slice/Operations.lean @@ -75,6 +75,14 @@ def toListRev (s : Slice γ) [ToIterator s Id β] [Iterator (ToIterator.State s [Finite (ToIterator.State s Id) Id] : List β := Internal.iter s |>.toListRev +instance {γ : Type u} {β : Type v} [∀ s : Slice γ, ToIterator s Id β] + [∀ s : Slice γ, Iterator (ToIterator.State s Id) Id β] + [∀ s : Slice γ, IteratorLoop (ToIterator.State s Id) Id m] + [∀ s : Slice γ, Finite (ToIterator.State s Id) Id] : + ForIn m (Slice γ) β where + forIn s init f := + forIn (Internal.iter s) init f + /-- 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 diff --git a/src/Std/Do/Triple/SpecLemmas.lean b/src/Std/Do/Triple/SpecLemmas.lean index 1def16982f..cd8db3b458 100644 --- a/src/Std/Do/Triple/SpecLemmas.lean +++ b/src/Std/Do/Triple/SpecLemmas.lean @@ -9,6 +9,8 @@ prelude public import Std.Do.Triple.Basic public import Init.Data.Range.Polymorphic.Iterators import Init.Data.Range.Polymorphic +public import Init.Data.Slice.Array +public import Init.Data.Iterators.ToIterator -- This public import is a workaround for #10652. -- Without it, adding the `spec` attribute for `instMonadLiftTOfMonadLift` will fail. @@ -1087,6 +1089,33 @@ theorem Spec.forIn_rii {α β : Type u} {m : Type u → Type v} {ps : PostShape} simp only [forIn] apply Spec.forIn'_rii inv step +open Std.Iterators in +@[spec] +theorem Spec.forIn_slice {m : Type w → Type x} {ps : PostShape} + [Monad m] [WPMonad m ps] + {γ : Type u} {β : Type w} + [LawfulMonad m] {δ : Type w} + [∀ s : Slice γ, ToIterator s Id β] + [∀ s : Slice γ, Iterator (ToIterator.State s Id) Id β] + [∀ s : Slice γ, IteratorLoop (ToIterator.State s Id) Id m] + [∀ s : Slice γ, LawfulIteratorLoop (ToIterator.State s Id) Id m] + [∀ s : Slice γ, IteratorCollect (ToIterator.State s Id) Id Id] + [∀ s : Slice γ, LawfulIteratorCollect (ToIterator.State s Id) Id Id] + [∀ s : Slice γ, Finite (ToIterator.State s Id) Id] + {init : δ} {f : β → δ → m (ForInStep δ)} + {xs : Slice γ} + (inv : Invariant xs.toList δ ps) + (step : ∀ pref cur suff (h : xs.toList = pref ++ cur :: suff) b, + Triple + (f cur b) + (inv.1 (⟨pref, cur::suff, h.symm⟩, b)) + (fun r => match r with + | .yield b' => inv.1 (⟨pref ++ [cur], suff, by simp [h]⟩, b') + | .done b' => inv.1 (⟨xs.toList, [], by simp⟩, b'), inv.2)) : + Triple (forIn xs init f) (inv.1 (⟨[], xs.toList, rfl⟩, init)) (fun b => inv.1 (⟨xs.toList, [], by simp⟩, b), inv.2) := by + simp only [← Slice.forIn_toList] + exact Spec.forIn_list inv step + @[spec] theorem Spec.forIn'_array {α β : Type u} {m : Type u → Type v} {ps : PostShape} [Monad m] [WPMonad m ps] diff --git a/tests/lean/run/doLogicTests.lean b/tests/lean/run/doLogicTests.lean index 050e27dab5..a7b3415510 100644 --- a/tests/lean/run/doLogicTests.lean +++ b/tests/lean/run/doLogicTests.lean @@ -4,8 +4,6 @@ Released under Apache 2.0 license as described in the file LICENSE. Authors: Sebastian Graf -/ -import Std.Tactic.Do -import Std.Tactic.Do.Syntax import Std import Lean.Elab.Tactic.Do.VCGen @@ -794,6 +792,23 @@ theorem mem_mergeWithAll [LawfulEqCmp cmp] {m₁ m₂ : ExtTreeMap α β cmp} {f end KimsUnivPolyUseCase +namespace Slices + +def subarraySum (xs : Subarray Nat) : Nat := Id.run do + let mut sum := 0 + for x in xs do + sum := sum + x + return sum + +theorem subarraySum_correct {xs : Subarray Nat} : subarraySum xs = xs.toList.sum := by + generalize h : subarraySum xs = r + apply Id.of_wp_run_eq h + mvcgen + case inv1 => exact ⇓⟨cursor, prefixSum⟩ => ⌜prefixSum = cursor.prefix.sum⌝ + all_goals simp_all + +end Slices + namespace PatricksFastExp def naive_expo (x n : Nat) : Nat := Id.run do @@ -817,8 +832,6 @@ def fast_expo (x n : Nat) : Nat := Id.run do return y -open Std.Do - theorem naive_expo_correct (x n : Nat) : naive_expo x n = x^n := by generalize h : naive_expo x n = r apply Id.of_wp_run_eq h