Type based collection partitioning in Scala

Question

Given the following data model:

sealed trait Fruit

case class Apple(id: Int, sweetness: Int) extends Fruit

case class Pear(id: Int, color: String) extends Fruit

I've been looking to implement a segregate basket function which for the given basket of fruits will return separate baskets of apples and pears:

def segregateBasket(fruitBasket: Set[Fruit]): (Set[Apple], Set[Pear])

I've attempted a couple of approaches, but none of them seems to be fitting the bill perfectly. Below are my attempts:

  def segregateBasket1(fruitBasket: Set[Fruit]): (Set[Apple], Set[Pear]) = fruitBasket
    .partition(_.isInstanceOf[Apple])
    .asInstanceOf[(Set[Apple], Set[Pear])]

This is the most concise solution I've found, but suffers from explicit type casting via asInstanceOf and is going to be a pain to extend should I decide to add additional types of fruits. Therefore:

  def segregateBasket2(fruitBasket: Set[Fruit]): (Set[Apple], Set[Pear]) = {
    val mappedFruits = fruitBasket.groupBy(_.getClass)
    val appleSet = mappedFruits.getOrElse(classOf[Apple], Set()).asInstanceOf[Set[Apple]]
    val pearSet = mappedFruits.getOrElse(classOf[Pear], Set()).asInstanceOf[Set[Pear]]
    (appleSet, pearSet)
  }

Resolves the problem of additional fruit types (extension really easy), but still strongly depends on risky type casting 'asInstanceOf' which I'd rather avoid. Therefore:

  def segregateBasket3(fruitBasket: Set[Fruit]): (Set[Apple], Set[Pear]) = {
    val appleSet = collection.mutable.Set[Apple]()
    val pearSet = collection.mutable.Set[Pear]()

    fruitBasket.foreach {
      case a: Apple => appleSet += a
      case p: Pear => pearSet += p
    }
    (appleSet.toSet, pearSet.toSet)
  }

Resolves the problem of explicit casting, but uses mutable collections and ideally I'd like to stick with immutable collections and idiomatic code.

I've looked here: Scala: Filtering based on type for some inspiration, but couldn't find a better approach either.

Does anyone have any suggestions on how this functionality can be better implemented in Scala?

Answer 1

  val emptyBaskets: (List[Apple], List[Pear]) = (Nil, Nil)

  def separate(fruits: List[Fruit]): (List[Apple], List[Pear]) = {
    fruits.foldRight(emptyBaskets) { case (f, (as, ps)) =>
      f match {
        case a @ Apple(_, _) => (a :: as, ps)
        case p @ Pear(_, _)  => (as, p :: ps)
      }
    }
  }