Avoid storing more transformed elements than necessary (#3376)

arrow-kt · Feb 18, 2024 · e2c53e8 · e2c53e8
1 parent 89b4c8e
commit e2c53e8
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 10 deletions.
diff --git a/arrow-libs/core/arrow-core/api/arrow-core.api b/arrow-libs/core/arrow-core/api/arrow-core.api
@@ -3614,6 +3614,8 @@ public final class arrow/core/raise/RaiseKt {
 	public static final fun forEachAccumulating (Larrow/core/raise/Raise;Ljava/util/Iterator;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function2;)V
 	public static final fun forEachAccumulating (Larrow/core/raise/Raise;Lkotlin/sequences/Sequence;Lkotlin/jvm/functions/Function2;)V
 	public static final fun forEachAccumulating (Larrow/core/raise/Raise;Lkotlin/sequences/Sequence;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function2;)V
+	public static final synthetic fun forEachAccumulatingImpl (Larrow/core/raise/Raise;Ljava/util/Iterator;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function3;)V
+	public static final synthetic fun forEachAccumulatingImpl (Larrow/core/raise/Raise;Ljava/util/Iterator;Lkotlin/jvm/functions/Function3;)V
 	public static final fun get (Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public static final fun get (Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun getOrElse (Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;

diff --git a/arrow-libs/core/arrow-core/src/commonMain/kotlin/arrow/core/raise/RaiseAccumulate.kt b/arrow-libs/core/arrow-core/src/commonMain/kotlin/arrow/core/raise/RaiseAccumulate.kt
@@ -14,7 +14,6 @@ import arrow.core.Validated
 import arrow.core.ValidatedDeprMsg
 import arrow.core.collectionSizeOrDefault
 import arrow.core.ValidatedNel
-import arrow.core.mapOrAccumulate
 import arrow.core.nonEmptyListOf
 import arrow.core.toNonEmptyListOrNull
 import arrow.core.toNonEmptySetOrNull
@@ -24,6 +23,7 @@ import kotlin.contracts.contract
 import kotlin.experimental.ExperimentalTypeInference
 import kotlin.jvm.JvmMultifileClass
 import kotlin.jvm.JvmName
+import kotlin.jvm.JvmSynthetic
 
 /**
  * Accumulate the errors from running both [action1] and [action2] using the given [combine] function.
@@ -521,10 +521,19 @@ public inline fun <Error, A> Raise<Error>.forEachAccumulating(
   iterator: Iterator<A>,
   combine: (Error, Error) -> Error,
   @BuilderInference block: RaiseAccumulate<Error>.(A) -> Unit
+): Unit = forEachAccumulatingImpl(iterator, combine) { item, _ -> block(item) }
+
+@PublishedApi @JvmSynthetic
+internal inline fun <Error, A> Raise<Error>.forEachAccumulatingImpl(
+  iterator: Iterator<A>,
+  combine: (Error, Error) -> Error,
+  @BuilderInference block: RaiseAccumulate<Error>.(item: A, hasErrors: Boolean) -> Unit
 ) {
   var error: Any? = EmptyValue
   for (item in iterator) {
-    recover({ block(RaiseAccumulate(this), item) }) { errors ->
+    recover({
+      block(RaiseAccumulate(this), item, error != EmptyValue)
+    }) { errors ->
       error = combine(error, errors.reduce(combine), combine)
     }
   }
@@ -547,10 +556,22 @@ public inline fun <Error, A> Raise<NonEmptyList<Error>>.forEachAccumulating(
 public inline fun <Error, A> Raise<NonEmptyList<Error>>.forEachAccumulating(
   iterator: Iterator<A>,
   @BuilderInference block: RaiseAccumulate<Error>.(A) -> Unit
+): Unit = forEachAccumulatingImpl(iterator) { item, _ -> block(item) }
+
+/**
+ * Allows to change what to do once the first error is raised.
+ * Used to provide more performant [mapOrAccumulate].
+ */
+@PublishedApi @JvmSynthetic
+internal inline fun <Error, A> Raise<NonEmptyList<Error>>.forEachAccumulatingImpl(
+  iterator: Iterator<A>,
+  @BuilderInference block: RaiseAccumulate<Error>.(item: A, hasErrors: Boolean) -> Unit
 ) {
   val error: MutableList<Error> = mutableListOf()
   for (item in iterator) {
-    recover({ block(RaiseAccumulate(this), item) }) {
+    recover({
+      block(RaiseAccumulate(this), item, error.isNotEmpty())
+    }) {
       error.addAll(it)
     }
   }
@@ -570,7 +591,9 @@ public inline fun <Error, A, B> Raise<Error>.mapOrAccumulate(
   combine: (Error, Error) -> Error,
   @BuilderInference transform: RaiseAccumulate<Error>.(A) -> B
 ): List<B> = buildList(iterable.collectionSizeOrDefault(10)) {
-  forEachAccumulating(iterable, combine) { add(transform(it)) }
+  forEachAccumulatingImpl(iterable.iterator(), combine) { item, hasErrors ->
+    transform(item).also { if (!hasErrors) add(it) }
+  }
 }
 
 /**
@@ -585,7 +608,9 @@ public inline fun <Error, A, B> Raise<NonEmptyList<Error>>.mapOrAccumulate(
   iterable: Iterable<A>,
   @BuilderInference transform: RaiseAccumulate<Error>.(A) -> B
 ): List<B> = buildList(iterable.collectionSizeOrDefault(10)) {
-  forEachAccumulating(iterable) { add(transform(it)) }
+  forEachAccumulatingImpl(iterable.iterator()) { item, hasErrors ->
+    transform(item).also { if (!hasErrors) add(it) }
+  }
 }
 
 /**
@@ -601,7 +626,9 @@ public inline fun <Error, A, B> Raise<Error>.mapOrAccumulate(
   combine: (Error, Error) -> Error,
   @BuilderInference transform: RaiseAccumulate<Error>.(A) -> B
 ): List<B> = buildList {
-  forEachAccumulating(sequence, combine) { add(transform(it)) }
+  forEachAccumulatingImpl(sequence.iterator(), combine) { item, hasErrors ->
+    transform(item).also { if (!hasErrors) add(it) }
+  }
 }
 
 /**
@@ -616,7 +643,9 @@ public inline fun <Error, A, B> Raise<NonEmptyList<Error>>.mapOrAccumulate(
   sequence: Sequence<A>,
   @BuilderInference transform: RaiseAccumulate<Error>.(A) -> B
 ): List<B> = buildList {
-  forEachAccumulating(sequence) { add(transform(it)) }
+  forEachAccumulatingImpl(sequence.iterator()) { item, hasErrors ->
+    transform(item).also { if (!hasErrors) add(it) }
+  }
 }
 
 /**
@@ -644,22 +673,30 @@ public inline fun <Error, A, B> Raise<NonEmptyList<Error>>.mapOrAccumulate(
   nonEmptySet: NonEmptySet<A>,
   @BuilderInference transform: RaiseAccumulate<Error>.(A) -> B
 ): NonEmptySet<B> = buildSet(nonEmptySet.size) {
-  forEachAccumulating(nonEmptySet) { add(transform(it)) }
+  forEachAccumulatingImpl(nonEmptySet.iterator()) { item, hasErrors ->
+    transform(item).also { if (!hasErrors) add(it) }
+  }
 }.toNonEmptySetOrNull()!!
 
+@RaiseDSL
 public inline fun <K, Error, A, B> Raise<Error>.mapOrAccumulate(
   map: Map<K, A>,
   combine: (Error, Error) -> Error,
   @BuilderInference transform: RaiseAccumulate<Error>.(Map.Entry<K, A>) -> B
 ): Map<K, B> = buildMap(map.size) {
-  forEachAccumulating(map.entries, combine) { put(it.key, transform(it)) }
+  forEachAccumulatingImpl(map.entries.iterator(), combine) { item, hasErrors ->
+    transform(item).also { if (!hasErrors) put(item.key, it) }
+  }
 }
 
+@RaiseDSL
 public inline fun <K, Error, A, B> Raise<NonEmptyList<Error>>.mapOrAccumulate(
   map: Map<K, A>,
   @BuilderInference transform: RaiseAccumulate<Error>.(Map.Entry<K, A>) -> B
 ): Map<K, B> = buildMap(map.size) {
-  forEachAccumulating(map.entries) { put(it.key, transform(it)) }
+  forEachAccumulatingImpl(map.entries.iterator()) { item, hasErrors ->
+    transform(item).also { if (!hasErrors) put(item.key, it) }
+  }
 }
 
 /**