Add flattenOrAccumulate (#2892)

arrow-kt · Jan 19, 2023 · 20a1905 · 20a1905
1 parent ad8ea21
commit 20a1905
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 0 deletions.
diff --git a/arrow-libs/core/arrow-core/api/arrow-core.api b/arrow-libs/core/arrow-core/api/arrow-core.api
@@ -587,6 +587,8 @@ public final class arrow/core/IterableKt {
 	public static final fun firstOrNone (Ljava/lang/Iterable;Lkotlin/jvm/functions/Function1;)Larrow/core/Option;
 	public static final fun flatten (Ljava/lang/Iterable;)Ljava/util/List;
 	public static final fun flattenOption (Ljava/lang/Iterable;)Ljava/util/List;
+	public static final fun flattenOrAccumulate (Ljava/lang/Iterable;)Larrow/core/Either;
+	public static final fun flattenOrAccumulate (Ljava/lang/Iterable;Lkotlin/jvm/functions/Function2;)Larrow/core/Either;
 	public static final fun fold (Ljava/lang/Iterable;Larrow/typeclasses/Monoid;)Ljava/lang/Object;
 	public static final fun foldMap (Ljava/lang/Iterable;Larrow/typeclasses/Monoid;Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public static final fun getListUnit ()Ljava/util/List;

diff --git a/arrow-libs/core/arrow-core/src/commonMain/kotlin/arrow/core/Iterable.kt b/arrow-libs/core/arrow-core/src/commonMain/kotlin/arrow/core/Iterable.kt
@@ -428,6 +428,45 @@ public inline fun <A, B> Iterable<A>.traverse(f: (A) -> B?): List<B>? {
 public fun <A> Iterable<A?>.sequenceNullable(): List<A>? =
   sequence()
 
+/**
+ * Flatten a list of [Either] into a single [Either] with a list of values, or accumulates all errors using [combine].
+ */
+public inline fun <Error, A> Iterable<Either<Error, A>>.flattenOrAccumulate(combine: (Error, Error) -> Error): Either<Error, List<A>> =
+  fold<Either<Error, A>, Either<Error, ArrayList<A>>>(Right(ArrayList(collectionSizeOrDefault(10)))) { acc, res ->
+    when (res) {
+      is Right -> when (acc) {
+        is Right -> acc.also { acc.value.add(res.value) }
+        is Left -> acc
+      }
+
+      is Left -> when (acc) {
+        is Right -> res
+        is Left -> Left(combine(acc.value, res.value))
+      }
+    }
+  }
+
+/**
+ * Flatten a list of [Either] into a single [Either] with a list of values, or accumulates all errors with into an [NonEmptyList].
+ */
+public fun <Error, A> Iterable<Either<Error, A>>.flattenOrAccumulate(): Either<NonEmptyList<Error>, List<A>> {
+  val buffer = mutableListOf<Error>()
+  val res = fold<Either<Error, A>, Either<MutableList<Error>, ArrayList<A>>>(Right(ArrayList(collectionSizeOrDefault(10)))) { acc, res ->
+    when (res) {
+      is Right -> when (acc) {
+        is Right -> acc.also { acc.value.add(res.value) }
+        is Left -> acc
+      }
+
+      is Left -> when (acc) {
+        is Right -> Left(buffer.also { it.add(res.value) })
+        is Left -> Left(buffer.also { it.add(res.value) })
+      }
+    }
+  }
+  return res.mapLeft { NonEmptyList(it[0], it.drop(1)) }
+}
+
 public fun <A> Iterable<A?>.sequence(): List<A>? =
   traverse(::identity)
 

diff --git a/arrow-libs/core/arrow-core/src/commonTest/kotlin/arrow/core/IterableTest.kt b/arrow-libs/core/arrow-core/src/commonTest/kotlin/arrow/core/IterableTest.kt
@@ -21,6 +21,28 @@ import kotlin.math.min
 
 class IterableTest : StringSpec({
 
+  "flattenOrAccumulate(combine)" {
+    checkAll(Arb.list(Arb.either(Arb.string(), Arb.int()))) { list ->
+      val expected =
+        if (list.any { it.isLeft() }) list.filterIsInstance<Either.Left<String>>()
+          .fold("") { acc, either -> "$acc${either.value}" }.left()
+        else list.filterIsInstance<Either.Right<Int>>().map { it.value }.right()
+
+      list.flattenOrAccumulate { a, b -> "$a$b" } shouldBe expected
+    }
+  }
+
+  "flattenOrAccumulate" {
+    checkAll(Arb.list(Arb.either(Arb.string(), Arb.int()))) { list ->
+      val expected =
+        if (list.any { it.isLeft() }) list.filterIsInstance<Either.Left<String>>()
+          .map { it.value }.toNonEmptyListOrNull().shouldNotBeNull().left()
+        else list.filterIsInstance<Either.Right<Int>>().map { it.value }.right()
+
+      list.flattenOrAccumulate() shouldBe expected
+    }
+  }
+
     "traverse Either stack-safe" {
       // also verifies result order and execution order (l to r)
       val acc = mutableListOf<Int>()