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

@@ -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.
 

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
 
 /--

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

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

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

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