Add partition function for map

Author: arsen-adzhiametov

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)
+}
+

libraries/stdlib/samples/test/samples/collections/maps.kt

@@ -412,5 +412,30 @@ class Maps {
             assertPrints(map.flatMap { (key, value) -> key.take(value).toList() }, "[1, 2, 3, 4, 5]")
         }
     }
+
+    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]})")
+        }
+    }
 }
 

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)
+    }
 }

libraries/tools/binary-compatibility-validator/reference-public-api/kotlin-stdlib-runtime-merged.txt

@@ -2579,6 +2579,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;