From 1329a264c814efdbe81747f0d44b9b355a198262 Mon Sep 17 00:00:00 2001
From: Kim Morrison <scott.morrison@gmail.com>
Date: Wed, 2 Oct 2024 14:23:27 +1000
Subject: [PATCH] feat: HashSet.all/any (#5582)

I think the overhead (runtime/later proving) of using `for` is paid off
by being able to short-circuit.

These functions are needed downstream to switch over the Std.HashSet.
---
 src/Std/Data/HashSet/Basic.lean | 12 ++++++++++++
 tests/lean/run/hashmap.lean     | 16 ++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/src/Std/Data/HashSet/Basic.lean b/src/Std/Data/HashSet/Basic.lean
index 081dd00740ad..dedad907249b 100644
--- a/src/Std/Data/HashSet/Basic.lean
+++ b/src/Std/Data/HashSet/Basic.lean
@@ -195,6 +195,18 @@ instance [BEq α] [Hashable α] {m : Type v → Type v} : ForM m (HashSet α) α
 instance [BEq α] [Hashable α] {m : Type v → Type v} : ForIn m (HashSet α) α where
   forIn m init f := m.forIn f init
 
+/-- Check if all elements satisfy the predicate, short-circuiting if a predicate fails. -/
+@[inline] def all (m : HashSet α) (p : α → Bool) : Bool := Id.run do
+  for a in m do
+    if ¬ p a then return false
+  return true
+
+/-- Check if any element satisfies the predicate, short-circuiting if a predicate succeeds. -/
+@[inline] def any (m : HashSet α) (p : α → Bool) : Bool := Id.run do
+  for a in m do
+    if p a then return true
+  return false
+
 /-- Transforms the hash set into a list of elements in some order. -/
 @[inline] def toList (m : HashSet α) : List α :=
   m.inner.keys
diff --git a/tests/lean/run/hashmap.lean b/tests/lean/run/hashmap.lean
index d9a3ace7ad94..774c9692068a 100644
--- a/tests/lean/run/hashmap.lean
+++ b/tests/lean/run/hashmap.lean
@@ -428,6 +428,22 @@ def addKeyToState (k : Nat) : StateM Nat PUnit := do
     ans := k :: ans
   return ans
 
+/-- info: true -/
+#guard_msgs in
+#eval m.any fun x => x > 4
+
+/-- info: false -/
+#guard_msgs in
+#eval m.any fun x => x = 0
+
+/-- info: true -/
+#guard_msgs in
+#eval m.all fun x => x < 2^30
+
+/-- info: false -/
+#guard_msgs in
+#eval m.all fun x => x > 4
+
 /-- info: [1000000000, 2, 1] -/
 #guard_msgs in
 #eval m.toList