zio · guizmaii · Apr 10, 2025 · Apr 7, 2025 · Apr 7, 2025 · Apr 7, 2025
diff --git a/streams-tests/shared/src/test/scala/zio/stream/ZPipelineSpec.scala b/streams-tests/shared/src/test/scala/zio/stream/ZPipelineSpec.scala
@@ -2,8 +2,9 @@ package zio.stream
 
 import zio._
 import zio.stream.encoding.EncodingException
-import zio.test.Assertion.{equalTo, fails}
+import zio.test.Assertion.{equalTo, fails, isEmpty}
 import zio.test._
+
 import scala.io.Source
 
 object ZPipelineSpec extends ZIOBaseSpec {
@@ -294,7 +295,62 @@ object ZPipelineSpec extends ZIOBaseSpec {
           }
           .runCollect
           .map(assert(_)(equalTo(Chunk.range(1, 21))))
-      }
+      },
+      suite("mapEitherChunk")(
+        test("with empty chunk") {
+          val chunk = Chunk.empty[Int]
+          for {
+            result <- ZStream
+                        .fromChunk(chunk)
+                        .via(ZPipeline.mapEitherChunked(i => Right(i)))
+                        .run(ZSink.collectAll)
+          } yield assert(result)(isEmpty)
+        },
+        test("with a 1 element chunk - Right") {
+          val chunk = Chunk(1)
+          for {
+            result <- ZStream
+                        .fromChunk(chunk)
+                        .via(ZPipeline.mapEitherChunked(i => Right(i)))
+                        .run(ZSink.collectAll)
+          } yield assert(result)(equalTo(Chunk(1)))
+        },
+        test("with a 1 element chunk - Left") {
+          val chunk = Chunk(1)
+          for {
+            collector <- Queue.unbounded[Int]
+            result <- ZStream
+                        .fromChunk(chunk)
+                        .via(ZPipeline.mapEitherChunked(i => Left(i)))
+                        .run(ZSink.fromQueue(collector))
+                        .exit
+            collected <- collector.takeAll
+          } yield assert(result)(fails(equalTo(1))) && assert(collected)(isEmpty)
+        },
+        test("Keeps order and values intact") {
+          val range = 1.to(10)
+          val chunk = Chunk.fromIterable(range)
+          for {
+            result <- ZStream
+                        .fromChunk(chunk)
+                        .via(ZPipeline.mapEitherChunked(i => Right(i)))
+                        .run(ZSink.collectAll)
+          } yield assert(result)(equalTo(Chunk(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)))
+        },
+        test("stop at the first Left") {
+          val range = 1.to(10)
+          val chunk = Chunk.fromIterable(range)
+          for {
+            collector <- Queue.unbounded[Int]
+            result <- ZStream
+                        .fromChunk(chunk)
+                        .via(ZPipeline.mapEitherChunked(i => if (i == 5) Left(i) else Right(i)))
+                        .run(ZSink.fromQueue(collector))
+                        .exit
+            collected <- collector.takeAll
+          } yield assert(result)(fails(equalTo(5))) && assert(collected)(equalTo(Chunk(1, 2, 3, 4)))
+        }
+      )
     )
 
   val weirdStringGenForSplitLines: Gen[Any, Chunk[String]] = Gen

diff --git a/streams/shared/src/main/scala/zio/stream/ZPipeline.scala b/streams/shared/src/main/scala/zio/stream/ZPipeline.scala
@@ -21,19 +21,9 @@ import zio.internal.SingleThreadedRingBuffer
 import zio.stacktracer.TracingImplicits.disableAutoTrace
 import zio.stream.encoding.EncodingException
 import zio.stream.internal.CharacterSet.{BOM, CharsetUtf32BE, CharsetUtf32LE}
-import zio.stream.internal.SingleProducerAsyncInput
-
-import java.nio.{Buffer, ByteBuffer, CharBuffer}
-import java.nio.charset.{
-  CharacterCodingException,
-  Charset,
-  CharsetDecoder,
-  CoderResult,
-  MalformedInputException,
-  StandardCharsets,
-  UnmappableCharacterException
-}
-import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import java.nio.charset._
+import java.nio.{ByteBuffer, CharBuffer}
 
 /**
  * A `ZPipeline[Env, Err, In, Out]` is a polymorphic stream transformer.
@@ -1788,6 +1778,53 @@ object ZPipeline extends ZPipelinePlatformSpecificConstructors {
       ZChannel.identity[Nothing, Chunk[In], Any].concatMap(_.map(f).map(_.channel).fold(ZChannel.unit)(_ *> _))
     )
 
+  /**
+   * Creates a pipeline that maps chunks of elements with the specified
+   * function.
+   *
+   * Will stop on the first Left found
+   */
+  def mapEitherChunked[Env, Err, In, Out](
+    f: In => Either[Err, Out]
+  )(implicit trace: Trace): ZPipeline[Env, Err, In, Out] = {
+    lazy val reader: ZChannel[Env, Err, Chunk[In], Any, Err, Chunk[Out], Any] =
+      ZChannel.readWithCause(
+        chunk => {
+          val size = chunk.size
+
+          if (size == 0) reader
+          else if (size == 1) {
+            val a = chunk.head
+
+            f(a) match {
+              case r: Right[?, Out] => ZChannel.write(Chunk.single(r.value)) *> reader
+              case l: Left[Err, ?]  => ZChannel.refailCause(Cause.fail(l.value))
+            }
+          } else {
+            val builder: ChunkBuilder[Out] = ChunkBuilder.make[Out](chunk.size)
+            val iterator: Iterator[In]     = chunk.iterator
+            var error: Err                 = null.asInstanceOf[Err]
+
+            while (iterator.hasNext && (error == null)) {
+              val a = iterator.next()
+              f(a) match {
+                case r: Right[?, Out] => builder.addOne(r.value)
+                case l: Left[Err, ?]  => error = l.value
+              }
+            }
+
+            val values = builder.result()
+            val next   = if (error == null) reader else ZChannel.refailCause(Cause.fail(error))
+            if (values.nonEmpty) ZChannel.write(values) *> next else next
+          }
+        },
+        err => ZChannel.refailCause(err),
+        done => ZChannel.succeed(done)
+      )
+
+    new ZPipeline(reader)
+  }
+
   /**
    * Creates a pipeline that maps elements with the specified effectful
    * function.