From 5cbeb22564c6caa0603612edfdabda4bdb34b780 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Mon, 24 Feb 2025 12:37:45 -0800
Subject: [PATCH] feat: add `ForIn` instance for `PHashSet` (#7214)

This PR adds a `ForIn` instance for the `PersistentHashSet` type.
---
 src/Lean/Data/PersistentHashSet.lean |  7 ++++++
 tests/lean/run/forIn_phashset.lean   | 35 ++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)
 create mode 100644 tests/lean/run/forIn_phashset.lean

diff --git a/src/Lean/Data/PersistentHashSet.lean b/src/Lean/Data/PersistentHashSet.lean
index 206e3fce36..1ead5c97f0 100644
--- a/src/Lean/Data/PersistentHashSet.lean
+++ b/src/Lean/Data/PersistentHashSet.lean
@@ -53,4 +53,11 @@ variable {_ : BEq α} {_ : Hashable α}
 def toList (s : PersistentHashSet α) : List α :=
   s.set.toList.map (·.1)
 
+protected def forIn {_ : BEq α} {_ : Hashable α} [Monad m]
+    (s : PersistentHashSet α) (init : σ) (f : α → σ → m (ForInStep σ)) : m σ := do
+  PersistentHashMap.forIn s.set init fun p s => f p.1 s
+
+instance {_ : BEq α} {_ : Hashable α} : ForIn m (PersistentHashSet α) α where
+  forIn := PersistentHashSet.forIn
+
 end PersistentHashSet
diff --git a/tests/lean/run/forIn_phashset.lean b/tests/lean/run/forIn_phashset.lean
new file mode 100644
index 0000000000..d2f27a51f3
--- /dev/null
+++ b/tests/lean/run/forIn_phashset.lean
@@ -0,0 +1,35 @@
+import Lean.Data.PersistentHashSet
+
+open Lean
+
+def sum (s : PHashSet Nat) : Nat := Id.run do
+  let mut r := 0
+  for a in s do
+    r := r + a
+  return r
+
+def sumIf (s : PHashSet Nat) (p : Nat → Bool) : Nat := Id.run do
+  let mut r := 0
+  for a in s do
+    unless p a do
+      continue
+    r := r + a
+  return r
+
+def mk [Hashable α] [BEq α] (f : Nat → α) (n : Nat) : PHashSet α := Id.run do
+  let mut s := {}
+  for i in [:n] do
+    s := s.insert (f i)
+  return s
+
+/-- info: 45 -/
+#guard_msgs in
+#eval sum (mk id 10)
+
+/-- info: 9900 -/
+#guard_msgs in
+#eval sum (mk (2*·) 100)
+
+/-- info: 2450 -/
+#guard_msgs in
+#eval sumIf (mk id 100) (· % 2 == 0)