From c3976ae1321d7dd27857abc9005009ea756a966b Mon Sep 17 00:00:00 2001
From: Arsen Adzhiametov <kalpazang@gmail.com>
Date: Sat, 5 Apr 2025 03:13:35 +0300
Subject: [PATCH] Add partition function for map

---
 .../stdlib/common/src/generated/_Maps.kt      | 17 +++++++++++++
 .../samples/test/samples/collections/maps.kt  | 25 +++++++++++++++++++
 libraries/stdlib/test/collections/MapTest.kt  | 20 +++++++++++++++
 .../kotlin-stdlib-runtime-merged.txt          |  1 +
 4 files changed, 63 insertions(+)

diff --git a/libraries/stdlib/common/src/generated/_Maps.kt b/libraries/stdlib/common/src/generated/_Maps.kt
index 40dcce2209846..344679b8aeea9 100644
--- a/libraries/stdlib/common/src/generated/_Maps.kt
+++ b/libraries/stdlib/common/src/generated/_Maps.kt
@@ -668,3 +668,20 @@ public fun <K, V> Map<out K, V>.asSequence(): Sequence<Map.Entry<K, V>> {
     return entries.asSequence()
 }
 
+/**
+ * Splits the original map into a pair of maps,
+ * where *first* map contains elements for which [predicate] yielded `true`,
+ * while *second* map contains elements for which [predicate] yielded `false`.
+ *
+ * @sample samples.collections.Maps.Operations.partition
+ */
+public inline fun <K, V> Map<K, V>.partition(predicate: (Map.Entry<K, V>) -> Boolean): Pair<Map<K, V>, Map<K, V>> {
+    val first = mutableMapOf<K, V>()
+    val second = mutableMapOf<K, V>()
+    forEach {
+        if (predicate(it)) first[it.key] = it.value
+        else second[it.key] = it.value
+    }
+    return Pair(first, second)
+}
+
diff --git a/libraries/stdlib/samples/test/samples/collections/maps.kt b/libraries/stdlib/samples/test/samples/collections/maps.kt
index d9a78e68e11e3..2d3a518354bec 100644
--- a/libraries/stdlib/samples/test/samples/collections/maps.kt
+++ b/libraries/stdlib/samples/test/samples/collections/maps.kt
@@ -552,5 +552,30 @@ class Maps {
             assertTrue(map.isEmpty())
         }
     }
+
+    class Operations {
+        private enum class Status {
+            NEW, IN_PROGRESS, COMPLETED
+        }
+
+        @Sample
+        fun partition() {
+            data class Mission(val id: String) {
+                override fun toString() = id
+            }
+
+            data class Activity(val status: Status) {
+                override fun toString() = "$status"
+            }
+
+            val map = mapOf(
+                Mission("1") to listOf(Activity(Status.NEW), Activity(Status.IN_PROGRESS)),
+                Mission("2") to listOf(Activity(Status.IN_PROGRESS), Activity(Status.COMPLETED)),
+                Mission("3") to listOf(Activity(Status.COMPLETED), Activity(Status.COMPLETED)),
+            )
+            val result = map.partition { it.value.all { activity -> activity.status == Status.COMPLETED } }
+            assertPrints(result, "({3=[COMPLETED, COMPLETED]}, {1=[NEW, IN_PROGRESS], 2=[IN_PROGRESS, COMPLETED]})")
+        }
+    }
 }
 
diff --git a/libraries/stdlib/test/collections/MapTest.kt b/libraries/stdlib/test/collections/MapTest.kt
index a99f51983c96a..744665d513c56 100644
--- a/libraries/stdlib/test/collections/MapTest.kt
+++ b/libraries/stdlib/test/collections/MapTest.kt
@@ -978,4 +978,24 @@ class MapTest {
         assertEquals(kclasses[A::class], 28)
         assertEquals(kclasses[B::class], 29)
     }
+
+    private enum class Status {
+        NEW, IN_PROGRESS, COMPLETED
+    }
+
+    @Test
+    fun partition() {
+        data class Mission(val id: String)
+        data class Activity(val status: Status)
+
+        val map = mapOf(
+            Mission("M1") to listOf(Activity(Status.NEW), Activity(Status.IN_PROGRESS)),
+            Mission("M2") to listOf(Activity(Status.IN_PROGRESS), Activity(Status.COMPLETED)),
+            Mission("M3") to listOf(Activity(Status.COMPLETED), Activity(Status.COMPLETED)),
+        )
+        val result = map.partition { it.value.all { activity -> activity.status == Status.COMPLETED } }
+
+        assertEquals(1, result.first.size)
+        assertEquals(2, result.second.size)
+    }
 }
diff --git a/libraries/tools/binary-compatibility-validator/reference-public-api/kotlin-stdlib-runtime-merged.txt b/libraries/tools/binary-compatibility-validator/reference-public-api/kotlin-stdlib-runtime-merged.txt
index 0620851e3fe3a..95cb1d8ed15a0 100644
--- a/libraries/tools/binary-compatibility-validator/reference-public-api/kotlin-stdlib-runtime-merged.txt
+++ b/libraries/tools/binary-compatibility-validator/reference-public-api/kotlin-stdlib-runtime-merged.txt
@@ -2582,6 +2582,7 @@ public final class kotlin/collections/MapsKt {
 	public static final fun none (Ljava/util/Map;Lkotlin/jvm/functions/Function1;)Z
 	public static final fun onEach (Ljava/util/Map;Lkotlin/jvm/functions/Function1;)Ljava/util/Map;
 	public static final fun onEachIndexed (Ljava/util/Map;Lkotlin/jvm/functions/Function2;)Ljava/util/Map;
+	public static final fun partition (Ljava/util/Map;Lkotlin/jvm/functions/Function1;)Lkotlin/Pair;
 	public static final fun plus (Ljava/util/Map;Ljava/lang/Iterable;)Ljava/util/Map;
 	public static final fun plus (Ljava/util/Map;Ljava/util/Map;)Ljava/util/Map;
 	public static final fun plus (Ljava/util/Map;Lkotlin/Pair;)Ljava/util/Map;