Scala 3 collection partitioning with subtypes

Question

In Scala 3, let's say I have a List[Try[String]]. Can I split it up into success and failures, such that each list has the appropriate subtype?

If I do the following:

import scala.util.{Try, Success, Failure}
val tries = List(Success("1"), Failure(Exception("2")))
val (successes, failures) = tries.partition(_.isSuccess)

then successes and failures are still of type List[Try[String]]. The same goes if I filter based on the type:

val successes = tries.filter(_.isInstanceOf[Success[String]])

I could of course cast to Success and Failure respectively, but is there a type-safe way to achieve this?

Use `tries.partitionMap(_.toEither)` - As O always say, the [**Scaladoc**](https://www.scala-lang.org/api/current/) is your friend. — Luis Miguel Mejía Suárez, Oct 07 '21 at 15:02
@LuisMiguelMejíaSuárez thanks! Although I don't know how I would have searched for this in the Scaladoc — mitchus, Oct 07 '21 at 15:12
@LuisMiguelMejíaSuárez ok the trick here is that `Try` has a `toEither` method which splits to the proper type. What if we have a regular sealed trait? — mitchus, Oct 07 '21 at 15:16
then you need to use pattern matching inside the `partitionMap` if you have more than two options then there isn't anything out of the box AFAIK. I would guess **Shapeless** would have something but I would just write my own little method using tail-recursion. — Luis Miguel Mejía Suárez, Oct 07 '21 at 15:28
@LuisMiguelMejíaSuárez *"I would guess Shapeless would have something"* OP asks about Scala 3. — Dmytro Mitin, Oct 07 '21 at 18:12
@DmytroMitin there is **Shapeless 3** AFAIK, although not sure what is its scope and what of that is already covered. — Luis Miguel Mejía Suárez, Oct 07 '21 at 18:25
@LuisMiguelMejíaSuárez [Shapeless-3](https://github.com/typelevel/shapeless-3) is currently about deriving and an only type class `Typeable` out of dozens of type classes existing in [Shapeless-2](https://github.com/milessabin/shapeless). — Dmytro Mitin, Oct 07 '21 at 19:01
@DmytroMitin good to know. Do you know if the plan to support everything they did? Or just focus on deriving this time? — Luis Miguel Mejía Suárez, Oct 07 '21 at 19:16
@LuisMiguelMejíaSuárez Well, not sure but I guess Shapeless-3 `Typeable` was just a proof of concept. It's an early time for many libraries in Scala-3 ecosphere (especially macro libraries). But it's better to look in chats https://gitter.im/milessabin/shapeless and https://discord.com/channels/632277896739946517/632312089616056350 — Dmytro Mitin, Oct 07 '21 at 19:36
@LuisMiguelMejíaSuárez I forgiot that there are also `Annotation`, `Annotations` besides `Typeable` there. — Dmytro Mitin, Oct 08 '21 at 14:59
@mitchus *"What if we have a regular sealed trait?"* In Scala 2 I would do something like https://scastie.scala-lang.org/DmytroMitin/uQp603sXT7WFYmYntDXmIw I guess it can be translated into Scala 3. — Dmytro Mitin, Oct 08 '21 at 15:01
@mitchus I managed to translate into Scala 3 although the translation turns to be wordy. — Dmytro Mitin, Oct 10 '21 at 02:24

Dmytro Mitin · Accepted Answer · 2022-10-24T03:33:50.787

@Luis Miguel Mejía Suárez:

Use tries.partitionMap(_.toEither)

@mitchus:

@LuisMiguelMejíaSuárez ok the trick here is that Try has a toEither method which splits to the proper type. What if we have a regular sealed trait?

In Scala 2 I would do something like

import shapeless.{:+:, ::, CNil, Coproduct, Generic, HList, HNil, Inl, Inr, Poly0}
import shapeless.ops.coproduct.ToHList
import shapeless.ops.hlist.{FillWith, Mapped, Tupler}

trait Loop[C <: Coproduct, L <: HList] {
  def apply(c: C, l: L): L
}
object Loop {
  implicit def recur[H, CT <: Coproduct, HT <: HList](implicit
    loop: Loop[CT, HT]
  ): Loop[H :+: CT, List[H] :: HT] = {
    case (Inl(h), hs :: ht) => (h :: hs) :: ht
    case (Inr(ct), hs :: ht) => hs :: loop(ct, ht)
  }

  implicit val base: Loop[CNil, HNil] = (_, l) => l
}

object nilPoly extends Poly0 {
  implicit def cse[A]: Case0[List[A]] = at(Nil)
}

def partition[A, C <: Coproduct, L <: HList, L1 <: HList](as: List[A])(implicit
  generic: Generic.Aux[A, C],
  toHList: ToHList.Aux[C, L],
  mapped: Mapped.Aux[L, List, L1],
  loop: Loop[C, L1],
  fillWith: FillWith[nilPoly.type, L1],
  tupler: Tupler[L1]
): tupler.Out = {
  val partitionHList: L1 = as.foldRight(fillWith())((a, l1) =>
    loop(generic.to(a), l1)
  )

  tupler(partitionHList)
}

sealed trait A
case class B(i: Int) extends A
case class C(i: Int) extends A
case class D(i: Int) extends A

partition(List[A](B(1), B(2), C(1), C(2), D(1), D(2), B(3), C(3))) 
// (List(B(1), B(2), B(3)),List(C(1), C(2), C(3)),List(D(1), D(2))): (List[B], List[C], List[D])

https://scastie.scala-lang.org/DmytroMitin/uQp603sXT7WFYmYntDXmIw/1

I managed to translate this code into Scala 3 although the translation turned to be wordy (I remplemented Generic and Coproduct)

import scala.annotation.tailrec
import scala.deriving.Mirror

object App1 {
  // ============= Generic =====================
  trait Generic[T] {
    type Repr
    def to(t: T): Repr
    def from(r: Repr): T
  }
  object Generic {
    type Aux[T, Repr0] = Generic[T] { type Repr = Repr0 }
    def instance[T, Repr0](f: T => Repr0, g: Repr0 => T): Aux[T, Repr0] =
      new Generic[T] {
        override type Repr = Repr0
        override def to(t: T): Repr0 = f(t)
        override def from(r: Repr0): T = g(r)
      }

    object ops {
      extension [A](a: A) {
        def toRepr(using g: Generic[A]): g.Repr = g.to(a)
      }

      extension [Repr](a: Repr) {
        def to[A](using g: Generic.Aux[A, Repr]): A = g.from(a)
      }
    }

    given [T <: Product](using
      m: Mirror.ProductOf[T]
    ): Aux[T, m.MirroredElemTypes] = instance(
      _.productIterator
       .foldRight[Tuple](EmptyTuple)(_ *: _)
       .asInstanceOf[m.MirroredElemTypes],
      m.fromProduct(_).asInstanceOf[T]
    )

    inline given [T, C <: Coproduct](using
      m: Mirror.SumOf[T],
      ev: Coproduct.ToCoproduct[m.MirroredElemTypes] =:= C
    ): Generic.Aux[T, C] =
      instance(
        matchExpr[T, C](_).asInstanceOf[C],
        Coproduct.unsafeFromCoproduct(_).asInstanceOf[T]
      )

    import scala.quoted.*

    inline def matchExpr[T, C <: Coproduct](ident: T): Coproduct =
      ${matchExprImpl[T, C]('ident)}

    def matchExprImpl[T: Type, C <: Coproduct : Type](
      ident: Expr[T]
    )(using Quotes): Expr[Coproduct] = {
      import quotes.reflect.*

      def unwrapCoproduct(typeRepr: TypeRepr): List[TypeRepr] = typeRepr match {
        case AppliedType(_, List(typ1, typ2)) => typ1 :: unwrapCoproduct(typ2)
        case _  => Nil
      }

      val typeReprs = unwrapCoproduct(TypeRepr.of[C])

      val methodIdent =
        Ident(TermRef(TypeRepr.of[Coproduct.type], "unsafeToCoproduct"))

      def caseDefs(ident: Term): List[CaseDef] =
        typeReprs.zipWithIndex.map { (typeRepr, i) =>
          CaseDef(
            Typed(ident, Inferred(typeRepr) /*TypeIdent(typeRepr.typeSymbol)*/),
            None,
            Block(
              Nil,
              Apply(
                methodIdent,
                List(Literal(IntConstant(i)), ident)
              )
            )
          )
        }

      def matchTerm(ident: Term): Term = Match(ident, caseDefs(ident))

      matchTerm(ident.asTerm).asExprOf[Coproduct]
    }
  }

  // ============= Coproduct =====================
  sealed trait Coproduct extends Product with Serializable
  sealed trait +:[+H, +T <: Coproduct] extends Coproduct
  final case class Inl[+H, +T <: Coproduct](head: H) extends (H +: T)
  final case class Inr[+H, +T <: Coproduct](tail: T) extends (H +: T)
  sealed trait CNil extends Coproduct

  object Coproduct {
    def unsafeToCoproduct(length: Int, value: Any): Coproduct =
      (0 until length).foldLeft[Coproduct](Inl(value))((c, _) => Inr(c))

    @tailrec
    def unsafeFromCoproduct(c: Coproduct): Any = c match {
      case Inl(h) => h
      case Inr(c) => unsafeFromCoproduct(c)
      case _: CNil => sys.error("impossible")
    }

    type ToCoproduct[T <: Tuple] <: Coproduct = T match {
      case EmptyTuple => CNil
      case h *: t => h +: ToCoproduct[t]
    }

//    type ToTuple[C <: Coproduct] <: Tuple = C match {
//      case CNil => EmptyTuple
//      case h +: t => h *: ToTuple[t]
//    }

    trait ToTuple[C <: Coproduct] {
      type Out <: Tuple
    }
    object ToTuple {
      type Aux[C <: Coproduct, Out0 <: Tuple] = ToTuple[C] { type Out = Out0 }
      def instance[C <: Coproduct, Out0 <: Tuple]: Aux[C, Out0] =
        new ToTuple[C] { override type Out = Out0 }

      given [H, T <: Coproduct](using 
        toTuple: ToTuple[T]
      ): Aux[H +: T, H *: toTuple.Out] = instance
      given Aux[CNil, EmptyTuple] = instance
    }
  }
}

// different file
import App1.{+:, CNil, Coproduct, Generic, Inl, Inr}

object App2 {    
  trait Loop[C <: Coproduct, L <: Tuple] {
    def apply(c: C, l: L): L
  }
  object Loop {
    given [H, CT <: Coproduct, HT <: Tuple](using 
      loop: Loop[CT, HT]
    ): Loop[H +: CT, List[H] *: HT] = {
      case (Inl(h), hs *: ht) => (h :: hs) *: ht
      case (Inr(ct), hs *: ht) => hs *: loop(ct, ht)
    }

    given Loop[CNil, EmptyTuple] = (_, l) => l
  }

  trait FillWithNil[L <: Tuple] {
    def apply(): L
  }
  object FillWithNil {
    given [H, T <: Tuple](using 
      fillWithNil: FillWithNil[T]
    ): FillWithNil[List[H] *: T] = () => Nil *: fillWithNil()
    given FillWithNil[EmptyTuple] = () => EmptyTuple
  }

  def partition[A, /*L <: Tuple,*/ L1 <: Tuple](as: List[A])(using
    generic: Generic.Aux[A, _ <: Coproduct],
    toTuple: Coproduct.ToTuple[generic.Repr],
    //ev0: Coproduct.ToTuple[generic.Repr] =:= L, // compile-time NPE
    ev: Tuple.Map[toTuple.Out/*L*/, List] =:= L1,
    loop: Loop[generic.Repr, L1],
    fillWith: FillWithNil[L1]
  ): L1 = as.foldRight(fillWith())((a, l1) => loop(generic.to(a), l1))

  sealed trait A
  case class B(i: Int) extends A
  case class C(i: Int) extends A
  case class D(i: Int) extends A

  def main(args: Array[String]): Unit = {
    println(partition(List[A](B(1), B(2), C(1), C(2), D(1), D(2), B(3), C(3))))
  // (List(B(1), B(2), B(3)),List(C(1), C(2), C(3)),List(D(1), D(2)))
  }
}

Scala 3.0.2

In the macro (generating pattern matching) Inferred(typeRepr) should be instead of TypeIdent(typeRepr.typeSymbol), otherwise this doesn't work for parametric case classes. Actually the macro can be removed at all if we use mirror.ordinal. Simplified version is

import scala.deriving.Mirror
import scala.util.NotGiven

trait Generic[T] {
  type Repr
  def to(t: T): Repr
  def from(r: Repr): T
}

object Generic {
  type Aux[T, Repr0] = Generic[T] {type Repr = Repr0}

  def instance[T, Repr0](f: T => Repr0, g: Repr0 => T): Aux[T, Repr0] =
    new Generic[T] {
      override type Repr = Repr0
      override def to(t: T): Repr0 = f(t)
      override def from(r: Repr0): T = g(r)
    }

  object ops {
    extension[A] (a: A) {
      def toRepr(using g: Generic[A]): g.Repr = g.to(a)
    }

    extension[Repr] (a: Repr) {
      def to[A](using g: Generic.Aux[A, Repr]): A = g.from(a)
    }
  }

  given [T <: Product](using
    // ev: NotGiven[T <:< Tuple],
    // ev1: NotGiven[T <:< Coproduct],
    m: Mirror.ProductOf[T],
    m1: Mirror.ProductOf[m.MirroredElemTypes]
  ): Aux[T, m.MirroredElemTypes] = instance(
    m1.fromProduct(_),
    m.fromProduct(_)
  )

  given[T, C <: Coproduct](using
    // ev: NotGiven[T <:< Tuple],
    // ev1: NotGiven[T <:< Coproduct],
    m: Mirror.SumOf[T],
    ev2: Coproduct.ToCoproduct[m.MirroredElemTypes] =:= C
  ): Generic.Aux[T, C/*Coproduct.ToCoproduct[m.MirroredElemTypes]*/] = {
    instance(
      t => Coproduct.unsafeToCoproduct(m.ordinal(t), t).asInstanceOf[C],
      Coproduct.unsafeFromCoproduct(_).asInstanceOf[T]
    )
  }
}

sealed trait Coproduct extends Product with Serializable
sealed trait +:[+H, +T <: Coproduct] extends Coproduct
final case class Inl[+H, +T <: Coproduct](head: H) extends (H +: T)
final case class Inr[+H, +T <: Coproduct](tail: T) extends (H +: T)
sealed trait CNil extends Coproduct

object Coproduct {
  def unsafeToCoproduct(length: Int, value: Any): Coproduct =
    (0 until length).foldLeft[Coproduct](Inl(value))((c, _) => Inr(c))

  @scala.annotation.tailrec
  def unsafeFromCoproduct(c: Coproduct): Any = c match {
    case Inl(h) => h
    case Inr(c) => unsafeFromCoproduct(c)
    case _: CNil => sys.error("impossible")
  }

  type ToCoproduct[T <: Tuple] <: Coproduct = T match {
    case EmptyTuple => CNil
    case h *: t => h +: ToCoproduct[t]
  }

  type ToTuple[C <: Coproduct] <: Tuple = C match {
    case CNil => EmptyTuple
    case h +: t => h *: ToTuple[t]
  }
}

Replacing type classes with compile-time/inline methods and match types

import scala.compiletime.erasedValue

inline def loop[C <: Coproduct, L <: Tuple](c: C, l: L): L = (inline erasedValue[C] match {
  case _: CNil => inline erasedValue[L] match {
    case _: EmptyTuple => EmptyTuple
  }
  case _: (h +: ct) => inline erasedValue[L] match {
    case _: (List[`h`] *: ht) => (c, l) match {
      case (Inl(h_v: `h`), (hs_v: List[`h`]) *: (ht_v: `ht`)) => 
        (h_v :: hs_v) *: ht_v
      case (Inr(ct_v: `ct`), (hs_v: List[`h`]) *: (ht_v: `ht`)) => 
        hs_v *: loop[ct, ht](ct_v, ht_v)
    }
  }
}).asInstanceOf[L]

inline def fillWithNil[L <: Tuple]: L = (inline erasedValue[L] match {
  case _: EmptyTuple => EmptyTuple
  case _: (List[h] *: t) => Nil *: fillWithNil[t]
}).asInstanceOf[L]

type TupleList[C <: Coproduct] = Tuple.Map[Coproduct.ToTuple[C], List]

inline def partition[A](as: List[A])(using
  generic: Generic.Aux[A, _ <: Coproduct]
): TupleList[generic.Repr] =
  as.foldRight(fillWithNil[TupleList[generic.Repr]])((a, l1) => loop(generic.to(a), l1))

sealed trait A
case class B(i: Int) extends A
case class C(i: Int) extends A
case class D(i: Int) extends A

@main def test = {
  println(partition(List[A](B(1), B(2), C(1), C(2), D(1), D(2), B(3), C(3))))
  // (List(B(1), B(2), B(3)),List(C(1), C(2), C(3)),List(D(1), D(2)))
}

Tested in 3.2.0 https://scastie.scala-lang.org/DmytroMitin/940QaiqDQQ2QegCyxTbEIQ/1

How to access parameter list of case class in a dotty macro

Alternative implementation of loop

//Loop[C, L] = L
type Loop[C <: Coproduct, L <: Tuple] <: Tuple = C match {
  case CNil    => CNilLoop[L]
  case h +: ct => CConsLoop[h, ct, L]
}
// match types seem not to support nested type matching
type CNilLoop[L <: Tuple] <: Tuple = L match {
  case EmptyTuple => EmptyTuple
}
type CConsLoop[H, CT <: Coproduct, L <: Tuple] <: Tuple = L match {
  case List[H] *: ht => List[H] *: Loop[CT, ht]
}
/*inline*/ def loop0[C <: Coproduct, L <: Tuple](c: C, l: L): Loop[C, L] = /*inline*/ c match {
  case _: CNil => /*inline*/ l match {
    case _: EmptyTuple => EmptyTuple
  }
  case c: (h +: ct) => /*inline*/ l match {
    case l: (List[`h`] *: ht) => (c, l) match {
      case (Inl(h_v/*: `h`*/), (hs_v/*: List[`h`]*/) *: (ht_v/*: `ht`*/)) =>
        (h_v :: hs_v) *: ht_v.asInstanceOf[Loop[ct, ht]]
      case (Inr(ct_v/*: `ct`*/), (hs_v/*: List[`h`]*/) *: (ht_v/*: `ht`*/)) => 
        hs_v *: loop0[ct, ht](ct_v, ht_v)
    }
  }
}
/*inline*/ def loop[C <: Coproduct, L <: Tuple](c: C, l: L): L = loop0(c, l).asInstanceOf[L]

Another implementation for Scala 2: Split list of algebraic date type to lists of branches?

That's cool! But I'm not sure I follow all of it. Could you add a high-level explanation of what happens? — mitchus, Oct 11 '21 at 11:07
@mitchus The type class `Generic` knows how to transform case classes into `Tuple` (former `HList`) and vice versa, sealed traits into `Coproduct` and vice versa. So firstly on type level for a sealed trait `A` I calculate its `Coproduct` representation `B +: C +: D +: CNil`, transform it into a tuple `(B, C, D)` (aka `B *: C *: D *: EmptyTuple`), mapping it with `List` to `(List[B], List[C], List[D])`, this will be the return type of `partition` taking an `List[A]`. — Dmytro Mitin, Oct 11 '21 at 13:14
@mitchus Secondly on value level we create `(Nil, Nil, Nil)` with `fillWith()` (of type `(List[B], List[C], List[D])`), then we're going through incoming `List[A]` element by element (`foldLeft`). For an element `a` its representation `generic.to(a)` can be either `Inl(b)` or `Inr(Inl(c))` or `Inr(Inr(Inl(d)))`. Depending on how many `Inr`s are in `Inr...(Inr(Inl(` the type class `Loop` adds `b`, `c` or `d` to the `List[B]`, `List[C]` or `List[D]`. — Dmytro Mitin, Oct 11 '21 at 13:14

Scala 3 collection partitioning with subtypes

1 Answers1

Linked