Refactor - #2812 sequence separate performance (#2818)

* Refactor: Use fold to separate Sequence of Either * Refactor: Use fold to separate Sequence of Validated * Feature: Add tests for separating Either & Validated Co-authored-by: Simon Vergauwen <[email protected]> Co-authored-by: Imran Malic Settuba <[email protected]>
arrow-kt · Oct 17, 2022 · 336c501 · 336c501
1 parent 8340c83
commit 336c501
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 10 deletions.
diff --git a/arrow-libs/core/arrow-core/src/commonMain/kotlin/arrow/core/Sequence.kt b/arrow-libs/core/arrow-core/src/commonMain/kotlin/arrow/core/Sequence.kt
@@ -593,23 +593,27 @@ public fun <A> Sequence<A>.salign(
  * @receiver Iterable of Validated
  * @return a tuple containing Sequence with [Either.Left] and another Sequence with its [Either.Right] values.
  */
-public fun <A, B> Sequence<Either<A, B>>.separateEither(): Pair<Sequence<A>, Sequence<B>> {
-  val asep = flatMap { gab -> gab.fold({ sequenceOf(it) }, { emptySequence() }) }
-  val bsep = flatMap { gab -> gab.fold({ emptySequence() }, { sequenceOf(it) }) }
-  return asep to bsep
-}
+public fun <A, B> Sequence<Either<A, B>>.separateEither(): Pair<Sequence<A>, Sequence<B>> =
+  fold(sequenceOf<A>() to sequenceOf<B>()) { (lefts, rights), either ->
+    when (either) {
+      is Left -> lefts + either.value to rights
+      is Right -> lefts to rights + either.value
+    }
+  }
 
 /**
  * Separate the inner [Validated] values into the [Validated.Invalid] and [Validated.Valid].
  *
  * @receiver Iterable of Validated
  * @return a tuple containing Sequence with [Validated.Invalid] and another Sequence with its [Validated.Valid] values.
  */
-public fun <A, B> Sequence<Validated<A, B>>.separateValidated(): Pair<Sequence<A>, Sequence<B>> {
-  val asep = flatMap { gab -> gab.fold({ sequenceOf(it) }, { emptySequence() }) }
-  val bsep = flatMap { gab -> gab.fold({ emptySequence() }, { sequenceOf(it) }) }
-  return asep to bsep
-}
+public fun <A, B> Sequence<Validated<A, B>>.separateValidated(): Pair<Sequence<A>, Sequence<B>> =
+  fold(sequenceOf<A>() to sequenceOf<B>()) { (invalids, valids), validated ->
+    when (validated) {
+      is Valid -> invalids to valids + validated.value
+      is Invalid -> invalids + validated.value to valids
+    }
+  }
 
 public fun <E, A> Sequence<Either<E, A>>.sequence(): Either<E, List<A>> =
   traverse(::identity)

diff --git a/arrow-libs/core/arrow-core/src/commonTest/kotlin/arrow/core/SequenceKTest.kt b/arrow-libs/core/arrow-core/src/commonTest/kotlin/arrow/core/SequenceKTest.kt
@@ -307,5 +307,31 @@ class SequenceKTest : UnitSpec() {
         ints.asSequence().filterOption().toList() shouldBe ints.filterOption()
       }
     }
+
+    "separateEither" {
+      checkAll(Arb.sequence(Arb.int())) { ints ->
+        val sequence = ints.map {
+          if (it % 2 == 0) it.left()
+          else it.right()
+        }
+
+        val (lefts, rights) = sequence.separateEither()
+
+        lefts.toList() to rights.toList() shouldBe ints.partition { it % 2 == 0 }
+      }
+    }
+
+    "separateValidated" {
+      checkAll(Arb.sequence(Arb.int())) { ints ->
+        val sequence = ints.map {
+          if (it % 2 == 0) it.invalid()
+          else it.valid()
+        }
+
+        val (invalids, valids) = sequence.separateValidated()
+
+        invalids.toList() to valids.toList() shouldBe ints.partition { it % 2 == 0 }
+      }
+    }
   }
 }