Can I get single-lookup default map behaviour in Kotlin?

Question

Coming from C++, Kotlin's withDefault feature doesn't quite work as expected:

val m = mutableMapOf<Int, MutableList<Int>>().withDefault { mutableListOf() }
...
m.getValue(1).add(2) // does not work if m[1] == null!

Now if for performance reasons I want to do only a single lookup, how can I do that?

A way with two lookups would be

if (m[1] == null) m[1] = mutableListOf(2)
else m[1].add(2)

Is there any way with a single lookup?

Answer 1

// Create the map with some example entries
val m = mutableMapOf(
  1 to mutableListOf(),
  2 to mutableListOf(1, 2, 3)
)

m.putIfAbsent(2, mutableListOf(789))?.add(5)
// Will append 5 to [1, 2, 3] as the key 2 exists

m.putIfAbsent(3, mutableListOf(789))?.add(5)
// Will create a new map entry 3 to [789] as the key 3 does not exist

m.toList().forEach(::println)

Output:

(1, [])
(2, [1, 2, 3, 5])
(3, [789])

Answer 2

TL;DR: Skip to the end of this answer if you just want to see a possible way to implement what you want.

---

Why does withDefault not work as you expected?

The [Mutable]Map.withDefault and [Mutable]Map.getValue extension functions seem designed for two things:

1. A way to define a default value at the declaration site rather than the use site. In other words, the following:

   val map = mutableMapOf<String, String>().withDefault { "default" }
   val foo = map.getValue("foo")
   

   Is an alternative for doing:

   val map = mutableMapOf<String, String>()
   val foo = map["foo"] ?: "default"
   

2. To be used as a property delegate where the delegated property must have a non-null value regardless of there being a corresponding entry in the map or not.

   // It's possible to do the 'withDefault' here instead of for each
   // delegated property, even if the default value depends on the
   // key (the default value function takes the key as an argument).
   val map = mutableMapOf<String, String>()
   
   // This assumes 'map' is not already a 'withDefault' decorated
   // map, otherwise calling 'withDefault' again just overrides the
   // default value function. 
   var foo by map.withDefault { "foo-default" }
   var bar by map.withDefault { "bar-default" }
   

   Without withDefault it would be possible to query foo and get a relatively unexpected exception.

All without side-effects. Why no side-effects? Well, for one, a Map after calling Map.withDefault has to remain read-only. But why not allow side-effects when you have a MutableMap and call MutableMap.withDefault? Not entirely sure. My guess is to maintain consistent behavior and to avoid code being able to "modify" a Map directly. For an example of the latter point, if you have:

class Foo {

    private val _map = mutableMapOf<String, Sring>().withDefault { "default" }
    val map: Map<String, String> get() = _map
}

The intent is to let instances of Foo mutate the map, but any outside code that gets the map via the map property will only see the read-only Map interface. Yet if withDefault had the side-effect of inserting the default value, then doing map.getValue("key") is effectively the same as doing map["key"] = "default" as if map was a MutableMap instead of just a Map.

Also, note that withDefault only works in conjunction with getValue. The default value function will not be used if you call get or use [] and the map doesn't contain the key.

---

A Solution

That said, it seems you do want the side-effect of the default value to being inserted into the map when no entry exists for the given key. I'm not aware of any built-in function for this. However, you can create your own extension function and map decorator to do what you want:

// possibly needs a better function name
fun <K, V> MutableMap<K, V>.ifAbsent(computeValue: (key: K) -> V): MutableMap<K, V> {
    return object : MutableMap<K, V> by this {

        override operator fun get(key: K): V? {
            val value = this@ifAbsent[key]
            return if (value != null || containsKey(key)) {
                value
            } else {
                computeValue(key).also { this[key] = it }
            }
        }

        override fun getOrDefault(key: K, defaultValue: V): V {
            return [email protected](key, defaultValue)
        }

        override fun equals(other: Any?) = [email protected](other)
        override fun hashCode() = [email protected]()
        override fun toString() = [email protected]()
    }
}

Which can be used like so:

fun main() {
    val map = mutableMapOf<Int, MutableList<Int>>().ifAbsent { mutableListOf() }

    map.getValue(2).add(42)
    // or: map[2]!!.add(42)

    println(map[2]) // prints "[42]"
}

JVM simplification

If you're only targeting the JVM or Android, then you can simplify the get function to just:

override operator fun get(key: K): V? {
    return [email protected](key, computeValue)
}