feat: Iterator find? consumer and variants (#10769)

This PR adds a `find?` consumer in analogy to `List.find?` and variants
thereof.

Author: datokrat

src/Init/Data/Iterators/Consumers/Loop.lean

@@ -202,6 +202,59 @@ def Iter.all {α β : Type w}
     (p : β → Bool) (it : Iter (α := α) β) : Bool :=
   (it.allM (fun x => pure (f := Id) (p x))).run
 
+@[inline]
+def Iter.findSomeM? {α β : Type w} {γ : Type x} {m : Type x → Type w'} [Monad m] [Iterator α Id β]
+    [IteratorLoop α Id m] [Finite α Id] (it : Iter (α := α) β) (f : β → m (Option γ)) :
+    m (Option γ) :=
+  ForIn.forIn it none (fun x _ => do
+    match ← f x with
+    | none => return .yield none
+    | some fx => return .done (some fx))
+
+@[inline]
+def Iter.Partial.findSomeM? {α β : Type w} {γ : Type x} {m : Type x → Type w'} [Monad m]
+    [Iterator α Id β] [IteratorLoopPartial α Id m] (it : Iter.Partial (α := α) β)
+    (f : β → m (Option γ)) :
+    m (Option γ) :=
+  ForIn.forIn it none (fun x _ => do
+    match ← f x with
+    | none => return .yield none
+    | some fx => return .done (some fx))
+
+@[inline]
+def Iter.findSome? {α β : Type w} {γ : Type x} [Iterator α Id β]
+    [IteratorLoop α Id Id] [Finite α Id] (it : Iter (α := α) β) (f : β → Option γ) :
+    Option γ :=
+  Id.run (it.findSomeM? (pure <| f ·))
+
+@[inline]
+def Iter.Partial.findSome? {α β : Type w} {γ : Type x} [Iterator α Id β]
+    [IteratorLoopPartial α Id Id] (it : Iter.Partial (α := α) β) (f : β → Option γ) :
+    Option γ :=
+  Id.run (it.findSomeM? (pure <| f ·))
+
+@[inline]
+def Iter.findM? {α β : Type w} {m : Type w → Type w'} [Monad m] [Iterator α Id β]
+    [IteratorLoop α Id m] [Finite α Id] (it : Iter (α := α) β) (f : β → m (ULift Bool)) :
+    m (Option β) :=
+  it.findSomeM? (fun x => return if (← f x).down then some x else none)
+
+@[inline]
+def Iter.Partial.findM? {α β : Type w} {m : Type w → Type w'} [Monad m] [Iterator α Id β]
+    [IteratorLoopPartial α Id m] (it : Iter.Partial (α := α) β) (f : β → m (ULift Bool)) :
+    m (Option β) :=
+  it.findSomeM? (fun x => return if (← f x).down then some x else none)
+
+@[inline]
+def Iter.find? {α β : Type w} [Iterator α Id β] [IteratorLoop α Id Id]
+    [Finite α Id] (it : Iter (α := α) β) (f : β → Bool) : Option β :=
+  Id.run (it.findM? (pure <| .up <| f ·))
+
+@[inline]
+def Iter.Partial.find? {α β : Type w} [Iterator α Id β] [IteratorLoopPartial α Id Id]
+    (it : Iter.Partial (α := α) β) (f : β → Bool) : Option β :=
+  Id.run (it.findM? (pure <| .up <| f ·))
+
 @[always_inline, inline, expose, inherit_doc IterM.size]
 def Iter.size {α : Type w} {β : Type w} [Iterator α Id β] [IteratorSize α Id]
     (it : Iter (α := α) β) : Nat :=

src/Init/Data/Iterators/Consumers/Monadic/Loop.lean

@@ -579,6 +579,60 @@ def IterM.Partial.all {α β : Type w} {m : Type w → Type w'} [Monad m]
     (p : β → Bool) (it : IterM.Partial (α := α) m β) : m (ULift Bool) := do
   it.allM (fun x => pure (.up (p x)))
 
+@[inline]
+def IterM.findSomeM? {α β γ : Type w} {m : Type w → Type w'} [Monad m] [Iterator α m β]
+    [IteratorLoop α m m] [Finite α m] (it : IterM (α := α) m β) (f : β → m (Option γ)) :
+    m (Option γ) :=
+  ForIn.forIn it none (fun x _ => do
+    match ← f x with
+    | none => return .yield none
+    | some fx => return .done (some fx))
+
+@[inline]
+def IterM.Partial.findSomeM? {α β γ : Type w} {m : Type w → Type w'} [Monad m] [Iterator α m β]
+    [IteratorLoopPartial α m m] (it : IterM.Partial (α := α) m β) (f : β → m (Option γ)) :
+    m (Option γ) :=
+  ForIn.forIn it none (fun x _ => do
+    match ← f x with
+    | none => return .yield none
+    | some fx => return .done (some fx))
+
+@[inline]
+def IterM.findSome? {α β γ : Type w} {m : Type w → Type w'} [Monad m] [Iterator α m β]
+    [IteratorLoop α m m] [Finite α m] (it : IterM (α := α) m β) (f : β → Option γ) :
+    m (Option γ) :=
+  it.findSomeM? (pure <| f ·)
+
+@[inline]
+def IterM.Partial.findSome? {α β γ : Type w} {m : Type w → Type w'} [Monad m] [Iterator α m β]
+    [IteratorLoopPartial α m m] (it : IterM.Partial (α := α) m β) (f : β → Option γ) :
+    m (Option γ) :=
+  it.findSomeM? (pure <| f ·)
+
+@[inline]
+def IterM.findM? {α β : Type w} {m : Type w → Type w'} [Monad m] [Iterator α m β]
+    [IteratorLoop α m m] [Finite α m] (it : IterM (α := α) m β) (f : β → m (ULift Bool)) :
+    m (Option β) :=
+  it.findSomeM? (fun x => return if (← f x).down then some x else none)
+
+@[inline]
+def IterM.Partial.findM? {α β : Type w} {m : Type w → Type w'} [Monad m] [Iterator α m β]
+    [IteratorLoopPartial α m m] (it : IterM.Partial (α := α) m β) (f : β → m (ULift Bool)) :
+    m (Option β) :=
+  it.findSomeM? (fun x => return if (← f x).down then some x else none)
+
+@[inline]
+def IterM.find? {α β : Type w} {m : Type w → Type w'} [Monad m] [Iterator α m β]
+    [IteratorLoop α m m] [Finite α m] (it : IterM (α := α) m β) (f : β → Bool) :
+    m (Option β) :=
+  it.findM? (pure <| .up <| f ·)
+
+@[inline]
+def IterM.Partial.find? {α β : Type w} {m : Type w → Type w'} [Monad m] [Iterator α m β]
+    [IteratorLoopPartial α m m] (it : IterM.Partial (α := α) m β) (f : β → Bool) :
+    m (Option β) :=
+  it.findM? (pure <| .up <| f ·)
+
 section Size
 
 /--

src/Init/Data/Iterators/Lemmas/Consumers/Loop.lean

@@ -726,4 +726,170 @@ theorem Iter.all_eq_not_any_not {α β : Type w} [Iterator α Id β]
   · simp [ihs ‹_›]
   · simp
 
+theorem Iter.findSomeM?_eq_match_step {α β : Type w} {γ : Type x} {m : Type x → Type w'} [Monad m]
+    [Iterator α Id β] [IteratorLoop α Id m] [LawfulMonad m] [Finite α Id] [LawfulIteratorLoop α Id m]
+    {it : Iter (α := α) β} {f : β → m (Option γ)} :
+    it.findSomeM? f = (do
+      match it.step.val with
+      | .yield it' out =>
+        match ← f out with
+        | none => it'.findSomeM? f
+        | some fx => return (some fx)
+      | .skip it' => it'.findSomeM? f
+      | .done => return none) := by
+  rw [findSomeM?, forIn_eq_match_step]
+  cases it.step using PlausibleIterStep.casesOn
+  · simp only [bind_assoc]
+    apply bind_congr; intro fx
+    split <;> simp [findSomeM?]
+  · simp [findSomeM?]
+  · simp
+
+theorem Iter.findSomeM?_toList {α β : Type w} {γ : Type x} {m : Type x → Type w'} [Monad m]
+    [Iterator α Id β] [IteratorLoop α Id m] [IteratorCollect α Id Id]
+    [LawfulMonad m] [Finite α Id] [LawfulIteratorLoop α Id m] [LawfulIteratorCollect α Id Id]
+    {it : Iter (α := α) β} {f : β → m (Option γ)} :
+    it.toList.findSomeM? f = it.findSomeM? f := by
+  induction it using Iter.inductSteps with | step it ihy ihs
+  rw [it.findSomeM?_eq_match_step, it.toList_eq_match_step]
+  cases it.step using PlausibleIterStep.casesOn
+  · simp only [List.findSomeM?_cons]
+    apply bind_congr; intro fx
+    split <;> simp [ihy ‹_›]
+  · simp [ihs ‹_›]
+  · simp
+
+theorem Iter.findSome?_eq_findSomeM? {α β : Type w} {γ : Type x}
+    [Iterator α Id β] [IteratorLoop α Id Id] [Finite α Id]
+    {it : Iter (α := α) β} {f : β → Option γ} :
+    it.findSome? f = Id.run (it.findSomeM? (pure <| f ·)) :=
+  (rfl)
+
+theorem Iter.findSome?_eq_findSome?_toIterM {α β γ : Type w}
+    [Iterator α Id β] [IteratorLoop α Id Id.{w}] [Finite α Id]
+    {it : Iter (α := α) β} {f : β → Option γ} :
+    it.findSome? f = (it.toIterM.findSome? f).run :=
+  (rfl)
+
+theorem Iter.findSome?_eq_match_step {α β : Type w} {γ : Type x}
+    [Iterator α Id β] [IteratorLoop α Id Id] [Finite α Id]
+    [LawfulIteratorLoop α Id Id] {it : Iter (α := α) β} {f : β → Option γ} :
+    it.findSome? f = (match it.step.val with
+      | .yield it' out =>
+        match f out with
+        | none => it'.findSome? f
+        | some fx => some fx
+      | .skip it' => it'.findSome? f
+      | .done => none) := by
+  rw [findSome?_eq_findSomeM?, findSomeM?_eq_match_step]
+  split
+  · simp only [pure_bind, findSome?_eq_findSomeM?]
+    split <;> simp
+  · simp [findSome?_eq_findSomeM?]
+  · simp
+
+theorem Iter.findSome?_toList {α β : Type w} {γ : Type x}
+    [Iterator α Id β] [IteratorLoop α Id Id] [IteratorCollect α Id Id]
+    [Finite α Id] [LawfulIteratorLoop α Id Id] [LawfulIteratorCollect α Id Id]
+    {it : Iter (α := α) β} {f : β → Option γ} :
+    it.toList.findSome? f = it.findSome? f := by
+  simp [findSome?_eq_findSomeM?, List.findSome?_eq_findSomeM?, findSomeM?_toList]
+
+theorem Iter.findSomeM?_pure {α β : Type w} {γ : Type x} {m : Type x → Type w'} [Monad m]
+    [Iterator α Id β] [IteratorLoop α Id m] [IteratorLoop α Id Id]
+    [LawfulMonad m] [Finite α Id] [LawfulIteratorLoop α Id m] [LawfulIteratorLoop α Id Id]
+    {it : Iter (α := α) β} {f : β → Option γ} :
+    it.findSomeM? (pure <| f ·) = pure (f := m) (it.findSome? f) := by
+  letI : IteratorCollect α Id Id := .defaultImplementation
+  simp [← findSomeM?_toList, ← findSome?_toList, List.findSomeM?_pure]
+
+theorem Iter.findM?_eq_findSomeM? {α β : Type w} {m : Type w → Type w'} [Monad m]
+    [Iterator α Id β] [IteratorLoop α Id m] [Finite α Id]
+    {it : Iter (α := α) β} {f : β → m (ULift Bool)} :
+    it.findM? f = it.findSomeM? (fun x => return if (← f x).down then some x else none) :=
+  (rfl)
+
+theorem Iter.findM?_eq_match_step {α β : Type w} {m : Type w → Type w'} [Monad m]
+    [Iterator α Id β] [IteratorLoop α Id m] [LawfulMonad m] [Finite α Id] [LawfulIteratorLoop α Id m]
+    {it : Iter (α := α) β} {f : β → m (ULift Bool)} :
+    it.findM? f = (do
+      match it.step.val with
+      | .yield it' out =>
+        if (← f out).down then return (some out) else it'.findM? f
+      | .skip it' => it'.findM? f
+      | .done => return none) := by
+  rw [findM?_eq_findSomeM?, findSomeM?_eq_match_step]
+  split
+  · simp only [bind_assoc]
+    apply bind_congr; intro fx
+    split <;> simp [findM?_eq_findSomeM?]
+  · simp [findM?_eq_findSomeM?]
+  · simp
+
+theorem Iter.findM?_toList {α β : Type} {m : Type → Type w'} [Monad m]
+    [Iterator α Id β] [IteratorLoop α Id m] [IteratorCollect α Id Id]
+    [LawfulMonad m] [Finite α Id] [LawfulIteratorLoop α Id m] [LawfulIteratorCollect α Id Id]
+    {it : Iter (α := α) β} {f : β → m Bool} :
+    it.toList.findM? f = it.findM? (.up <$> f ·) := by
+  simp [findM?_eq_findSomeM?, List.findM?_eq_findSomeM?, findSomeM?_toList]
+
+theorem Iter.findM?_eq_findM?_toList {α β : Type} {m : Type → Type w'} [Monad m]
+    [Iterator α Id β] [IteratorLoop α Id m] [IteratorCollect α Id Id]
+    [LawfulMonad m] [Finite α Id] [LawfulIteratorLoop α Id m] [LawfulIteratorCollect α Id Id]
+    {it : Iter (α := α) β} {f : β → m (ULift Bool)} :
+    it.findM? f = it.toList.findM? (ULift.down <$> f ·) := by
+  simp [findM?_toList]
+
+theorem Iter.find?_eq_findM? {α β : Type w} [Iterator α Id β]
+    [IteratorLoop α Id Id] [Finite α Id] {it : Iter (α := α) β} {f : β → Bool} :
+    it.find? f = Id.run (it.findM? (pure <| .up <| f ·)) :=
+  (rfl)
+
+theorem Iter.find?_eq_find?_toIterM {α β : Type w} [Iterator α Id β]
+    [IteratorLoop α Id Id] [Finite α Id] {it : Iter (α := α) β} {f : β → Bool} :
+    it.find? f = (it.toIterM.find? f).run :=
+  (rfl)
+
+theorem Iter.find?_eq_findSome? {α β : Type w} [Iterator α Id β]
+    [IteratorLoop α Id Id] [Finite α Id] {it : Iter (α := α) β} {f : β → Bool} :
+    it.find? f = it.findSome? (fun x => if f x then some x else none) := by
+  simp [find?_eq_findM?, findSome?_eq_findSomeM?, findM?_eq_findSomeM?]
+
+theorem Iter.find?_eq_match_step {α β : Type w}
+    [Iterator α Id β] [IteratorLoop α Id Id] [Finite α Id] [LawfulIteratorLoop α Id Id]
+    {it : Iter (α := α) β} {f : β → Bool} :
+    it.find? f = (match it.step.val with
+      | .yield it' out =>
+        if f out then some out else it'.find? f
+      | .skip it' => it'.find? f
+      | .done => none) := by
+  rw [find?_eq_findM?, findM?_eq_match_step]
+  split
+  · simp only [pure_bind]
+    split <;> simp [find?_eq_findM?]
+  · simp [find?_eq_findM?]
+  · simp
+
+theorem Iter.find?_toList {α β : Type w}
+    [Iterator α Id β] [IteratorLoop α Id Id] [IteratorCollect α Id Id]
+    [Finite α Id] [LawfulIteratorLoop α Id Id] [LawfulIteratorCollect α Id Id]
+    {it : Iter (α := α) β} {f : β → Bool} :
+    it.toList.find? f = it.find? f := by
+  simp [find?_eq_findSome?, List.find?_eq_findSome?_guard, findSome?_toList, Option.guard_def]
+
+theorem Iter.findM?_pure {α β : Type w} {m : Type w → Type w'} [Monad m]
+    [Iterator α Id β] [IteratorLoop α Id m] [IteratorLoop α Id Id]
+    [LawfulMonad m] [Finite α Id] [LawfulIteratorLoop α Id m] [LawfulIteratorLoop α Id Id]
+    {it : Iter (α := α) β} {f : β → ULift Bool} :
+    it.findM? (pure (f := m) <| f ·) = pure (f := m) (it.find? (ULift.down <| f ·)) := by
+  induction it using Iter.inductSteps with | step it ihy ihs
+  rw [findM?_eq_match_step, find?_eq_match_step]
+  cases it.step using PlausibleIterStep.casesOn
+  · simp only [pure_bind]
+    split
+    · simp
+    · simp [ihy ‹_›]
+  · simp [ihs ‹_›]
+  · simp
+
 end Std.Iterators