ZStream.distributedWith on a sum type

zstream.scala

type DistributedSumType[E, Tup <: Tuple] <: Tuple = Tup match {
  case EmptyTuple => EmptyTuple
  case h *: t => Dequeue[Exit[Option[E], h]] *: DistributedSumType[E, t]
}

extension [R, E, O](v: ZStream[R, E, O]) {

  inline def distributedSumType[E1 >: E](maximumLag: Int)(using
    m: Mirror.SumOf[O],
  ): ZManaged[R, Nothing, DistributedSumType[E1, m.MirroredElemTypes]] = {
    val n = constValue[Tuple.Size[m.MirroredElemTypes]]

    def decide(o: O): UIO[Int => Boolean] = ZIO.succeed(_ == m.ordinal(o))
    v.distributedWith[E1](n, maximumLag = maximumLag, decide).map { ls =>
      Tuple.fromArray(ls.toArray).asInstanceOf[DistributedSumType[E1, m.MirroredElemTypes]]
    }
  }
}