From b41019e8e8c967d5aa008c245d0f3733cef42efa Mon Sep 17 00:00:00 2001
From: Kim Morrison <scott.morrison@gmail.com>
Date: Fri, 20 Sep 2024 16:05:10 +1000
Subject: [PATCH] feat: HashSet.partition (unverified) (#5370)

`Aesop` is depending on (a custom-rolled) `Set` data structure with
`.partition`.
---
 src/Std/Data/DHashMap/Basic.lean | 9 +++++++++
 src/Std/Data/HashMap/Basic.lean  | 5 +++++
 src/Std/Data/HashSet/Basic.lean  | 5 +++++
 3 files changed, 19 insertions(+)

diff --git a/src/Std/Data/DHashMap/Basic.lean b/src/Std/Data/DHashMap/Basic.lean
index 472c73f703..42daaa7743 100644
--- a/src/Std/Data/DHashMap/Basic.lean
+++ b/src/Std/Data/DHashMap/Basic.lean
@@ -186,6 +186,15 @@ section Unverified
     (init : δ) (b : DHashMap α β) : δ :=
   b.1.fold f init
 
+/-- Partition a hashset into two hashsets based on a predicate. -/
+@[inline] def partition (f : (a : α) → β a → Bool)
+    (m : DHashMap α β) : DHashMap α β × DHashMap α β :=
+  m.fold (init := (∅, ∅)) fun ⟨l, r⟩  a b =>
+    if f a b then
+      (l.insert a b, r)
+    else
+      (l, r.insert a b)
+
 @[inline, inherit_doc Raw.forM] def forM (f : (a : α) → β a → m PUnit)
     (b : DHashMap α β) : m PUnit :=
   b.1.forM f
diff --git a/src/Std/Data/HashMap/Basic.lean b/src/Std/Data/HashMap/Basic.lean
index 973889e797..f68cc9520b 100644
--- a/src/Std/Data/HashMap/Basic.lean
+++ b/src/Std/Data/HashMap/Basic.lean
@@ -190,6 +190,11 @@ section Unverified
     (m : HashMap α β) : HashMap α β :=
   ⟨m.inner.filter f⟩
 
+@[inline, inherit_doc DHashMap.partition] def partition (f : α → β → Bool)
+    (m : HashMap α β) : HashMap α β × HashMap α β :=
+  let ⟨l, r⟩ := m.inner.partition f
+  ⟨⟨l⟩, ⟨r⟩⟩
+
 @[inline, inherit_doc DHashMap.foldM] def foldM {m : Type w → Type w}
     [Monad m] {γ : Type w} (f : γ → α → β → m γ) (init : γ) (b : HashMap α β) : m γ :=
   b.inner.foldM f init
diff --git a/src/Std/Data/HashSet/Basic.lean b/src/Std/Data/HashSet/Basic.lean
index a065598981..4c2ee8fcf6 100644
--- a/src/Std/Data/HashSet/Basic.lean
+++ b/src/Std/Data/HashSet/Basic.lean
@@ -158,6 +158,11 @@ section Unverified
 @[inline] def filter (f : α → Bool) (m : HashSet α) : HashSet α :=
   ⟨m.inner.filter fun a _ => f a⟩
 
+/-- Partition a hashset into two hashsets based on a predicate. -/
+@[inline] def partition (f : α → Bool) (m : HashSet α) : HashSet α × HashSet α :=
+  let ⟨l, r⟩ := m.inner.partition fun a _ => f a
+  ⟨⟨l⟩, ⟨r⟩⟩
+
 /--
 Monadically computes a value by folding the given function over the elements in the hash set in some
 order.