feat: any/all predicates for iterators (#10686)

This PR introduces `any`, `anyM`, `all` and `allM` for pure and monadic
iterators. It also provides lemmas about them.

Author: datokrat

src/Init/Control/Lawful/Basic.lean

@@ -252,6 +252,7 @@ instance : LawfulMonad Id := by
 @[simp] theorem run_map (x : Id α) (f : α → β) : (f <$> x).run = f x.run := rfl
 @[simp] theorem run_bind (x : Id α) (f : α → Id β) : (x >>= f).run = (f x.run).run := rfl
 @[simp] theorem run_pure (a : α) : (pure a : Id α).run = a := rfl
+@[simp] theorem pure_run (a : Id α) : pure a.run = a := rfl
 @[simp] theorem run_seqRight (x y : Id α) : (x *> y).run = y.run := rfl
 @[simp] theorem run_seqLeft (x y : Id α) : (x <* y).run = x.run := rfl
 @[simp] theorem run_seq (f : Id (α → β)) (x : Id α) : (f <*> x).run = f.run x.run := rfl

src/Init/Data/Iterators/Combinators/Monadic/FilterMap.lean

@@ -463,7 +463,7 @@ For each value emitted by the base iterator `it`, this combinator calls `f`.
 @[inline, expose]
 def IterM.mapM {α β γ : Type w} {m : Type w → Type w'} {n : Type w → Type w''} [Iterator α m β]
     [Monad n] [MonadLiftT m n] (f : β → n γ) (it : IterM (α := α) m β) :=
-  (it.filterMapWithPostcondition (fun b => some <$> PostconditionT.lift (f b)) : IterM n γ)
+  (it.mapWithPostcondition (fun b => PostconditionT.lift (f b)) : IterM n γ)
 
 /--
 If `it` is an iterator, then `it.filterM f` is another iterator that applies a monadic

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

@@ -139,6 +139,70 @@ def Iter.Partial.fold {α : Type w} {β : Type w} {γ : Type x} [Iterator α Id
     (init : γ) (it : Iter.Partial (α := α) β) : γ :=
   ForIn.forIn (m := Id) it init (fun x acc => ForInStep.yield (f acc x))
 
+set_option doc.verso true in
+/--
+Returns {lean}`true` if the monadic predicate {name}`p` returns {lean}`true` for
+any element emitted by the iterator {name}`it`.
+
+{lit}`O(|xs|)`. Short-circuits upon encountering the first match. The elements in {name}`it` are
+examined in order of iteration.
+-/
+@[specialize]
+def Iter.anyM {α β : Type w} {m : Type → Type w'} [Monad m]
+    [Iterator α Id β] [IteratorLoop α Id m] [Finite α Id]
+    (p : β → m Bool) (it : Iter (α := α) β) : m Bool :=
+  ForIn.forIn it false (fun x _ => do
+    if ← p x then
+      return .done true
+    else
+      return .yield false)
+
+set_option doc.verso true in
+/--
+Returns {lean}`true` if the pure predicate {name}`p` returns {lean}`true` for
+any element emitted by the iterator {name}`it`.
+
+{lit}`O(|xs|)`. Short-circuits upon encountering the first match. The elements in {name}`it` are
+examined in order of iteration.
+-/
+@[inline]
+def Iter.any {α β : Type w}
+    [Iterator α Id β] [IteratorLoop α Id Id] [Finite α Id]
+    (p : β → Bool) (it : Iter (α := α) β) : Bool :=
+  (it.anyM (fun x => pure (f := Id) (p x))).run
+
+set_option doc.verso true in
+/--
+Returns {lean}`true` if the monadic predicate {name}`p` returns {lean}`true` for
+all elements emitted by the iterator {name}`it`.
+
+{lit}`O(|xs|)`. Short-circuits upon encountering the first mismatch. The elements in {name}`it` are
+examined in order of iteration.
+-/
+@[specialize]
+def Iter.allM {α β : Type w} {m : Type → Type w'} [Monad m]
+    [Iterator α Id β] [IteratorLoop α Id m] [Finite α Id]
+    (p : β → m Bool) (it : Iter (α := α) β) : m Bool :=
+  ForIn.forIn it true (fun x _ => do
+    if ← p x then
+      return .yield true
+    else
+      return .done false)
+
+set_option doc.verso true in
+/--
+Returns {lean}`true` if the pure predicate {name}`p` returns {lean}`true` for
+all elements emitted by the iterator {name}`it`.
+
+{lit}`O(|xs|)`. Short-circuits upon encountering the first mismatch. The elements in {name}`it` are
+examined in order of iteration.
+-/
+@[inline]
+def Iter.all {α β : Type w}
+    [Iterator α Id β] [IteratorLoop α Id Id] [Finite α Id]
+    (p : β → Bool) (it : Iter (α := α) β) : Bool :=
+  (it.allM (fun x => pure (f := Id) (p x))).run
+
 @[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

@@ -416,6 +416,169 @@ def IterM.Partial.drain {α : Type w} {m : Type w → Type w'} [Monad m] {β : T
     m PUnit :=
   it.fold (γ := PUnit) (fun _ _ => .unit) .unit
 
+set_option doc.verso true in
+/--
+Returns {lean}`ULift.up true` if the monadic predicate {name}`p` returns {lean}`ULift.up true` for
+any element emitted by the iterator {name}`it`.
+
+{lit}`O(|xs|)`. Short-circuits upon encountering the first match. The elements in {name}`it` are
+examined in order of iteration.
+
+This function requires a {name}`Finite` instance proving that the iterator will finish after a
+finite number of steps. If the iterator is not finite or such an instance is not available,
+consider using {lit}`it.allowNontermination.anyM` instead of {lean}`it.anyM`. However, it is not
+possible to formally verify the behavior of the partial variant.
+-/
+@[specialize]
+def IterM.anyM {α β : Type w} {m : Type w → Type w'} [Monad m]
+    [Iterator α m β] [IteratorLoop α m m] [Finite α m]
+    (p : β → m (ULift Bool)) (it : IterM (α := α) m β) : m (ULift Bool) :=
+  ForIn.forIn it (ULift.up false) (fun x _ => do
+    if (← p x).down then
+      return .done (.up true)
+    else
+      return .yield (.up false))
+
+set_option doc.verso true in
+/--
+Returns {lean}`ULift.up true` if the monadic predicate {name}`p` returns {lean}`ULift.up true` for
+any element emitted by the iterator {name}`it`.
+
+{lit}`O(|xs|)`. Short-circuits upon encountering the first match. The elements in {name}`it` are
+examined in order of iteration.
+
+This is a partial, potentially nonterminating, function. It is not possible to formally verify
+its behavior. If the iterator has a {name}`Finite` instance, consider using {name}`IterM.anyM`
+instead.
+-/
+@[specialize]
+def IterM.Partial.anyM {α β : Type w} {m : Type w → Type w'} [Monad m]
+    [Iterator α m β] [IteratorLoopPartial α m m]
+    (p : β → m (ULift Bool)) (it : IterM.Partial (α := α) m β) : m (ULift Bool) :=
+  ForIn.forIn it (ULift.up false) (fun x _ => do
+    if (← p x).down then
+      return .done (.up true)
+    else
+      return .yield (.up false))
+
+set_option doc.verso true in
+/--
+Returns {lean}`ULift.up true` if the pure predicate {name}`p` returns {lean}`true` for
+any element emitted by the iterator {name}`it`.
+
+{lit}`O(|xs|)`. Short-circuits upon encountering the first match. The elements in {name}`it` are
+examined in order of iteration.
+
+This function requires a {name}`Finite` instance proving that the iterator will finish after a
+finite number of steps. If the iterator is not finite or such an instance is not available,
+consider using {lit}`it.allowNontermination.any` instead of {lean}`it.any`. However, it is not
+possible to formally verify the behavior of the partial variant.
+-/
+@[inline]
+def IterM.any {α β : Type w} {m : Type w → Type w'} [Monad m]
+    [Iterator α m β] [IteratorLoop α m m] [Finite α m]
+    (p : β → Bool) (it : IterM (α := α) m β) : m (ULift Bool) := do
+  it.anyM (fun x => pure (.up (p x)))
+
+set_option doc.verso true in
+/--
+Returns {lean}`ULift.up true` if the pure predicate {name}`p` returns {lean}`true` for
+any element emitted by the iterator {name}`it`.
+
+{lit}`O(|xs|)`. Short-circuits upon encountering the first match. The elements in {name}`it` are
+examined in order of iteration.
+
+This is a partial, potentially nonterminating, function. It is not possible to formally verify
+its behavior. If the iterator has a {name}`Finite` instance, consider using {name}`IterM.any`
+instead.
+-/
+@[inline]
+def IterM.Partial.any {α β : Type w} {m : Type w → Type w'} [Monad m]
+    [Iterator α m β] [IteratorLoopPartial α m m]
+    (p : β → Bool) (it : IterM.Partial (α := α) m β) : m (ULift Bool) := do
+  it.anyM (fun x => pure (.up (p x)))
+
+set_option doc.verso true in
+/--
+Returns {lean}`ULift.up true` if the monadic predicate {name}`p` returns {lean}`ULift.up true` for
+all elements emitted by the iterator {name}`it`.
+
+{lit}`O(|xs|)`. Short-circuits upon encountering the first mismatch. The elements in {name}`it` are
+examined in order of iteration.
+
+This function requires a {name}`Finite` instance proving that the iterator will finish after a
+finite number of steps. If the iterator is not finite or such an instance is not available,
+consider using {lit}`it.allowNontermination.allM` instead of {lean}`it.allM`. However, it is not
+possible to formally verify the behavior of the partial variant.
+-/
+@[specialize]
+def IterM.allM {α β : Type w} {m : Type w → Type w'} [Monad m]
+    [Iterator α m β] [IteratorLoop α m m] [Finite α m]
+    (p : β → m (ULift Bool)) (it : IterM (α := α) m β) : m (ULift Bool) := do
+  ForIn.forIn it (ULift.up true) (fun x _ => do
+    if (← p x).down then
+      return .yield (.up true)
+    else
+      return .done (.up false))
+set_option doc.verso true in
+/--
+Returns {lean}`ULift.up true` if the monadic predicate {name}`p` returns {lean}`ULift.up true` for
+all elements emitted by the iterator {name}`it`.
+
+{lit}`O(|xs|)`. Short-circuits upon encountering the first mismatch. The elements in {name}`it` are
+examined in order of iteration.
+
+This is a partial, potentially nonterminating, function. It is not possible to formally verify
+its behavior. If the iterator has a {name}`Finite` instance, consider using {name}`IterM.allM`
+instead.
+-/
+@[specialize]
+def IterM.Partial.allM {α β : Type w} {m : Type w → Type w'} [Monad m]
+    [Iterator α m β] [IteratorLoopPartial α m m]
+    (p : β → m (ULift Bool)) (it : IterM.Partial (α := α) m β) : m (ULift Bool) := do
+  ForIn.forIn it (ULift.up true) (fun x _ => do
+    if (← p x).down then
+      return .yield (.up true)
+    else
+      return .done (.up false))
+
+set_option doc.verso true in
+/--
+Returns {lean}`ULift.up true` if the pure predicate {name}`p` returns {lean}`true` for
+all elements emitted by the iterator {name}`it`.
+
+{lit}`O(|xs|)`. Short-circuits upon encountering the first mismatch. The elements in {name}`it` are
+examined in order of iteration.
+
+This function requires a {name}`Finite` instance proving that the iterator will finish after a
+finite number of steps. If the iterator is not finite or such an instance is not available,
+consider using {lit}`it.allowNontermination.all` instead of {lean}`it.all`. However, it is not
+possible to formally verify the behavior of the partial variant.
+-/
+@[inline]
+def IterM.all {α β : Type w} {m : Type w → Type w'} [Monad m]
+    [Iterator α m β] [IteratorLoop α m m] [Finite α m]
+    (p : β → Bool) (it : IterM (α := α) m β) : m (ULift Bool) := do
+  it.allM (fun x => pure (.up (p x)))
+
+set_option doc.verso true in
+/--
+Returns {lean}`ULift.up true` if the pure predicate {name}`p` returns {lean}`true` for
+all elements emitted by the iterator {name}`it`.
+
+{lit}`O(|xs|)`. Short-circuits upon encountering the first mismatch. The elements in {name}`it` are
+examined in order of iteration.
+
+This is a partial, potentially nonterminating, function. It is not possible to formally verify
+its behavior. If the iterator has a {name}`Finite` instance, consider using {name}`IterM.all`
+instead.
+-/
+@[inline]
+def IterM.Partial.all {α β : Type w} {m : Type w → Type w'} [Monad m]
+    [Iterator α m β] [IteratorLoopPartial α m m]
+    (p : β → Bool) (it : IterM.Partial (α := α) m β) : m (ULift Bool) := do
+  it.allM (fun x => pure (.up (p x)))
+
 section Size
 
 /--