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

This PR adds a `find?` consumer in analogy to `List.find?` and variants
thereof.
This commit is contained in:
Paul Reichert 2025-10-20 11:12:53 +02:00 committed by GitHub
parent 0d5869bb71
commit 71f1a6c164
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 414 additions and 0 deletions

View file

@ -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 :=

View file

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

View file

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

View file

@ -514,4 +514,126 @@ theorem IterM.all_eq_not_any_not {α β : Type w} {m : Type w → Type w'} [Iter
· simp [ihs _]
· simp
theorem IterM.findSomeM?_eq_match_step {α β γ : Type w} {m : Type w → Type w'} [Monad m]
[Iterator α m β] [IteratorLoop α m m] [LawfulMonad m] [Finite α m] [LawfulIteratorLoop α m m]
{it : IterM (α := α) m β} {f : β → m (Option γ)} :
it.findSomeM? f = (do
match (← it.step).inflate.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]
apply bind_congr; intro step
cases step.inflate using PlausibleIterStep.casesOn
· simp only [bind_assoc]
apply bind_congr; intro fx
split <;> simp [findSomeM?]
· simp [findSomeM?]
· simp
theorem IterM.findSome?_eq_findSomeM? {α β γ : Type w} {m : Type w → Type w'} [Monad m]
[Iterator α m β] [IteratorLoop α m m] [Finite α m]
{it : IterM (α := α) m β} {f : β → Option γ} :
it.findSome? f = it.findSomeM? (pure <| f ·) :=
(rfl)
theorem IterM.findSome?_eq_match_step {α β γ : Type w} {m : Type w → Type w'} [Monad m]
[Iterator α m β] [IteratorLoop α m m] [LawfulMonad m] [Finite α m] [LawfulIteratorLoop α m m]
{it : IterM (α := α) m β} {f : β → Option γ} :
it.findSome? f = (do
match (← it.step).inflate.val with
| .yield it' out =>
match f out with
| none => it'.findSome? f
| some fx => return (some fx)
| .skip it' => it'.findSome? f
| .done => return none) := by
rw [findSome?_eq_findSomeM?, findSomeM?_eq_match_step]
apply bind_congr; intro step
split <;> simp [findSome?_eq_findSomeM?]
theorem IterM.findSomeM?_pure {α β γ : Type w} {m : Type w → Type w'} [Monad m]
[Iterator α m β] [IteratorLoop α m m]
[LawfulMonad m] [Finite α m] [LawfulIteratorLoop α m m]
{it : IterM (α := α) m β} {f : β → Option γ} :
it.findSomeM? (pure <| f ·) = it.findSome? f := by
induction it using IterM.inductSteps with | step it ihy ihs
rw [findSomeM?_eq_match_step, findSome?_eq_match_step]
apply bind_congr; intro step
cases step.inflate using PlausibleIterStep.casesOn
· simp only [pure_bind]
split <;> simp [ihy _]
· simp [ihs _]
· simp
theorem IterM.findM?_eq_findSomeM? {α β : Type w} {m : Type w → Type w'} [Monad m]
[Iterator α m β] [IteratorLoop α m m] [Finite α m]
{it : IterM (α := α) m β} {f : β → m (ULift Bool)} :
it.findM? f = it.findSomeM? (fun x => return if (← f x).down then some x else none) :=
(rfl)
theorem IterM.findM?_eq_match_step {α β : Type w} {m : Type w → Type w'} [Monad m]
[Iterator α m β] [IteratorLoop α m m] [LawfulMonad m] [Finite α m] [LawfulIteratorLoop α m m]
{it : IterM (α := α) m β} {f : β → m (ULift Bool)} :
it.findM? f = (do
match (← it.step).inflate.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]
apply bind_congr; intro step
split
· simp only [bind_assoc]
apply bind_congr; intro fx
split <;> simp [findM?_eq_findSomeM?]
· simp [findM?_eq_findSomeM?]
· simp
theorem IterM.find?_eq_findM? {α β : Type w} {m : Type w → Type w'} [Monad m] [Iterator α m β]
[IteratorLoop α m m] [Finite α m] {it : IterM (α := α) m β} {f : β → Bool} :
it.find? f = it.findM? (pure <| .up <| f ·) :=
(rfl)
theorem IterM.find?_eq_findSome? {α β : Type w} {m : Type w → Type w'} [Monad m] [Iterator α m β]
[IteratorLoop α m m] [LawfulMonad m] [Finite α m] {it : IterM (α := α) m β} {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 IterM.find?_eq_match_step {α β : Type w} {m : Type w → Type w'} [Monad m]
[Iterator α m β] [IteratorLoop α m m] [LawfulMonad m] [Finite α m] [LawfulIteratorLoop α m m]
{it : IterM (α := α) m β} {f : β → Bool} :
it.find? f = (do
match (← it.step).inflate.val with
| .yield it' out =>
if f out then return (some out) else it'.find? f
| .skip it' => it'.find? f
| .done => return none) := by
rw [find?_eq_findM?, findM?_eq_match_step]
apply bind_congr; intro step
split
· simp only [pure_bind]
split <;> simp [find?_eq_findM?]
· simp [find?_eq_findM?]
· simp
theorem IterM.findM?_pure {α β : Type w} {m : Type w → Type w'} [Monad m]
[Iterator α m β] [IteratorLoop α m m]
[LawfulMonad m] [Finite α m] [LawfulIteratorLoop α m m]
{it : IterM (α := α) m β} {f : β → ULift Bool} :
it.findM? (pure (f := m) <| f ·) = it.find? (ULift.down <| f ·) := by
induction it using IterM.inductSteps with | step it ihy ihs
rw [findM?_eq_match_step, find?_eq_match_step]
apply bind_congr; intro step
cases step.inflate using PlausibleIterStep.casesOn
· simp only [pure_bind]
split
· simp
· simp [ihy _]
· simp [ihs _]
· simp
end Std.Iterators

View file

@ -403,6 +403,21 @@ def findSomeM? {m : Type u → Type v} [Monad m] {α : Type w} {β : Type u} (f
| some b => pure (some b)
| none => findSomeM? f as
@[simp, grind =]
theorem findSomeM?_nil [Monad m] {α : Type w} {β : Type u}
{f : α → m (Option β)} :
([] : List α).findSomeM? f = pure none :=
(rfl)
@[grind =]
theorem findSomeM?_cons [Monad m] {α : Type w} {β : Type u}
{f : α → m (Option β)} {a : α} {as : List α} :
(a::as).findSomeM? f = (do
match ← f a with
| some b => return some b
| none => as.findSomeM? f) :=
(rfl)
@[simp]
theorem findSomeM?_pure [Monad m] [LawfulMonad m] {f : α → Option β} {as : List α} :
findSomeM? (m := m) (pure <| f ·) as = pure (as.findSome? f) := by
@ -424,6 +439,10 @@ theorem findSomeM?_id (f : α → Id (Option β)) (as : List α) :
findSomeM? (m := Id) f as = as.findSome? f :=
findSomeM?_pure
theorem findSome?_eq_findSomeM? {f : α → Option β} {as : List α} :
as.findSome? f = (as.findSomeM? (pure (f := Id) <| f ·)).run := by
simp
theorem findM?_eq_findSomeM? [Monad m] [LawfulMonad m] {p : α → m Bool} {as : List α} :
as.findM? p = as.findSomeM? fun a => return if (← p a) then some a else none := by
induction as with