From ba0e755adc73d771420afaa47811a4e5f1bfcb56 Mon Sep 17 00:00:00 2001 From: Markus Himmel Date: Tue, 27 Jan 2026 13:10:16 +0100 Subject: [PATCH] feat: `Std.Iter.first?` (#12162) This PR adds the function `Std.Iter.first?` and proves the specification lemma `Std.Iter.first?_eq_match_step` if the iterator is productive. The monadic variant on `Std.IterM` is also provided. We use this new function to fix the default implementation for `startsWith` and `dropPrefix` on `String` patterns, which used to fail if the searcher returned a `skip` at the beginning. None of the patterns we ship out of the box were affected by this, but user-defined patterns were vulnerable. --------- Co-authored-by: Paul Reichert <6992158+datokrat@users.noreply.github.com> --- src/Init/Data/Iterators/Consumers/Loop.lean | 37 +++++++++++++ .../Iterators/Consumers/Monadic/Loop.lean | 52 ++++++++++++++++++- .../Internal/LawfulMonadLiftFunction.lean | 5 ++ .../Data/Iterators/Lemmas/Consumers/Loop.lean | 22 ++++++++ .../Lemmas/Consumers/Monadic/Loop.lean | 31 ++++++++++- src/Init/Data/String/Pattern/Basic.lean | 39 +++++++------- 6 files changed, 165 insertions(+), 21 deletions(-) diff --git a/src/Init/Data/Iterators/Consumers/Loop.lean b/src/Init/Data/Iterators/Consumers/Loop.lean index b5191d4309..34353d0dd3 100644 --- a/src/Init/Data/Iterators/Consumers/Loop.lean +++ b/src/Init/Data/Iterators/Consumers/Loop.lean @@ -630,6 +630,43 @@ def Iter.Total.find? {α β : Type w} [Iterator α Id β] [IteratorLoop α Id Id (it : Iter.Total (α := α) β) (f : β → Bool) : Option β := it.it.find? f +/-- +Returns the first output of the iterator, or `none` if no such output is found. + +`O(|it|)` since the iterator may skip an unknown number of times before returning a result. +Short-circuits upon encountering the first result. Only the first element of `it` is examined. + +If the iterator is not productive, this function might run forever. The variant +`it.ensureTermination.first?` always terminates after finitely many steps. + +Examples: +* `[7, 6].iter.first? = some 7` +* `[].iter.first? = none` +-/ +@[inline] +def Iter.first? {α β : Type w} [Iterator α Id β] [IteratorLoop α Id Id] + (it : Iter (α := α) β) : Option β := + it.toIterM.first?.run + +/-- +Returns the first output of the iterator, or `none` if no such output is found. + +`O(|it|)` since the iterator may skip an unknown number of times before returning a result. +Short-circuits upon encountering the first result. The elements in `it` are examined in order of +iteration. + +This variant terminates after finitely many steps and requires a proof that the iterator is +productive. If such a proof is not available, consider using `Iter.first?`. + +Examples: +* `[7, 6].iter.first? = some 7` +* `[].iter.first? = none` +-/ +@[inline] +def Iter.Total.first? {α β : Type w} [Iterator α Id β] [IteratorLoop α Id Id] [Productive α Id] + (it : Iter.Total (α := α) β) : Option β := + it.it.first? + /-- Steps through the whole iterator, counting the number of outputs emitted. diff --git a/src/Init/Data/Iterators/Consumers/Monadic/Loop.lean b/src/Init/Data/Iterators/Consumers/Monadic/Loop.lean index 6c89c180cf..3bbb537f2f 100644 --- a/src/Init/Data/Iterators/Consumers/Monadic/Loop.lean +++ b/src/Init/Data/Iterators/Consumers/Monadic/Loop.lean @@ -185,8 +185,8 @@ instance instLawfulIteratorLoopDefaultImplementation (α : Type w) (m : Type w constructor; simp theorem IteratorLoop.wellFounded_of_finite {m : Type w → Type w'} - {α β : Type w} {γ : Type x} [Iterator α m β] [Finite α m] : - WellFounded α m (γ := γ) fun _ _ _ => True := by + {α β : Type w} {γ : Type x} [Iterator α m β] [Finite α m] {P : β → γ → ForInStep γ → Prop} : + WellFounded α m (γ := γ) P := by apply Subrelation.wf (r := InvImage IterM.TerminationMeasures.Finite.Rel (fun p => p.1.finitelyManySteps)) · intro p' p h @@ -197,6 +197,16 @@ theorem IteratorLoop.wellFounded_of_finite {m : Type w → Type w'} · apply InvImage.wf exact WellFoundedRelation.wf +theorem IteratorLoop.wellFounded_of_productive {α β : Type w} {m : Type w → Type w'} + [Iterator α m β] [IteratorLoop α m m] [Productive α m] {P : β → γ → ForInStep γ → Prop} + (hp : ∀ {b g s}, P b g s → s matches ForInStep.done ..) : + WellFounded α m (γ := γ) P := by + rw [WellFounded] + unfold IteratorLoop.rel + have {b g q} : ¬ P b g (ForInStep.yield q) := fun h => by simpa using hp h + simp only [and_false, exists_false, false_or, this] + exact Subrelation.wf And.left (InvImage.wf Prod.fst Productive.wf) + /-- This `ForIn'`-style loop construct traverses a finite iterator using an `IteratorLoop` instance. -/ @@ -902,6 +912,44 @@ def IterM.Total.find? {α β : Type w} {m : Type w → Type w'} [Monad m] [Itera m (Option β) := it.it.find? f +/-- +Returns the first output of the iterator, or `none` if no such output is found. + +`O(|it|)` since the iterator may skip an unknown number of times before returning a result. +Short-circuits upon encountering the first result. Only the first element of `it` is examined. + +If the iterator is not productive, this function might run forever. The variant +`it.ensureTermination.first?` always terminates after finitely many steps. + +Examples: +* `([7, 6].iterM Id).first? = pure (some 7)` +* `([].iterM Id).first? = pure none` +-/ +@[inline] +def IterM.first? {α β : Type w} {m : Type w → Type w'} [Monad m] [Iterator α m β] + [IteratorLoop α m m] (it : IterM (α := α) m β) : m (Option β) := + IteratorLoop.forIn (fun _ _ => flip Bind.bind) _ (fun b _ s => s = ForInStep.done (some b)) it + none (fun b _ _ => pure ⟨ForInStep.done (some b), rfl⟩) + +/-- +Returns the first output of the iterator, or `none` if no such output is found. + +`O(|it|)` since the iterator may skip an unknown number of times before returning a result. +Short-circuits upon encountering the first result. The elements in `it` are examined in order of +iteration. + +This variant terminates after finitely many steps and requires a proof that the iterator is +productive. If such a proof is not available, consider using `IterM.first?`. + +Examples: +* `([7, 6].iterM Id).first? = pure (some 7)` +* `([].iterM Id).first? = pure none` +-/ +@[inline] +def IterM.Total.first? {α β : Type w} {m : Type w → Type w'} [Monad m] [Iterator α m β] + [IteratorLoop α m m] [Productive α m] (it : IterM.Total (α := α) m β) : m (Option β) := + it.it.first? + section Count /-- diff --git a/src/Init/Data/Iterators/Internal/LawfulMonadLiftFunction.lean b/src/Init/Data/Iterators/Internal/LawfulMonadLiftFunction.lean index aa1be20a8c..7bdc8f04e1 100644 --- a/src/Init/Data/Iterators/Internal/LawfulMonadLiftFunction.lean +++ b/src/Init/Data/Iterators/Internal/LawfulMonadLiftFunction.lean @@ -111,6 +111,11 @@ instance {n : Type u → Type w} [Monad n] [LawfulMonad n] : liftBind_pure := by simp liftBind_bind := by simp +instance {m : Type u → Type v} [Monad m] [LawfulMonad m] : + LawfulMonadLiftBindFunction (m := m) (n := m) (fun _ _ => flip Bind.bind) where + liftBind_pure := by simp [flip] + liftBind_bind := by simp [flip] + end LiftBind end Std.Internal diff --git a/src/Init/Data/Iterators/Lemmas/Consumers/Loop.lean b/src/Init/Data/Iterators/Lemmas/Consumers/Loop.lean index f871300c8b..74e0e47384 100644 --- a/src/Init/Data/Iterators/Lemmas/Consumers/Loop.lean +++ b/src/Init/Data/Iterators/Lemmas/Consumers/Loop.lean @@ -915,4 +915,26 @@ theorem Iter.findM?_pure {α β : Type w} {m : Type w → Type w'} [Monad m] · simp [ihs ‹_›] · simp +theorem Iter.first?_eq_first?_toIterM {α β : Type w} [Iterator α Id β] [IteratorLoop α Id Id] + {it : Iter (α := α) β} : + it.first? = it.toIterM.first?.run := (rfl) + +theorem Iter.first?_eq_match_step {α β : Type w} [Iterator α Id β] [IteratorLoop α Id Id] + [Productive α Id] [LawfulIteratorLoop α Id Id] {it : Iter (α := α) β} : + it.first? = match it.step.val with + | .yield _ out => some out + | .skip it' => it'.first? + | .done => none := by + rw [Iter.first?_eq_first?_toIterM, IterM.first?_eq_match_step] + simp only [Id.run_bind, step] + generalize it.toIterM.step.run.inflate = s + rcases s with ⟨_|_|_, _⟩ <;> simp [Iter.first?_eq_first?_toIterM] + +theorem Iter.first?_eq_head?_toList {α β : Type w} [Iterator α Id β] [IteratorLoop α Id Id] + [Finite α Id] [LawfulIteratorLoop α Id Id] {it : Iter (α := α) β} : + it.first? = it.toList.head? := by + induction it using Iter.inductSteps with | step it ihy ihs + rw [first?_eq_match_step, toList_eq_match_step] + cases it.step using PlausibleIterStep.casesOn <;> simp [*] + end Std diff --git a/src/Init/Data/Iterators/Lemmas/Consumers/Monadic/Loop.lean b/src/Init/Data/Iterators/Lemmas/Consumers/Monadic/Loop.lean index f6c514c15e..1599b1e560 100644 --- a/src/Init/Data/Iterators/Lemmas/Consumers/Monadic/Loop.lean +++ b/src/Init/Data/Iterators/Lemmas/Consumers/Monadic/Loop.lean @@ -15,6 +15,15 @@ public section namespace Std open Std.Iterators +theorem IterM.DefaultConsumers.forIn_eq {α β : Type w} {m : Type w → Type w'} + {n : Type x → Type x'} [Monad n] [Iterator α m β] + {lift : (γ : Type w) → (δ : Type x) → (γ → n δ) → m γ → n δ} + {plausible_forInStep : β → γ → ForInStep γ → Prop} {it : IterM (α := α) m β} {init : γ} + {f : (b : β) → it.IsPlausibleIndirectOutput b → (c : γ) → n (Subtype (plausible_forInStep b c))} : + letI : IteratorLoop α m n := .defaultImplementation + IteratorLoop.forIn lift γ plausible_forInStep it init f = + IterM.DefaultConsumers.forIn' lift γ plausible_forInStep it init _ (fun _ => id) f := rfl + theorem IterM.DefaultConsumers.forIn'_eq_match_step {α β : Type w} {m : Type w → Type w'} [Iterator α m β] {n : Type x → Type x'} [Monad n] [LawfulMonad n] {lift : ∀ γ δ, (γ → n δ) → m γ → n δ} {γ : Type x} @@ -731,7 +740,7 @@ theorem IterM.findSomeM?_eq_match_step {α β γ : Type w} {m : Type w → Type · simp theorem IterM.findSome?_eq_findSomeM? {α β γ : Type w} {m : Type w → Type w'} [Monad m] - [Iterator α m β] [IteratorLoop α m m] [Finite α m] + [Iterator α m β] [IteratorLoop α m m] {it : IterM (α := α) m β} {f : β → Option γ} : it.findSome? f = it.findSomeM? (pure <| f ·) := (rfl) @@ -832,4 +841,24 @@ theorem IterM.findM?_pure {α β : Type w} {m : Type w → Type w'} [Monad m] · simp [ihs ‹_›] · simp +theorem IterM.first?_eq_match_step {α β : Type w} {m : Type w → Type w'} [Monad m] + [Iterator α m β] [IteratorLoop α m m] [LawfulMonad m] [Productive α m] + [LawfulIteratorLoop α m m] {it : IterM (α := α) m β} : + it.first? = (do + match (← it.step).inflate.val with + | .yield _ out => return (some out) + | .skip it' => it'.first? + | .done => return none) := by + simp only [first?] + have := IteratorLoop.wellFounded_of_productive (α := α) (β := β) (m := m) + (P := fun b g s => s = ForInStep.done (some b)) (by simp) + simp only [LawfulIteratorLoop.lawful _ _ _ _ _ this] + rw [IterM.DefaultConsumers.forIn_eq, IterM.DefaultConsumers.forIn'_eq_match_step _ this] + simp only [flip, pure_bind] + congr + ext s + split <;> try (simp [*]; done) + simp only [DefaultConsumers.forIn_eq, *] + exact IterM.DefaultConsumers.forIn'_eq_forIn' _ this (by simp) + end Std diff --git a/src/Init/Data/String/Pattern/Basic.lean b/src/Init/Data/String/Pattern/Basic.lean index 45c3a2e27e..c8d7f7d553 100644 --- a/src/Init/Data/String/Pattern/Basic.lean +++ b/src/Init/Data/String/Pattern/Basic.lean @@ -8,6 +8,7 @@ module prelude public import Init.Data.String.Basic public import Init.Data.Iterators.Basic +public import Init.Data.Iterators.Consumers.Loop set_option doc.verso true @@ -129,23 +130,24 @@ variable [∀ s, Std.Iterator (σ s) Id (SearchStep s)] variable (pat : ρ) [ToForwardSearcher pat σ] @[specialize pat] -def defaultStartsWith (s : Slice) : Bool := +def defaultStartsWith (s : Slice) [Std.IteratorLoop (σ s) Id Id] : Bool := let searcher := ToForwardSearcher.toSearcher pat s - match searcher.step with - | .yield _ (.matched start ..) _ => s.startPos = start + match searcher.first? with + | some (.matched start ..) => s.startPos = start | _ => false @[specialize pat] -def defaultDropPrefix? (s : Slice) : Option s.Pos := +def defaultDropPrefix? (s : Slice) [Std.IteratorLoop (σ s) Id Id] : Option s.Pos := let searcher := ToForwardSearcher.toSearcher pat s - match searcher.step with - | .yield _ (.matched _ endPos) _ => some endPos + match searcher.first? with + | some (.matched _ endPos) => some endPos | _ => none @[always_inline, inline] -def defaultImplementation {pat : ρ} [ToForwardSearcher pat σ] : ForwardPattern pat where - startsWith := defaultStartsWith pat - dropPrefix? := defaultDropPrefix? pat +def defaultImplementation {pat : ρ} [ToForwardSearcher pat σ] [∀ s, Std.IteratorLoop (σ s) Id Id] : + ForwardPattern pat where + startsWith s := defaultStartsWith pat s + dropPrefix? s := defaultDropPrefix? pat s end ForwardPattern @@ -188,23 +190,24 @@ variable [∀ s, Std.Iterator (σ s) Id (SearchStep s)] variable (pat : ρ) [ToBackwardSearcher pat σ] @[specialize pat] -def defaultEndsWith (s : Slice) : Bool := +def defaultEndsWith (s : Slice) [Std.IteratorLoop (σ s) Id Id] : Bool := let searcher := ToBackwardSearcher.toSearcher pat s - match searcher.step with - | .yield _ (.matched _ endPos) _ => s.endPos = endPos + match searcher.first? with + | some (.matched _ endPos) => s.endPos = endPos | _ => false @[specialize pat] -def defaultDropSuffix? (s : Slice) : Option s.Pos := +def defaultDropSuffix? (s : Slice) [Std.IteratorLoop (σ s) Id Id] : Option s.Pos := let searcher := ToBackwardSearcher.toSearcher pat s - match searcher.step with - | .yield _ (.matched startPos _) _ => some startPos + match searcher.first? with + | some (.matched startPos _) => some startPos | _ => none @[always_inline, inline] -def defaultImplementation {pat : ρ} [ToBackwardSearcher pat σ] : BackwardPattern pat where - endsWith := defaultEndsWith pat - dropSuffix? := defaultDropSuffix? pat +def defaultImplementation {pat : ρ} [ToBackwardSearcher pat σ] [∀ s, Std.IteratorLoop (σ s) Id Id] : + BackwardPattern pat where + endsWith s := defaultEndsWith pat s + dropSuffix? s := defaultDropSuffix? pat s end ToBackwardSearcher