Better way to create state transformation for adding set elements to Map[K, Set[V]]

Question

I have Map[K, Set[V]] and I am using Scalaz Lenses and State to add elements to it.

So far, I see myself doing this repeatedly:

myMapLens.member(key) %= {
    case Some(vals) => Some(vals + newValue)
    case None       => Some(Set(newValue))
}

Is there a better way to do this using Scalaz ? Casting my value set into Some(...) every time seems wasteful.

Specifically, is there a way to compose Scalaz MapLens and SetLens to achieve this ?

Answer 1

You can write an adapter to "flatten" the Option :

import scalaz._, Scalaz._

def noneZero[A: Monoid]: Lens[Option[A], A] = Lens.lensg(_ => Some(_), _.orZero)

This is a little more generic than you need, but has the same behavior for your use case:

val myMapLens = Lens.lensId[Map[String, Set[Int]]]

val myLens = myMapLens.member("foo").andThen(noneZero).contains(1)

You could of course use any of the other methods on SetLens — contains just makes for a nice demonstration:

scala> myLens.get(Map("foo" -> Set(1)))
res0: Boolean = true

scala> myLens.get(Map("bar" -> Set(1)))
res1: Boolean = false

scala> myLens.set(Map("foo" -> Set(2)), true)
res2: Map[String,Set[Int]] = Map(foo -> Set(2, 1))

scala> myLens.set(Map("bar" -> Set(2)), true)
res3: Map[String,Set[Int]] = Map(bar -> Set(2), foo -> Set(1))

scala> myLens.set(Map("foo" -> Set(1)), false)
res4: Map[String,Set[Int]] = Map(foo -> Set())

The following is arguably a slightly more principled way to write the adapter:

def noneZero[A: Monoid: Equal]: Lens[Option[A], A] = Lens.lensg(
  _ => a => a.ifEmpty[Option[A]](none)(some(a)),
  _.orZero
)

This behaves the same except that unsetting the last value in a set removes it from the map:

scala> myLens.set(Map("foo" -> Set(1)), false)
res5: Map[String,Set[Int]] = Map()

This may not be what you want, though.

Answer 2

Vanilla

myMap + (key -> myMap.get(key).fold(Set(newValue))(_ + newValue))

seems easier.

So does writing an extension method, and there it's worth a little extra work to avoid needless reconstruction of anything:

implicit class MapsToSetsCanAdd[K,V](map: Map[K, Set[V]]) {
  def setAdd(key: K, value: V) = map.get(key) match {
    case Some(set) => if (set contains value) map else map + (key -> (set + value))
    case None => map + (key -> Set(value))
  }
}

Now you can merrily myMap setAdd (key, newValue) .