From ed4252f8c93f591ae753cc963cfd69657ebd3e6a Mon Sep 17 00:00:00 2001 From: Paul Reichert <6992158+datokrat@users.noreply.github.com> Date: Fri, 30 May 2025 20:17:53 +0200 Subject: [PATCH] feat: array iterators, `repeat`/`unfold`, `ForM` for iterators (#8552) This PR provides array iterators (`Array.iter(M)`, `Array.iterFromIdx(M)`), infinite iterators produced by a step function (`Iter.repeat`), and a `ForM` instance for finite iterators that is implemented in terms of `ForIn`. --- src/Std/Data/Iterators/Consumers/Loop.lean | 10 ++ .../Iterators/Consumers/Monadic/Loop.lean | 15 ++- .../Lemmas/Combinators/Monadic/Take.lean | 8 ++ .../Iterators/Lemmas/Combinators/Take.lean | 9 ++ .../Lemmas/Consumers/Monadic/Loop.lean | 25 ++++ src/Std/Data/Iterators/Lemmas/Producers.lean | 2 + .../Iterators/Lemmas/Producers/Array.lean | 87 +++++++++++++ .../Iterators/Lemmas/Producers/Monadic.lean | 1 + .../Lemmas/Producers/Monadic/Array.lean | 120 ++++++++++++++++++ .../Iterators/Lemmas/Producers/Repeat.lean | 55 ++++++++ src/Std/Data/Iterators/Producers.lean | 2 + src/Std/Data/Iterators/Producers/Array.lean | 49 +++++++ src/Std/Data/Iterators/Producers/Monadic.lean | 1 + .../Iterators/Producers/Monadic/Array.lean | 120 ++++++++++++++++++ src/Std/Data/Iterators/Producers/Repeat.lean | 83 ++++++++++++ tests/lean/run/iterators.lean | 50 ++++++++ 16 files changed, 636 insertions(+), 1 deletion(-) create mode 100644 src/Std/Data/Iterators/Lemmas/Producers/Array.lean create mode 100644 src/Std/Data/Iterators/Lemmas/Producers/Monadic/Array.lean create mode 100644 src/Std/Data/Iterators/Lemmas/Producers/Repeat.lean create mode 100644 src/Std/Data/Iterators/Producers/Array.lean create mode 100644 src/Std/Data/Iterators/Producers/Monadic/Array.lean create mode 100644 src/Std/Data/Iterators/Producers/Repeat.lean diff --git a/src/Std/Data/Iterators/Consumers/Loop.lean b/src/Std/Data/Iterators/Consumers/Loop.lean index 4250bbfe26..35bc0e095a 100644 --- a/src/Std/Data/Iterators/Consumers/Loop.lean +++ b/src/Std/Data/Iterators/Consumers/Loop.lean @@ -35,6 +35,16 @@ instance (α : Type w) (β : Type w) (n : Type w → Type w') [Monad n] letI : MonadLift Id n := ⟨pure⟩ ForIn.forIn it.it.toIterM.allowNontermination init f +instance {m : Type w → Type w'} + {α : Type w} {β : Type w} [Iterator α Id β] [Finite α Id] [IteratorLoop α Id m] : + ForM m (Iter (α := α) β) β where + forM it f := forIn it PUnit.unit (fun out _ => do f out; return .yield .unit) + +instance {m : Type w → Type w'} + {α : Type w} {β : Type w} [Iterator α Id β] [Finite α Id] [IteratorLoopPartial α Id m] : + ForM m (Iter.Partial (α := α) β) β where + forM it f := forIn it PUnit.unit (fun out _ => do f out; return .yield .unit) + /-- Folds a monadic function over an iterator from the left, accumulating a value starting with `init`. The accumulated value is combined with the each element of the list in order, using `f`. diff --git a/src/Std/Data/Iterators/Consumers/Monadic/Loop.lean b/src/Std/Data/Iterators/Consumers/Monadic/Loop.lean index 548c88377d..018db4909a 100644 --- a/src/Std/Data/Iterators/Consumers/Monadic/Loop.lean +++ b/src/Std/Data/Iterators/Consumers/Monadic/Loop.lean @@ -216,7 +216,20 @@ instance {m : Type w → Type w'} {n : Type w → Type w''} 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 := IteratorLoopPartial.forInPartial (α := α) (m := m) (fun _ => monadLift) it.it init f + 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''} + {α : Type w} {β : Type w} [Iterator α m β] [Finite α m] [IteratorLoop α m n] + [MonadLiftT m n] : + ForM n (IterM (α := α) m β) β where + forM it f := forIn it PUnit.unit (fun out _ => do f out; return .yield .unit) + +instance {m : Type w → Type w'} {n : Type w → Type w''} + {α : Type w} {β : Type w} [Iterator α m β] [Finite α m] [IteratorLoopPartial α m n] + [MonadLiftT m n] : + ForM n (IterM.Partial (α := α) m β) β where + forM it f := forIn it PUnit.unit (fun out _ => do f out; return .yield .unit) /-- Folds a monadic function over an iterator from the left, accumulating a value starting with `init`. diff --git a/src/Std/Data/Iterators/Lemmas/Combinators/Monadic/Take.lean b/src/Std/Data/Iterators/Lemmas/Combinators/Monadic/Take.lean index 10a36f9805..8d272a4470 100644 --- a/src/Std/Data/Iterators/Lemmas/Combinators/Monadic/Take.lean +++ b/src/Std/Data/Iterators/Lemmas/Combinators/Monadic/Take.lean @@ -27,4 +27,12 @@ theorem IterM.step_take {α m β} [Monad m] [Iterator α m β] {n : Nat} obtain ⟨step, h⟩ := step cases step <;> rfl +theorem IterM.toList_take_zero {α m β} [Monad m] [LawfulMonad m] [Iterator α m β] + [Finite (Take α m β) m] + [IteratorCollect (Take α m β) m m] [LawfulIteratorCollect (Take α m β) m m] + {it : IterM (α := α) m β} : + (it.take 0).toList = pure [] := by + rw [toList_eq_match_step] + simp [step_take] + end Std.Iterators diff --git a/src/Std/Data/Iterators/Lemmas/Combinators/Take.lean b/src/Std/Data/Iterators/Lemmas/Combinators/Take.lean index b492f1b542..aacd6f29a1 100644 --- a/src/Std/Data/Iterators/Lemmas/Combinators/Take.lean +++ b/src/Std/Data/Iterators/Lemmas/Combinators/Take.lean @@ -83,4 +83,13 @@ theorem Iter.toArray_take_of_finite {α β} [Iterator α Id β] {n : Nat} (it.take n).toArray = it.toArray.take n := by rw [← toArray_toList, ← toArray_toList, List.take_toArray, toList_take_of_finite] +@[simp] +theorem Iter.toList_take_zero {α β} [Iterator α Id β] + [Finite (Take α Id β) Id] + [IteratorCollect (Take α Id β) Id Id] [LawfulIteratorCollect (Take α Id β) Id Id] + {it : Iter (α := α) β} : + (it.take 0).toList = [] := by + rw [toList_eq_match_step] + simp [step_take] + end Std.Iterators diff --git a/src/Std/Data/Iterators/Lemmas/Consumers/Monadic/Loop.lean b/src/Std/Data/Iterators/Lemmas/Consumers/Monadic/Loop.lean index cd4efbe15a..6f939b41be 100644 --- a/src/Std/Data/Iterators/Lemmas/Consumers/Monadic/Loop.lean +++ b/src/Std/Data/Iterators/Lemmas/Consumers/Monadic/Loop.lean @@ -65,6 +65,31 @@ theorem IterM.forIn_eq_match_step {α β : Type w} {m : Type w → Type w'} [Ite · simp [IterM.forIn_eq] · simp +theorem IterM.forM_eq_forIn {α β : Type w} {m : Type w → Type w'} [Iterator α m β] + [Finite α m] {n : Type w → Type w''} [Monad n] [LawfulMonad n] + [IteratorLoop α m n] [LawfulIteratorLoop α m n] + [MonadLiftT m n] {it : IterM (α := α) m β} + {f : β → n PUnit} : + ForM.forM it f = ForIn.forIn it PUnit.unit (fun out _ => do f out; return .yield .unit) := + rfl + +theorem IterM.forM_eq_match_step {α β : Type w} {m : Type w → Type w'} [Iterator α m β] + [Finite α m] {n : Type w → Type w''} [Monad n] [LawfulMonad n] + [IteratorLoop α m n] [LawfulIteratorLoop α m n] + [MonadLiftT m n] {it : IterM (α := α) m β} + {f : β → n PUnit} : + ForM.forM it f = (do + match ← it.step with + | .yield it' out _ => + f out + ForM.forM it' f + | .skip it' _ => ForM.forM it' f + | .done _ => return) := by + rw [forM_eq_forIn, forIn_eq_match_step] + apply bind_congr + intro step + cases step using PlausibleIterStep.casesOn <;> simp [forM_eq_forIn] + theorem IterM.foldM_eq_forIn {α β γ : Type w} {m : Type w → Type w'} [Iterator α m β] [Finite α m] {n : Type w → Type w''} [Monad n] [IteratorLoop α m n] [MonadLiftT m n] {f : γ → β → n γ} {init : γ} {it : IterM (α := α) m β} : diff --git a/src/Std/Data/Iterators/Lemmas/Producers.lean b/src/Std/Data/Iterators/Lemmas/Producers.lean index 9913584074..baeab5010c 100644 --- a/src/Std/Data/Iterators/Lemmas/Producers.lean +++ b/src/Std/Data/Iterators/Lemmas/Producers.lean @@ -5,4 +5,6 @@ Authors: Paul Reichert -/ prelude import Std.Data.Iterators.Lemmas.Producers.Monadic +import Std.Data.Iterators.Lemmas.Producers.Array import Std.Data.Iterators.Lemmas.Producers.List +import Std.Data.Iterators.Lemmas.Producers.Repeat diff --git a/src/Std/Data/Iterators/Lemmas/Producers/Array.lean b/src/Std/Data/Iterators/Lemmas/Producers/Array.lean new file mode 100644 index 0000000000..5aac42b3a6 --- /dev/null +++ b/src/Std/Data/Iterators/Lemmas/Producers/Array.lean @@ -0,0 +1,87 @@ +/- +Copyright (c) 2025 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Paul Reichert +-/ +prelude +import Std.Data.Iterators.Lemmas.Consumers.Collect +import Std.Data.Iterators.Lemmas.Producers.Monadic.Array + +/-! +# Lemmas about array iterators + +This module provides lemmas about the interactions of `Array.iter` with `Iter.step` and various +collectors. +-/ + +namespace Std.Iterators + +variable {β : Type w} + +theorem _root_.Array.iter_eq_toIter_iterM {array : Array β} : + array.iter = (array.iterM Id).toIter := + rfl + +theorem _root_.Array.iter_eq_iterFromIdx {array : Array β} : + array.iter = array.iterFromIdx 0 := + rfl + +theorem _root_.Array.iterFromIdx_eq_toIter_iterFromIdxM {array : Array β} {pos : Nat} : + array.iterFromIdx pos = (array.iterFromIdxM Id pos).toIter := + rfl + +theorem _root_.Array.step_iterFromIdx {array : Array β} {pos : Nat} : + (array.iterFromIdx pos).step = if h : pos < array.size then + .yield + (array.iterFromIdx (pos + 1)) + array[pos] + ⟨rfl, rfl, h, rfl⟩ + else + .done (Nat.not_lt.mp h) := by + simp only [Array.iterFromIdx_eq_toIter_iterFromIdxM, Iter.step, Iter.toIterM_toIter, + Array.step_iterFromIdxM, Id.run_pure] + split <;> rfl + +theorem _root_.Array.step_iter {array : Array β} : + array.iter.step = if h : 0 < array.size then + .yield + (array.iterFromIdx 1) + array[0] + ⟨rfl, rfl, h, rfl⟩ + else + .done (Nat.not_lt.mp h) := by + simp only [Array.iter_eq_iterFromIdx, Array.step_iterFromIdx] + +@[simp] +theorem _root_.Array.toList_iterFromIdx {array : Array β} + {pos : Nat} : + (array.iterFromIdx pos).toList = array.toList.drop pos := by + simp [Iter.toList, Array.iterFromIdx_eq_toIter_iterFromIdxM, Iter.toIterM_toIter, + Array.toList_iterFromIdxM] + +@[simp] +theorem _root_.Array.toList_iter {array : Array β} : + array.iter.toList = array.toList := by + simp [Array.iter_eq_iterFromIdx, Array.toList_iterFromIdx] + +@[simp] +theorem _root_.Array.toArray_iterFromIdx {array : Array β} {pos : Nat} : + (array.iterFromIdx pos).toArray = array.extract pos := by + simp [Array.iterFromIdx_eq_toIter_iterFromIdxM, Iter.toArray] + +@[simp] +theorem _root_.Array.toArray_toIter {array : Array β} : + array.iter.toArray = array := by + simp [Array.iter_eq_iterFromIdx, Array.toArray_iterFromIdxM] + +@[simp] +theorem _root_.Array.toListRev_iterFromIdx {array : Array β} {pos : Nat} : + (array.iterFromIdx pos).toListRev = (array.toList.drop pos).reverse := by + simp [Iter.toListRev_eq, Array.toList_iterFromIdx] + +@[simp] +theorem _root_.Array.toListRev_toIter {array : Array β} : + array.iter.toListRev = array.toListRev := by + simp [Array.iter_eq_iterFromIdx] + +end Std.Iterators diff --git a/src/Std/Data/Iterators/Lemmas/Producers/Monadic.lean b/src/Std/Data/Iterators/Lemmas/Producers/Monadic.lean index 7c070e36e7..82882a799d 100644 --- a/src/Std/Data/Iterators/Lemmas/Producers/Monadic.lean +++ b/src/Std/Data/Iterators/Lemmas/Producers/Monadic.lean @@ -4,4 +4,5 @@ Released under Apache 2.0 license as described in the file LICENSE. Authors: Paul Reichert -/ prelude +import Std.Data.Iterators.Lemmas.Producers.Monadic.Array import Std.Data.Iterators.Lemmas.Producers.Monadic.List diff --git a/src/Std/Data/Iterators/Lemmas/Producers/Monadic/Array.lean b/src/Std/Data/Iterators/Lemmas/Producers/Monadic/Array.lean new file mode 100644 index 0000000000..a391dea8e6 --- /dev/null +++ b/src/Std/Data/Iterators/Lemmas/Producers/Monadic/Array.lean @@ -0,0 +1,120 @@ +/- +Copyright (c) 2025 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Paul Reichert +-/ +prelude +import Std.Data.Iterators.Producers.Monadic.Array +import Std.Data.Iterators.Consumers +import Std.Data.Iterators.Lemmas.Consumers.Monadic +import Std.Data.Internal.LawfulMonadLiftFunction + +/-! +# Lemmas about array iterators + +This module provides lemmas about the interactions of `Array.iterM` with `IterM.step` and various +collectors. +-/ + +namespace Std.Iterators + +variable {m : Type w → Type w'} [Monad m] {β : Type w} + +theorem _root_.Array.iterM_eq_iterFromIdxM {array : Array β} : + array.iterM m = array.iterFromIdxM m 0 := + rfl + +theorem _root_.Array.step_iterFromIdxM {array : Array β} {pos : Nat} : + (array.iterFromIdxM m pos).step = (pure <| if h : pos < array.size then + .yield + (array.iterFromIdxM m (pos + 1)) + array[pos] + ⟨rfl, rfl, h, rfl⟩ + else + .done (Nat.not_lt.mp h)) := by + rfl + +theorem _root_.Array.step_iterM {array : Array β} : + (array.iterM m).step = (pure <| if h : 0 < array.size then + .yield + (array.iterFromIdxM m 1) + array[0] + ⟨rfl, rfl, h, rfl⟩ + else + .done (Nat.not_lt.mp h)) := by + rfl + +@[simp] +theorem _root_.Array.toList_iterFromIdxM [LawfulMonad m] {array : Array β} + {pos : Nat} : + (array.iterFromIdxM m pos).toList = pure (array.toList.drop pos) := by + by_cases h : pos < array.size + · suffices h' : ∀ n p, p ≥ array.size - n → (array.iterFromIdxM m p).toList = pure (array.toList.drop p) by + apply h' array.size + omega + intro n + induction n + · intro p hp + rw [IterM.toList_eq_match_step] + simp [Array.step_iterFromIdxM] + rw [List.drop_eq_nil_iff.mpr] + · simp [show ¬ p < array.size by omega] + · simp only [Array.length_toList] + omega + · rename_i n ih + intro p hp + by_cases h : p ≥ array.size - n + · apply ih + assumption + · rw [IterM.toList_eq_match_step, Array.step_iterFromIdxM] + simp [show p < array.size by omega] + rw [ih _ (by omega)] + simp + apply congrArg pure + rw (occs := [2]) [← List.getElem_cons_drop_succ_eq_drop] + simp + rw [Array.getElem_toList] + · rw [IterM.toList_eq_match_step, List.drop_eq_nil_iff.mpr] + · simp [Array.step_iterFromIdxM, h] + · simp only [Array.length_toList] + omega + +@[simp] +theorem _root_.Array.toList_iterM [LawfulMonad m] {array : Array β} : + (array.iterM m).toList = pure array.toList := by + simp [Array.iterM_eq_iterFromIdxM, Array.toList_iterFromIdxM] + +-- TODO: move to Init.Data.Array.Lemmas in a separate PR afterwards +private theorem _root_.List.drop_toArray' {l : List α} {k : Nat} : + l.toArray.drop k = (l.drop k).toArray := by + induction l generalizing k + case nil => simp + case cons l' ih => + match k with + | 0 => simp + | k' + 1 => simp [List.drop_succ_cons, ← ih] + +@[simp] +theorem _root_.Array.toArray_iterFromIdxM [LawfulMonad m] {array : Array β} {pos : Nat} : + (array.iterFromIdxM m pos).toArray = pure (array.extract pos) := by + simp [← IterM.toArray_toList, Array.toList_iterFromIdxM] + rw [← Array.drop_eq_extract] + rw (occs := [2]) [← Array.toArray_toList (xs := array)] + rw [List.drop_toArray'] + +@[simp] +theorem _root_.Array.toArray_toIterM [LawfulMonad m] {array : Array β} : + (array.iterM m).toArray = pure array := by + simp [Array.iterM_eq_iterFromIdxM, Array.toArray_iterFromIdxM] + +@[simp] +theorem _root_.Array.toListRev_iterFromIdxM [LawfulMonad m] {array : Array β} {pos : Nat} : + (array.iterFromIdxM m pos).toListRev = pure (array.toList.drop pos).reverse := by + simp [IterM.toListRev_eq, Array.toList_iterFromIdxM] + +@[simp] +theorem _root_.Array.toListRev_toIterM [LawfulMonad m] {array : Array β} : + (array.iterM m).toListRev = pure array.toListRev := by + simp [Array.iterM_eq_iterFromIdxM, Array.toListRev_iterFromIdxM] + +end Std.Iterators diff --git a/src/Std/Data/Iterators/Lemmas/Producers/Repeat.lean b/src/Std/Data/Iterators/Lemmas/Producers/Repeat.lean new file mode 100644 index 0000000000..9cbbdd16a3 --- /dev/null +++ b/src/Std/Data/Iterators/Lemmas/Producers/Repeat.lean @@ -0,0 +1,55 @@ +/- +Copyright (c) 2025 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Paul Reichert +-/ +prelude +import Init.Data.Option.Lemmas +import Std.Data.Iterators.Producers.Repeat +import Std.Data.Iterators.Consumers.Access +import Std.Data.Iterators.Consumers.Collect +import Std.Data.Iterators.Combinators.Take +import Std.Data.Iterators.Lemmas.Combinators.Take + +namespace Std.Iterators + +variable {α : Type w} {f : α → α} {init : α} + +theorem Iter.step_repeat : + (Iter.repeat f init).step = .yield (Iter.repeat f (f init)) init ⟨rfl, rfl⟩ := by + rfl + +theorem Iter.atIdxSlow?_zero_repeat : + (Iter.repeat f init).atIdxSlow? 0 = some init := by + rw [atIdxSlow?, step_repeat] + +theorem Iter.atIdxSlow?_succ_repeat {k : Nat} : + (Iter.repeat f init).atIdxSlow? (k + 1) = (Iter.repeat f (f init)).atIdxSlow? k := by + rw [atIdxSlow?, step_repeat] + +theorem Iter.atIdxSlow?_succ_repeat_eq_map {k : Nat} : + (Iter.repeat f init).atIdxSlow? (k + 1) = f <$> ((Iter.repeat f init).atIdxSlow? k) := by + rw [atIdxSlow?, step_repeat] + simp only + induction k generalizing init + · simp [atIdxSlow?_zero_repeat, Functor.map] + · simp [*, atIdxSlow?_succ_repeat] + +@[simp] +theorem Iter.atIdxSlow?_repeat {n : Nat} : + (Iter.repeat f init).atIdxSlow? n = some (Nat.repeat f n init) := by + induction n generalizing init + · apply atIdxSlow?_zero_repeat + · rename_i _ ih + simp [atIdxSlow?_succ_repeat_eq_map, ih, Nat.repeat] + +theorem Iter.isSome_atIdxSlow?_repeat {k : Nat} : + ((Iter.repeat f init).atIdxSlow? k).isSome := by + induction k generalizing init <;> simp [*, atIdxSlow?_succ_repeat] + +@[simp] +theorem Iter.toList_take_repeat_succ {k : Nat} : + ((Iter.repeat f init).take (k + 1)).toList = init :: ((Iter.repeat f (f init)).take k).toList := by + rw [toList_eq_match_step, step_take, step_repeat] + +end Std.Iterators diff --git a/src/Std/Data/Iterators/Producers.lean b/src/Std/Data/Iterators/Producers.lean index 5053bfde4b..8a8c3598a6 100644 --- a/src/Std/Data/Iterators/Producers.lean +++ b/src/Std/Data/Iterators/Producers.lean @@ -5,4 +5,6 @@ Authors: Paul Reichert -/ prelude import Std.Data.Iterators.Producers.Monadic +import Std.Data.Iterators.Producers.Array import Std.Data.Iterators.Producers.List +import Std.Data.Iterators.Producers.Repeat diff --git a/src/Std/Data/Iterators/Producers/Array.lean b/src/Std/Data/Iterators/Producers/Array.lean new file mode 100644 index 0000000000..c600f69854 --- /dev/null +++ b/src/Std/Data/Iterators/Producers/Array.lean @@ -0,0 +1,49 @@ +/- +Copyright (c) 2025 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Paul Reichert +-/ +prelude +import Std.Data.Iterators.Producers.Monadic.Array + +/-! +# Array iterator + +This module provides an iterator for arrays that is accessible via `Array.iter`. +-/ + +namespace Std.Iterators + +/-- +Returns a finite iterator for the given array starting at the given index. +The iterator yields the elements of the array in order and then terminates. + +The monadic version of this iterator is `Array.iterFromIdxM`. + +**Termination properties:** + +* `Finite` instance: always +* `Productive` instance: always +-/ +@[always_inline, inline] +def _root_.Array.iterFromIdx {α : Type w} (l : Array α) (pos : Nat) : + Iter (α := ArrayIterator α) α := + ((l.iterFromIdxM Id pos).toIter : Iter α) + +/-- +Returns a finite iterator for the given array. +The iterator yields the elements of the array in order and then terminates. + +The monadic version of this iterator is `Array.iterM`. + +**Termination properties:** + +* `Finite` instance: always +* `Productive` instance: always +-/ +@[always_inline, inline] +def _root_.Array.iter {α : Type w} (l : Array α) : + Iter (α := ArrayIterator α) α := + ((l.iterM Id).toIter : Iter α) + +end Std.Iterators diff --git a/src/Std/Data/Iterators/Producers/Monadic.lean b/src/Std/Data/Iterators/Producers/Monadic.lean index 9c869fc322..d7db666229 100644 --- a/src/Std/Data/Iterators/Producers/Monadic.lean +++ b/src/Std/Data/Iterators/Producers/Monadic.lean @@ -4,4 +4,5 @@ Released under Apache 2.0 license as described in the file LICENSE. Authors: Paul Reichert -/ prelude +import Std.Data.Iterators.Producers.Monadic.Array import Std.Data.Iterators.Producers.Monadic.List diff --git a/src/Std/Data/Iterators/Producers/Monadic/Array.lean b/src/Std/Data/Iterators/Producers/Monadic/Array.lean new file mode 100644 index 0000000000..db3a3161b4 --- /dev/null +++ b/src/Std/Data/Iterators/Producers/Monadic/Array.lean @@ -0,0 +1,120 @@ +/- +Copyright (c) 2025 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Paul Reichert +-/ +prelude +import Init.Data.Nat.Lemmas +import Init.RCases +import Std.Data.Iterators.Consumers +import Std.Data.Iterators.Internal.Termination + +/-! +# Array iterator + +This module provides an iterator for arrays that is accessible via `Array.iterM`. +-/ + +namespace Std.Iterators + +variable {α : Type w} {m : Type w → Type w'} + +/-- +The underlying state of a list iterator. Its contents are internal and should +not be used by downstream users of the library. +-/ +@[unbox] +structure ArrayIterator (α : Type w) where + /-- Internal implementation detail of the iterator library. -/ + array : Array α + /-- Internal implementation detail of the iterator library. -/ + pos : Nat + +/-- +Returns a finite monadic iterator for the given array starting at the given index. +The iterator yields the elements of the array in order and then terminates. + +The pure version of this iterator is `Array.iterFromIdx`. + +**Termination properties:** + +* `Finite` instance: always +* `Productive` instance: always +-/ +@[always_inline, inline] +def _root_.Array.iterFromIdxM {α : Type w} (array : Array α) (m : Type w → Type w') (pos : Nat) + [Pure m] : + IterM (α := ArrayIterator α) m α := + toIterM { array := array, pos := pos } m α + +/-- +Returns a finite monadic iterator for the given array. +The iterator yields the elements of the array in order and then terminates. There are no side +effects. + +The pure version of this iterator is `Array.iter`. + +**Termination properties:** + +* `Finite` instance: always +* `Productive` instance: always +-/ +@[always_inline, inline] +def _root_.Array.iterM {α : Type w} (array : Array α) (m : Type w → Type w') [Pure m] : + IterM (α := ArrayIterator α) m α := + array.iterFromIdxM m 0 + +@[always_inline, inline] +instance {α : Type w} [Pure m] : Iterator (ArrayIterator α) m α where + IsPlausibleStep it + | .yield it' out => it.internalState.array = it'.internalState.array ∧ + it'.internalState.pos = it.internalState.pos + 1 ∧ + ∃ _ : it.internalState.pos < it.internalState.array.size, + it.internalState.array[it.internalState.pos] = out + | .skip _ => False + | .done => it.internalState.pos ≥ it.internalState.array.size + step it := pure <| if h : it.internalState.pos < it.internalState.array.size then + .yield + ⟨⟨it.internalState.array, it.internalState.pos + 1⟩⟩ + it.internalState.array[it.internalState.pos] + ⟨rfl, rfl, h, rfl⟩ + else + .done (Nat.not_lt.mp h) + +private def ArrayIterator.finitenessRelation [Pure m] : + FinitenessRelation (ArrayIterator α) m where + rel := InvImage WellFoundedRelation.rel + (fun it => it.internalState.array.size - it.internalState.pos) + wf := InvImage.wf _ WellFoundedRelation.wf + subrelation {it it'} h := by + simp_wf + obtain ⟨step, h, h'⟩ := h + cases step + · cases h + obtain ⟨h, h', h'', rfl⟩ := h' + rw [h] at h'' + rw [h, h'] + omega + · cases h' + · cases h + +instance [Pure m] : Finite (ArrayIterator α) m := + Finite.of_finitenessRelation ArrayIterator.finitenessRelation + +@[always_inline, inline] +instance {α : Type w} [Monad m] [Monad n] : IteratorCollect (ArrayIterator α) m n := + .defaultImplementation + +@[always_inline, inline] +instance {α : Type w} [Monad m] [Monad n] : IteratorCollectPartial (ArrayIterator α) m n := + .defaultImplementation + +@[always_inline, inline] +instance {α : Type w} [Monad m] [Monad n] : IteratorLoop (ArrayIterator α) m n := + .defaultImplementation + +@[always_inline, inline] +instance {α : Type w} [Monad m] [Monad n] : IteratorLoopPartial (ArrayIterator α) m n := + .defaultImplementation + +end Std.Iterators diff --git a/src/Std/Data/Iterators/Producers/Repeat.lean b/src/Std/Data/Iterators/Producers/Repeat.lean new file mode 100644 index 0000000000..4b1b24eb76 --- /dev/null +++ b/src/Std/Data/Iterators/Producers/Repeat.lean @@ -0,0 +1,83 @@ +/- +Copyright (c) 2025 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Paul Reichert +-/ +prelude +import Std.Data.Iterators.Consumers.Monadic +import Std.Data.Iterators.Internal.Termination + +/-! +# Function-unfolding iterator + +This module provides an infinite iterator that given an initial value `init` function `f` emits +the iterates `init`, `f init`, `f (f init)`, ... . +-/ + +namespace Std.Iterators + +universe u v + +variable {α : Type w} {m : Type w → Type w'} {f : α → α} + +/-- +Internal state of the `repeat` combinator. Do not depend on its internals. +-/ +@[unbox] +structure RepeatIterator (α : Type u) (f : α → α) where + /-- Internal implementation detail of the iterator library. -/ + next : α + +@[always_inline, inline] +instance : Iterator (RepeatIterator α f) Id α where + IsPlausibleStep it + | .yield it' out => out = it.internalState.next ∧ it' = ⟨⟨f it.internalState.next⟩⟩ + | .skip _ => False + | .done => False + step it := pure <| .yield ⟨⟨f it.internalState.next⟩⟩ it.internalState.next (by simp) + +/-- +Creates an infinite iterator from an initial value `init` and a function `f : α → α`. +First it yields `init`, and in each successive step, the iterator applies `f` to the previous value. +So the iterator just emitted `a`, in the next step it will yield `f a`. In other words, the +`n`-th value is `Nat.repeat f n init`. + +For example, if `f := (· + 1)` and `init := 0`, then the iterator emits all natural numbers in +order. + +**Termination properties:** + +* `Finite` instance: not available and never possible +* `Productive` instance: always +-/ +@[always_inline, inline] +def Iter.repeat {α : Type w} (f : α → α) (init : α) := + (⟨RepeatIterator.mk (f := f) init⟩ : Iter α) + +private def RepeatIterator.instProductivenessRelation : + ProductivenessRelation (RepeatIterator α f) Id where + rel := emptyWf.rel + wf := emptyWf.wf + subrelation {it it'} h := by cases h + +instance RepeatIterator.instProductive : + Productive (RepeatIterator α f) Id := + Productive.of_productivenessRelation instProductivenessRelation + +instance RepeatIterator.instIteratorLoop {α : Type w} {f : α → α} {n : Type w → Type w'} [Monad n] : + IteratorLoop (RepeatIterator α f) Id n := + .defaultImplementation + +instance RepeatIterator.instIteratorLoopPartial {α : Type w} {f : α → α} {n : Type w → Type w'} + [Monad n] : IteratorLoopPartial (RepeatIterator α f) Id n := + .defaultImplementation + +instance RepeatIterator.instIteratorCollect {α : Type w} {f : α → α} {n : Type w → Type w'} + [Monad n] : IteratorCollect (RepeatIterator α f) Id n := + .defaultImplementation + +instance RepeatIterator.instIteratorCollectPartial {α : Type w} {f : α → α} {n : Type w → Type w'} + [Monad n] : IteratorCollectPartial (RepeatIterator α f) Id n := + .defaultImplementation + +end Std.Iterators diff --git a/tests/lean/run/iterators.lean b/tests/lean/run/iterators.lean index 509cf4450a..bd25bbdc42 100644 --- a/tests/lean/run/iterators.lean +++ b/tests/lean/run/iterators.lean @@ -55,6 +55,22 @@ example : ([1, 2, 3].iterM (StateM Nat)).toListRev = pure [3, 2, 1] := by end ListIteratorBasic +section Array + +example : #[1, 2, 3].iter.toArray = #[1, 2, 3] := by + simp + +example : #[1, 2, 3].iter.toList = [1, 2, 3] := by + simp + +example : #[1, 2, 3].iter.toListRev = [3, 2, 1] := by + simp + +example : (#[1, 2, 3].iterFromIdx 2).toList = [3] := by + simp + +end Array + section WellFoundedRecursion def sum (l : List Nat) : Nat := @@ -165,6 +181,21 @@ fun #guard_msgs in #eval ["Lean", "is", "fun"].iter.mapM (IO.println s!"{·}") |>.drain +-- This example demonstrates that chained `mapM` calls are executed in a different order than with `List.mapM`. +def chainedMapM (l : List Nat) : IO Unit := + l.iterM IO |>.mapM (IO.println <| s!"1st {.}") |>.mapM (IO.println <| s!"2nd {·}") |>.drain + +/-- +info: 1st 1 +2nd () +1st 2 +2nd () +1st 3 +2nd () +-/ +#guard_msgs in +#eval! chainedMapM [1, 2, 3] + end FilterMap section Zip @@ -208,3 +239,22 @@ example : ([1, 2, 3, 4].iter.dropWhile (· ≠ 3)).toListRev = [4, 3] := by simp end DropWhile + +section Repeat + +@[simp] +def positives := Std.Iterators.Iter.repeat (init := 1) (· + 1) + +example : (positives.take 5).toList = [1, 2, 3, 4, 5] := by + simp + +@[simp] +def divisors (n : Nat) := (positives.take n |>.filter (n % · = 0)) + +@[simp] +def isPrime (n : Nat) : Bool := (divisors n).toList.length == 2 + +example : isPrime 5 := by + simp + +end Repeat