Fun with types: building a type-level Map in Scala

Everyone knows what a Mapis (or a dictionnary, or whatever they are called in your preferred language). It’s a structure that associate values of some type (e.g., a string), the key, to values of a given type (e.g., an integer). Instead of having values as keys, though, could we rather have types as keys, and a value of that type as value? For example, such a map could associate Stringto "hello"and Intto 0. In this blog post, we are going to see how we can do that.

But let us first give some motivation for such enterprise. Let us keep it brief, because if the prospect of having fun with types is not enough to keep you reading, you are probably better spending your time elsewhere. A motivation that comes to mind could be Dependency Injection (DI). If you think about it, DI is precisely a type level map. You ask for some service A to build some other service B, and your DI framework will give you an instance (hopefully the right one) of that type Aat runtime. However, most DI frameworks “cheat” in the sense that they don’t give you a proof at compile time that everything will be wired up properly. We want that.

Before we dive in to the implementation of our type-level map, let’s first build a “usual” map. That is, a map from run-time values of some type to run-time values of another type. However, in order to get a proper parallel with our type-level map, we are going to have a weird-ish implementation (read: far from optimal).

Note: the following assumes Scala 3, as we will use features that are not present in Scala 2.

Computing the types! (By MidJourney)

A run-time map

Let us first build a map from instances of some type K(the key) to instances of some other type V(the value). We will call this map RMap[K, V](“R” for “run-time”). What do we want to do with such thing? At least two things: (1) being able to add a new key-value pair, and (2) given some key, retrieve the value associated to that key. Said in code, we want an implementation for the following trait:

trait RMap[K, V] {
  def add(pair: (K, V)): RMap[K, V] // (1)
  def access(key: K): V             // (2)
}

You can think of many implementations for that (optimized for many different use cases). We are going to implement it in a way that is sub-optimal (atrocious, even) from all run-time performance considerations, but one that, I hope, easen the understanding of our real goal, the type-level map.

We will store the keys in a List, and the values in a Vector, in such a way that the value at some index in the vector is associated to the key at the same index in the list. We thus have something like this:

class RMap[K, V] private (keys: List[K], values: Vector[V]) {
  def add(pair: (K, V)): RMap[K, V] = ???
  def access(key: K): V = ???
}

In order to implement the access method, we need to (1) retrieve the index in the list where the key is and (2) retrieve the value at that index. So we need a (private) method def indexOfKey(key: K): Int, that we are going to implement using tail recursion:

def indexOfKey[K](keys: List[K], key: K): Int = {
  def indexOfKeyAcc(remainingKeys: List[K], currentIndex: Int): Int = 
    remainingKeys match {
      case head :: _ if head == key => currentIndex
      case _ :: tail => indexOfKeyAcc(tail, currentIndex + 1)
      case Nil => throw RuntimeException(
        s"The key $key was not present in Map"
      )
    }
  indexOfKeyAcc(keys, 0)
}
def access(key: K): V = values(indexOfKey(keys, key))

The add method requires us to add the key to the list of keys and the values to the vector of values, at the same index. There are only two reasonable choices: either append or prepend. Usually, we want that newly inserted pairs erase previous ones. Given our implementation of the access method, we thus need to prepend:

def add(pair: (K, V)): RMap[K, V] = 
  RMap(pair._1 +: keys, pair._2 +: values)

Finally, we need something to create an empty map, and we have “everything” we can dream of:

object RMap {
  def empty[K, V]: RMap[K, V] = RMap(List.empty, Vector.empty)
}

Notice what happens when we access a key that is not present in the map: we crash at run-time. This is inevitable to the nature of a run-time map, we can only know what is in there at run-time. Hence, an exception must be raised. This will not be the case with our type-level map.

The type-level map

Our run-time map was parametrized by two types, the type for the keys and the types for the values. The type-level map will be a bit different, it will be parametrized by all the types of all the values contained in it. Previously, the keys where stored in a List of run-time values. In this case, we will need to “store” the keys in a “list” of types. In Scala 3, a list of types is no other than a Tuple. A Tuple is constructed with a type for the head, and another Tuple for the tail (which is exactly the same structure as a List), written like this: Head *: Tail. And we have the EmptyTuple as base type. So, if we have a String and an Int in our map, it will be of type String *: Int *: EmptyTuple. You might sometimes see this kind of construction called a heterogeneous list.

Again, the question is: what do we want exactly? The answer is roughly the same: we want (1) a method to add a new key-value pair and (2) given a type, retrieve the value of that type contained in the map. As a class, this would be

class TMap[TypesOfValues <: Tuple](val values: TypesOfValues) {
  def add[Key](value: Key): TMap[Key *: TypesOfValues] // (1)
  def access[Key]: Key                                 // (2)
}

As we can see, everything that was related to the keys is now encoded in the types. Note that the actual signature of the access method will be more involved, because we need some type-level machinery to make it work. The add method, however, can be readily implemented:

def add[Key](value: Key): TMap[Key *: TypesOfValues] = 
  TMap(value *: values)

Nothing more, nothing less. At this point, we already made the choice that new values will erase previous ones.

What about the access method? If we look at our run-time map implementation, given a type Key, we need a way to find the index of Key in the TypesOfValues tuple. But now, this has to be done at the level of the types! So that would look like

type IndexOfKey[Keys <: Tuple, Key] <: Int

We are again implementing that using some tail recursion mechanism, with a Scala 3 feature called “match types”. So we need our “inner” accumulator function/type, which is

import scala.compiletime.ops.int.+
type IndexOfKeyAcc[RemainingKeys <: Tuple, Key, CurrentIndex <: Int] <: Int = 
  RemainingKeys match {
    case Key *: _ => CurrentIndex
    case _ *: tail => IndexOfKeyAcc[tail, Key, CurrentIndex + 1]
  }

The imported + brings some compile-time arithmetics for Int types. If you are interested in that, I recommand reading this piece which covers that subject beautifully.

And now we have

type IndexOfKey[Keys <: Tuple, Key] = IndexOfKeyAcc[Keys, Key, 0]

The compiler will then be able to compute the type index of our Key in the TypesOfValues tuple, but a type is not that useful. We need the actual run-time value corresponding to the type. Luckily for us, the compiler can also give that to us, if we kindly ask for a given (implicit) ValueOf for that index type. So now the signature of the access method is

def access[Key](
  using indexValue: ValueOf[IndexOfKey[TypesOfValues, Key]]
): Key

The beautiful thing that is not obvious from here is that, if the Key is not present in our map, then accessing that Key will not compile! That is a great guarantee that we get from this, basically for free.

What does the implementation look like, though? From the index, we can now get the value like this:

def access[Key](
  using indexValue: ValueOf[IndexOfKey[TypesOfValues, Key]]
): Key = values.toArray(indexValue.value).asInstanceOf[Key]

(Unfortunately we need this asInstanceOf, but remember that at this point we are safe, because we have the proof that the value at that index is of type Key.) If it’s the first time you are seeing match types and an implicit ValueOf in action, you might be scratching your head wondering what is happening. So here is an example. Let’s say we have a

val map: TMap[String *: Int *: Double *: EmptyTuple]

and we want to retrieve the Int. When you do map.access[Int], the compiler will first compute the type

IndexOfKey[String *: Int *: Double *: EmptyTuple, Int]

and correctly find that it’s the type 1. Then, it sees that we require an implicit ValueOf[1]. And now we are in luck, because there is indeed only one run-time value of type 1, it’s precisely 1! Hence, it can provide us the indexValue that we asked for.

And finally we can have our empty map constructor

object TMap {
  def empty: TMap[EmptyTuple] = TMap(EmptyTuple)
}

We can somewhat test our implementation with the following asserts

assert(TMap.empty.add(3).access[Int] == 3)
assert(TMap.empty.add(3).add("hello").access[Int] == 3)
assert(TMap.empty.add(3).add("hello").access[String] == "hello")
assert(TMap.empty.add(3).add(4).access[Int] == 4)
// TMap.empty.add(3).access[String] // would not compile

So that’s it! We have both the methods that we wanted. Are we really done though? If you start using this thing, you will actually quickly stumble into usability issues that we should fix.

Make it actually usable

We have a problem with the current state of affairs. Let’s say you have a function foo with the following signature

def foo(map: TMap[String *: Int *: EmptyTuple]): String

When you use this function (perhaps you don’t realize this by just reading this post), but you want to be able to give a TMap[Double *: String *: Int *: EmptyTuple] or, even, a TMap[Int *: String *: EmptyTuple]. Because, when you write foo, you simply think “I need a String and an Int”, but you should not care whether the map as argument has “too much” or in a different order. Currently, you simply can’t give foo anything else than precisely a TMap[String *: Int *: EmptyTuple]. Said otherwise, you want to consider the keys of your TMap as a set, rather than an ordered list.

Intuitively, you want to consider that given two tuple types T1 and T2, TMap[T1] would be a subtype of TMap[T2] if and only if all types in T2 are in T1. This would be a covariant-ish functor, if that speaks to you. There is no way to do that though. That is not such a big deal. What we can do instead is to provide an implicit conversion in that case. That way, we could use a map that has at least the types we want, no matter the order.

In order to do that, we need to give the compiler a proof that all types in T2 are indeed in T1. A proof is enough in theory, but in practice, we actually need to build a TMap[T2] from a TMap[T1], by reshuffling and deleting some types. Concretely, we need to find, for every type in T2, the (first) index of that type in T1. And again, we need to do that at the level of the types. But thankfully, it is not that hard. We already have a type that compute the index of some type in a tuple. So, if we consider the tuple T2 as a list, we have to “map” each of these types to the index. The beautiful thing is that there is a built-in type for that, Tuple.Map:

type IndexesOfKeys[SubsetKeys <: Tuple, Keys <: Tuple] = 
  Tuple.Map[SubsetKeys, [key] =>> IndexOfKey[Keys, key]]

We have used a “higher order” function at the level of the types. Quite marvellous, when you think about it. The type IndexesOfKeys will then always be a Tuple, where all its types are (concrete) Ints. Keeping that in mind, we can create a function that will turn an instance of a tuple T1 into an instance of tuple T2:

def transform[From <: Tuple, To <: Tuple](
  from: From, indexes: IndexesOfKeys[To, From]
): To = {
  val fromArray = from.toArray
  val toArray = indexes.toArray
    .map(_.asInstanceOf[Int])
    .map(index => fromArray(index))
  runtime.Tuples.fromArray(toArray).asInstanceOf[To]
}

We used two type casts there. One because (unfortunately) the compiler does not know that all types in the IndexOfKeys tuple are Ints. The other one, again, because the compiler can’t understand that the toArray array indeed represents an array for the tuple To. But again, both these casts are perfectly safe. From the signature of the function, and our definition of the IndexesOfKeys type, if it compiles, it is correct.

The last thing is to say “if the compiler can (implicitly) give me a value for a type IndexesOfKeys for a tuple type From to To, then I can construct a given/implicit conversion from TMap[From] to TMap[To]. You might be tempted to ask for a ValueOf[IndexesOfKeys[To, From]], but unfortunately that does not work. It is a tuple and, even though the compiler can give us a ValueOf for each of the types in that tuple (they are all concrete Ints), it does not want to give us a ValueOf for the tuple. Never fear, though, we can build it ourselves:

case class TupleValueOf[T <: Tuple](value: T)
object TupleValueOf {
  given TupleValueOf[EmptyTuple] = TupleValueOf(EmptyTuple)
  given [Head, Tail <: Tuple](
    using headValue: ValueOf[Head], tailValue: TupleValueOf[Tail]
  ): TupleValueOf[Head *: Tail] = TupleValueOf(
    headValue.value *: tailValue.value
  )
}

The first given is telling the compiler that it can give us a TupleValueOf for the EmptyTuple. The second given instructs that, if it can have the ValueOf for the Head , and the TupleValueOf for the Tail, then it can give us a TupleValueOf for the tuple Head *: Tail. Recursively, that can create the TupleValueOf for any Tuple where we have the ValueOf for each of its types.

Equipped with that, we can now provide our given conversion (in the companion object of TMap)

given [From <: Tuple, To <: Tuple](
  using indexesValue: TupleValueOf[IndexesOfKeys[To, From]]
): Conversion[TMap[From], TMap[To]] = 
  (from: TMap[From]) => TMap[To](transform(from.values, indexesValue.value))

With all of the above, we have a working and practical implementation of our TMap. Several exercises for the interested reader could be (1) add a method to delete a key (don’t forget that the key could appear several times in the TypesOfValues tuple), (2) add a method to concatenate two TMaps (instead of only adding an element) or (3) (hard) create a polymorphic function to map a TMap to another one.

If you want to tinker with that, you can head over to this scastie where all these things are implemented.

Closing words

I hope this was fun! It is not often that you need to do this kind of nonsense. But when you do, it is extremely rewarding and mind blowing to see what the compiler is able to understand. Sometimes, you get a little bit disappointed with the limitations, though. But something too powerful would probably be broken (as much as being unsound).