From 336c501b74d3fbf4ab72db22b12f20a7c0415a25 Mon Sep 17 00:00:00 2001
From: Giorgos Makris <makris.giwrgos1234@gmail.com>
Date: Mon, 17 Oct 2022 08:52:28 +0300
Subject: [PATCH] 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 <nomisRev@users.noreply.github.com>
Co-authored-by: Imran Malic Settuba <46971368+i-walker@users.noreply.github.com>
---
 .../commonMain/kotlin/arrow/core/Sequence.kt  | 24 ++++++++++-------
 .../kotlin/arrow/core/SequenceKTest.kt        | 26 +++++++++++++++++++
 2 files changed, 40 insertions(+), 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
index 8e00e19fc78..7f67780684d 100644
--- 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,11 +593,13 @@ 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].
@@ -605,11 +607,13 @@ public fun <A, B> Sequence<Either<A, B>>.separateEither(): Pair<Sequence<A>, Seq
  * @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
index 1ad6bffe6f2..8e77cffad6f 100644
--- 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 }
+      }
+    }
   }
 }