From a12f89aefa10812fda01275d87bfc18d97914899 Mon Sep 17 00:00:00 2001 From: Paul Reichert <6992158+datokrat@users.noreply.github.com> Date: Fri, 30 May 2025 12:34:12 +0200 Subject: [PATCH] feat: introduce `take` iterator combinator (#8418) This PR provides the `take` iterator combinator that transforms any iterator into an iterator that stops after a given number of steps. The change contains the implementation and lemmas. `take` has a special implementation of `IteratorLoop` that relies on a potentially more efficient `forIn` implementation of the inner iterator. The mysterious `@[specialize]` on a test has been removed because it is not necessary anymore according to a manual inspection of the IR. Either I erroneously concluded from experiments that it was necessary of something has changed in the meantime that makes it unnecessary. --- src/Std/Data/Iterators.lean | 3 +- src/Std/Data/Iterators/Combinators.lean | 8 + .../Data/Iterators/Combinators/Monadic.lean | 7 + .../Iterators/Combinators/Monadic/Take.lean | 198 ++++++++++++++++++ src/Std/Data/Iterators/Combinators/Take.lean | 39 ++++ src/Std/Data/Iterators/Lemmas.lean | 1 + .../Data/Iterators/Lemmas/Combinators.lean | 8 + .../Iterators/Lemmas/Combinators/Monadic.lean | 7 + .../Lemmas/Combinators/Monadic/Take.lean | 30 +++ .../Iterators/Lemmas/Combinators/Take.lean | 86 ++++++++ tests/lean/run/iterators.lean | 38 ++++ 11 files changed, 424 insertions(+), 1 deletion(-) create mode 100644 src/Std/Data/Iterators/Combinators.lean create mode 100644 src/Std/Data/Iterators/Combinators/Monadic.lean create mode 100644 src/Std/Data/Iterators/Combinators/Monadic/Take.lean create mode 100644 src/Std/Data/Iterators/Combinators/Take.lean create mode 100644 src/Std/Data/Iterators/Lemmas/Combinators.lean create mode 100644 src/Std/Data/Iterators/Lemmas/Combinators/Monadic.lean create mode 100644 src/Std/Data/Iterators/Lemmas/Combinators/Monadic/Take.lean create mode 100644 src/Std/Data/Iterators/Lemmas/Combinators/Take.lean diff --git a/src/Std/Data/Iterators.lean b/src/Std/Data/Iterators.lean index fa3de7db30..067844988b 100644 --- a/src/Std/Data/Iterators.lean +++ b/src/Std/Data/Iterators.lean @@ -7,8 +7,9 @@ prelude import Std.Data.Iterators.Basic import Std.Data.Iterators.Producers import Std.Data.Iterators.Consumers -import Std.Data.Iterators.Internal +import Std.Data.Iterators.Combinators import Std.Data.Iterators.Lemmas +import Std.Data.Iterators.Internal /-! # Iterators diff --git a/src/Std/Data/Iterators/Combinators.lean b/src/Std/Data/Iterators/Combinators.lean new file mode 100644 index 0000000000..5fb02bf61c --- /dev/null +++ b/src/Std/Data/Iterators/Combinators.lean @@ -0,0 +1,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 +-/ +prelude +import Std.Data.Iterators.Combinators.Monadic +import Std.Data.Iterators.Combinators.Take diff --git a/src/Std/Data/Iterators/Combinators/Monadic.lean b/src/Std/Data/Iterators/Combinators/Monadic.lean new file mode 100644 index 0000000000..ae87ad2b80 --- /dev/null +++ b/src/Std/Data/Iterators/Combinators/Monadic.lean @@ -0,0 +1,7 @@ +/- +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.Combinators.Monadic.Take diff --git a/src/Std/Data/Iterators/Combinators/Monadic/Take.lean b/src/Std/Data/Iterators/Combinators/Monadic/Take.lean new file mode 100644 index 0000000000..a654294432 --- /dev/null +++ b/src/Std/Data/Iterators/Combinators/Monadic/Take.lean @@ -0,0 +1,198 @@ +/- +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.Basic +import Std.Data.Iterators.Consumers.Monadic.Collect +import Std.Data.Iterators.Consumers.Monadic.Loop +import Std.Data.Iterators.Internal.Termination + +/-! +This module provides the iterator combinator `IterM.take`. +-/ + +namespace Std.Iterators + +variable {α : Type w} {m : Type w → Type w'} {β : Type w} + +/-- +The internal state of the `IterM.take` iterator combinator. +-/ +@[unbox] +structure Take (α : Type w) (m : Type w → Type w') (β : Type w) where + /-- Internal implementation detail of the iterator library -/ + remaining : Nat + /-- Internal implementation detail of the iterator library -/ + inner : IterM (α := α) m β + +/-- +Given an iterator `it` and a natural number `n`, `it.take n` is an iterator that outputs +up to the first `n` of `it`'s values in order and then terminates. + +**Marble diagram:** + +```text +it ---a----b---c--d-e--⊥ +it.take 3 ---a----b---c⊥ + +it ---a--⊥ +it.take 3 ---a--⊥ +``` + +**Termination properties:** + +* `Finite` instance: only if `it` is productive +* `Productive` instance: only if `it` is productive + +**Performance:** + +This combinator incurs an additional O(1) cost with each output of `it`. +-/ +@[inline] +def IterM.take (n : Nat) (it : IterM (α := α) m β) := + toIterM (Take.mk n it) m β + +theorem IterM.take.surjective {α : Type w} {m : Type w → Type w'} {β : Type w} + (it : IterM (α := Take α m β) m β) : + ∃ (it₀ : IterM (α := α) m β) (k : Nat), it = it₀.take k := by + refine ⟨it.internalState.inner, it.internalState.remaining, rfl⟩ + +inductive Take.PlausibleStep [Iterator α m β] (it : IterM (α := Take α m β) m β) : + (step : IterStep (IterM (α := Take α m β) m β) β) → Prop where + | yield : ∀ {it' out k}, it.internalState.inner.IsPlausibleStep (.yield it' out) → + it.internalState.remaining = k + 1 → PlausibleStep it (.yield (it'.take k) out) + | skip : ∀ {it' k}, it.internalState.inner.IsPlausibleStep (.skip it') → + it.internalState.remaining = k + 1 → PlausibleStep it (.skip (it'.take (k + 1))) + | done : it.internalState.inner.IsPlausibleStep .done → PlausibleStep it .done + | depleted : it.internalState.remaining = 0 → + PlausibleStep it .done + +@[always_inline, inline] +instance Take.instIterator [Monad m] [Iterator α m β] : Iterator (Take α m β) m β where + IsPlausibleStep := Take.PlausibleStep + step it := + match h : it.internalState.remaining with + | 0 => pure <| .done (.depleted h) + | k + 1 => do + match ← it.internalState.inner.step with + | .yield it' out h' => pure <| .yield (it'.take k) out (.yield h' h) + | .skip it' h' => pure <| .skip (it'.take (k + 1)) (.skip h' h) + | .done h' => pure <| .done (.done h') + +def Take.Rel (m : Type w → Type w') [Monad m] [Iterator α m β] [Productive α m] : + IterM (α := Take α m β) m β → IterM (α := Take α m β) m β → Prop := + InvImage (Prod.Lex Nat.lt_wfRel.rel IterM.TerminationMeasures.Productive.Rel) + (fun it => (it.internalState.remaining, it.internalState.inner.finitelyManySkips)) + +theorem Take.rel_of_remaining [Monad m] [Iterator α m β] [Productive α m] + {it it' : IterM (α := Take α m β) m β} + (h : it'.internalState.remaining < it.internalState.remaining) : Take.Rel m it' it := + Prod.Lex.left _ _ h + +theorem Take.rel_of_inner [Monad m] [Iterator α m β] [Productive α m] {remaining : Nat} + {it it' : IterM (α := α) m β} + (h : it'.finitelyManySkips.Rel it.finitelyManySkips) : + Take.Rel m (it'.take remaining) (it.take remaining) := + Prod.Lex.right _ h + +private def Take.instFinitenessRelation [Monad m] [Iterator α m β] + [Productive α m] : + FinitenessRelation (Take α m β) m where + rel := Take.Rel m + wf := by + apply InvImage.wf + refine ⟨fun (a, b) => Prod.lexAccessible (WellFounded.apply ?_ a) (WellFounded.apply ?_) b⟩ + · exact WellFoundedRelation.wf + · exact WellFoundedRelation.wf + subrelation {it it'} h := by + obtain ⟨step, h, h'⟩ := h + cases h' + case yield it' out k h' h'' => + cases h + apply rel_of_remaining + simp_all [IterM.take] + case skip it' out k h' h'' => + cases h + obtain ⟨it, k, rfl⟩ := IterM.take.surjective it + cases h'' + apply Take.rel_of_inner + exact IterM.TerminationMeasures.Productive.rel_of_skip h' + case done _ => + cases h + case depleted _ => + cases h + +instance Take.instFinite [Monad m] [Iterator α m β] [Productive α m] : + Finite (Take α m β) m := + Finite.of_finitenessRelation instFinitenessRelation + +instance Take.instIteratorCollect [Monad m] [Iterator α m β] [Productive α m] : + IteratorCollect (Take α m β) m := + .defaultImplementation + +instance Take.instIteratorCollectPartial [Monad m] [Iterator α m β] : + IteratorCollectPartial (Take α m β) m := + .defaultImplementation + +private def Take.PlausibleForInStep {β : Type u} {γ : Type v} + (f : β → γ → ForInStep γ → Prop) : + β → γ × Nat → (ForInStep (γ × Nat)) → Prop + | out, (c, n), ForInStep.yield (c', n') => n = n' + 1 ∧ f out c (.yield c') + | _, _, .done _ => True + +private def Take.wellFounded_plausibleForInStep {α β : Type w} {m : Type w → Type w'} + [Monad m] [Iterator α m β] {γ : Type x} + {f : β → γ → ForInStep γ → Prop} (wf : IteratorLoop.WellFounded (Take α m β) m f) : + IteratorLoop.WellFounded α m (PlausibleForInStep f) := by + simp only [IteratorLoop.WellFounded] at ⊢ wf + letI : WellFoundedRelation _ := ⟨_, wf⟩ + apply Subrelation.wf + (r := InvImage WellFoundedRelation.rel fun p => (p.1.take (p.2.2 + 1), p.2.1)) + (fun {p q} h => by + simp only [InvImage, WellFoundedRelation.rel, this, IteratorLoop.rel, + IterM.IsPlausibleStep, Iterator.IsPlausibleStep] + obtain ⟨out, h, h'⟩ | ⟨h, h'⟩ := h + · apply Or.inl + exact ⟨out, .yield h (by simp only [IterM.take, internalState_toIterM, + Nat.add_right_cancel_iff, this]; exact h'.1), h'.2⟩ + · apply Or.inr + refine ⟨?_, by rw [h']⟩ + rw [h'] + apply PlausibleStep.skip + · exact h + · rfl) + apply InvImage.wf + exact WellFoundedRelation.wf + +instance Take.instIteratorFor [Monad m] [Monad n] [Iterator α m β] + [IteratorLoop α m n] [MonadLiftT m n] : + IteratorLoop (Take α m β) m n where + forIn lift {γ} Plausible wf it init f := by + refine Prod.fst <$> IteratorLoop.forIn lift (γ := γ × Nat) + (PlausibleForInStep Plausible) + (wellFounded_plausibleForInStep wf) + it.internalState.inner + (init, it.internalState.remaining) + fun out acc => + match h : acc.snd with + | 0 => pure <| ⟨.done acc, True.intro⟩ + | n + 1 => (fun + | ⟨.yield x, hp⟩ => ⟨.yield ⟨x, n⟩, ⟨h, hp⟩⟩ + | ⟨.done x ,hp⟩ => ⟨.done ⟨x, n⟩, .intro⟩) <$> f out acc.fst + +instance Take.instIteratorForPartial [Monad m] [Monad n] [Iterator α m β] + [IteratorLoopPartial α m n] [MonadLiftT m n] : + IteratorLoopPartial (Take α m β) m n where + forInPartial lift {γ} it init f := do + Prod.fst <$> IteratorLoopPartial.forInPartial lift it.internalState.inner (γ := γ × Nat) + (init, it.internalState.remaining) + fun out acc => + match acc.snd with + | 0 => pure <| .done acc + | n + 1 => (fun | .yield x => .yield ⟨x, n⟩ | .done x => .done ⟨x, n⟩) <$> f out acc.fst + +end Std.Iterators diff --git a/src/Std/Data/Iterators/Combinators/Take.lean b/src/Std/Data/Iterators/Combinators/Take.lean new file mode 100644 index 0000000000..8e361a6029 --- /dev/null +++ b/src/Std/Data/Iterators/Combinators/Take.lean @@ -0,0 +1,39 @@ +/- +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.Combinators.Monadic.Take + +namespace Std.Iterators + +/-- +Given an iterator `it` and a natural number `n`, `it.take n` is an iterator that outputs +up to the first `n` of `it`'s values in order and then terminates. + +**Marble diagram:** + +```text +it ---a----b---c--d-e--⊥ +it.take 3 ---a----b---c⊥ + +it ---a--⊥ +it.take 3 ---a--⊥ +``` + +**Termination properties:** + +* `Finite` instance: only if `it` is productive +* `Productive` instance: only if `it` is productive + +**Performance:** + +This combinator incurs an additional O(1) cost with each output of `it`. +-/ +@[always_inline, inline] +def Iter.take {α : Type w} {β : Type w} (n : Nat) (it : Iter (α := α) β) : + Iter (α := Take α Id β) β := + it.toIterM.take n |>.toIter + +end Std.Iterators diff --git a/src/Std/Data/Iterators/Lemmas.lean b/src/Std/Data/Iterators/Lemmas.lean index 4f2229633b..24daf75518 100644 --- a/src/Std/Data/Iterators/Lemmas.lean +++ b/src/Std/Data/Iterators/Lemmas.lean @@ -7,4 +7,5 @@ prelude import Std.Data.Iterators.Lemmas.Basic import Std.Data.Iterators.Lemmas.Monadic import Std.Data.Iterators.Lemmas.Consumers +import Std.Data.Iterators.Lemmas.Combinators import Std.Data.Iterators.Lemmas.Producers diff --git a/src/Std/Data/Iterators/Lemmas/Combinators.lean b/src/Std/Data/Iterators/Lemmas/Combinators.lean new file mode 100644 index 0000000000..1d910d9649 --- /dev/null +++ b/src/Std/Data/Iterators/Lemmas/Combinators.lean @@ -0,0 +1,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 +-/ +prelude +import Std.Data.Iterators.Lemmas.Combinators.Monadic +import Std.Data.Iterators.Lemmas.Combinators.Take diff --git a/src/Std/Data/Iterators/Lemmas/Combinators/Monadic.lean b/src/Std/Data/Iterators/Lemmas/Combinators/Monadic.lean new file mode 100644 index 0000000000..b8922fcc19 --- /dev/null +++ b/src/Std/Data/Iterators/Lemmas/Combinators/Monadic.lean @@ -0,0 +1,7 @@ +/- +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.Combinators.Monadic.Take diff --git a/src/Std/Data/Iterators/Lemmas/Combinators/Monadic/Take.lean b/src/Std/Data/Iterators/Lemmas/Combinators/Monadic/Take.lean new file mode 100644 index 0000000000..10a36f9805 --- /dev/null +++ b/src/Std/Data/Iterators/Lemmas/Combinators/Monadic/Take.lean @@ -0,0 +1,30 @@ +/- +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.Combinators.Monadic.Take +import Std.Data.Iterators.Lemmas.Consumers.Monadic + +namespace Std.Iterators + +theorem IterM.step_take {α m β} [Monad m] [Iterator α m β] {n : Nat} + {it : IterM (α := α) m β} : + (it.take n).step = (match n with + | 0 => pure <| .done (.depleted rfl) + | k + 1 => do + match ← it.step with + | .yield it' out h => pure <| .yield (it'.take k) out (.yield h rfl) + | .skip it' h => pure <| .skip (it'.take (k + 1)) (.skip h rfl) + | .done h => pure <| .done (.done h)) := by + simp only [take, step, Iterator.step, internalState_toIterM, Nat.succ_eq_add_one] + cases n + case zero => rfl + case succ k => + apply bind_congr + intro step + obtain ⟨step, h⟩ := step + cases step <;> rfl + +end Std.Iterators diff --git a/src/Std/Data/Iterators/Lemmas/Combinators/Take.lean b/src/Std/Data/Iterators/Lemmas/Combinators/Take.lean new file mode 100644 index 0000000000..58836cf851 --- /dev/null +++ b/src/Std/Data/Iterators/Lemmas/Combinators/Take.lean @@ -0,0 +1,86 @@ +/- +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.Combinators.Take +import Std.Data.Iterators.Consumers.Access +import Std.Data.Iterators.Lemmas.Combinators.Monadic.Take +import Std.Data.Iterators.Lemmas.Consumers + +namespace Std.Iterators + +theorem Iter.take_eq {α β} [Iterator α Id β] {n : Nat} + {it : Iter (α := α) β} : + it.take n = (it.toIterM.take n).toIter := + rfl + +theorem Iter.step_take {α β} [Iterator α Id β] {n : Nat} + {it : Iter (α := α) β} : + (it.take n).step = (match n with + | 0 => .done (.depleted rfl) + | k + 1 => + match it.step with + | .yield it' out h => .yield (it'.take k) out (.yield h rfl) + | .skip it' h => .skip (it'.take (k + 1)) (.skip h rfl) + | .done h => .done (.done h)) := by + simp only [Iter.step, Iter.step, Iter.take_eq, IterM.step_take, toIterM_toIter] + cases n + case zero => simp [PlausibleIterStep.done] + case succ k => + simp only [Id.run_bind] + generalize it.toIterM.step.run = step + cases step using PlausibleIterStep.casesOn <;> + simp [PlausibleIterStep.yield, PlausibleIterStep.skip, PlausibleIterStep.done] + +theorem Iter.atIdxSlow?_take {α β} + [Iterator α Id β] [Productive α Id] {k l : Nat} + {it : Iter (α := α) β} : + (it.take k).atIdxSlow? l = if l < k then it.atIdxSlow? l else none := by + fun_induction it.atIdxSlow? l generalizing k + case case1 it it' out h h' => + simp only [atIdxSlow?.eq_def (it := it.take k), step_take, h'] + cases k <;> simp + case case2 it it' out h h' l ih => + simp only [Nat.succ_eq_add_one, atIdxSlow?.eq_def (it := it.take k), step_take, h'] + cases k <;> cases l <;> simp [ih] + case case3 l it it' h h' ih => + simp only [atIdxSlow?.eq_def (it := it.take k), step_take, h'] + cases k <;> cases l <;> simp [ih] + case case4 l it h h' => + simp only [atIdxSlow?.eq_def (it := it.take k), atIdxSlow?.eq_def (it := it), step_take, h'] + cases k <;> cases l <;> simp + +@[simp] +theorem Iter.toList_take_of_finite {α β} [Iterator α Id β] {n : Nat} + [Finite α Id] [IteratorCollect α Id] [LawfulIteratorCollect α Id] + {it : Iter (α := α) β} : + (it.take n).toList = it.toList.take n := by + induction it using Iter.inductSteps generalizing n with | step it ihy ihs => + rw [Iter.toList_eq_match_step, Iter.toList_eq_match_step, Iter.step_take] + cases n + case zero => simp + case succ k => + simp + obtain ⟨step, h⟩ := it.step + cases step + · simp [ihy h] + · simp [ihs h] + · simp + +@[simp] +theorem Iter.toListRev_take_of_finite {α β} [Iterator α Id β] {n : Nat} + [Finite α Id] [IteratorCollect α Id] [LawfulIteratorCollect α Id] + {it : Iter (α := α) β} : + (it.take n).toListRev = it.toListRev.drop (it.toList.length - n) := by + rw [toListRev_eq, toList_take_of_finite, List.reverse_take, toListRev_eq] + +@[simp] +theorem Iter.toArray_take_of_finite {α β} [Iterator α Id β] {n : Nat} + [Finite α Id] [IteratorCollect α Id] [LawfulIteratorCollect α Id] + {it : Iter (α := α) β} : + (it.take n).toArray = it.toArray.take n := by + rw [← toArray_toList, ← toArray_toList, List.take_toArray, toList_take_of_finite] + +end Std.Iterators diff --git a/tests/lean/run/iterators.lean b/tests/lean/run/iterators.lean index b3baf07a38..1717538aa9 100644 --- a/tests/lean/run/iterators.lean +++ b/tests/lean/run/iterators.lean @@ -106,3 +106,41 @@ example (l : List Nat) : simp end Loop + +section Take + +def sumTakeRec (l : List Nat) : Nat := + go (l.iter.take 2) 0 +where + go it acc := + match it.step with + | .yield it' out _ => go it' (acc + out) + | .skip it' _ => go it' acc + | .done _ => acc + termination_by it.finitelyManySteps + +def sumTakeFold (l : List Nat) : Nat := + l.iter.take 2 |>.fold (init := 0) (· + ·) + +/-- info: [1, 2] -/ +#guard_msgs in +#eval [1, 2, 3].iter.take 2 |>.toList + +/-- info: 3 -/ +#guard_msgs in +#eval sumTakeRec [1, 2, 3] + +/-- info: 3 -/ +#guard_msgs in +#eval sumTakeFold [1, 2, 3] + +example : ([1, 2, 3].iter.take 2).toList = [1, 2] := by + simp + +example : ([1, 2, 3].iter.take 2).toArray = #[1, 2] := by + simp + +example : ([1, 2, 3].iter.take 2).toListRev = [2, 1] := by + simp + +end Take