Mapping Every Key, Part 3

In Part 1, I surveyed common data structures for partial and total maps and settled on lazily-generated tries. In Part 2, I wrote a denotational specification and derived the representation and operations from it. Now, in Part 3, it's finally time to turn all of that into real Haskell code. We'll translate the derived functions into a type class, use GHC's generics to make instances nearly free to write, and then take the result for a spin.

Real Code, At Last

The functions are polytypic, meaning they are defined by induction on types, so type classes will certainly be involved. Since even the representation of a trie depends on the type of the keys, we will need either associated type families or associated data types. Associated data types may be a better choice for defining a "proper" abstraction, since you can hide the representation. Associated type families will be more terse and should allow us to reuse some existing types, along with their existing instances without having to write or derive any wrappers, so here I will use type families.

We start by defining our type class. Since the structure of a trie only depends on the function's input type, never its output type, I am making the associated type family have arity 1, not 2. The upshot is that instances will always define the type in a point-free style with functors. I've decided to name the type constructor (->>) to reflect the math notation from Part 2.

I'm also defining a type synonym Endo, short for endomorphism, so that I can emphasize certain places where Type -> Type can be intuitively thought of as an atomic thing, Endo Type, instead of as a type constructor that needs a type parameter. It won't be immediately obvious why this is nice, but I'll highlight it later when it's more interesting.

{-# LANGUAGE TypeFamilies #-}

infixr 0 ->>
infixl 9 !

type Endo a = a -> a

class TrieKey a where
  type (->>) a :: Endo Type
  
  (!) :: (a ->> b) -> a -> b
  trie :: (a -> b) -> (a ->> b)
  adjust :: a -> (b -> b) -> (a ->> b) -> (a ->> b)

This also seems like a good place to implement replace, since it works for all instances of TrieKey.

replace :: TrieKey a => a -> b -> (a ->> b) -> (a ->> b)
replace a = adjust a . const

It is straightforward to port all of our derived implementations to some basic Haskell types. We will build up the trie representations using the existing Const, Identity, Product, and Compose types from base, but I'll redefine them so that you don't have to go read the documentation. To hopefully make them easier to understand, I'm giving them some kind signatures (types are to values as kinds are to types), using the Endo type synonym we defined earlier to emphasize an intuition that these type constructors return and sometimes transform parameterized types, when partially applied. Some type class instances will be useful later, so we'll derive them here. Also, I've made the unusual decision that the Show instances should not include newtype constructors (conventionally, they should have the same syntax as the source code for defining the values), so that when we view tries in GHCi later the memory representation is easier to see, without syntactic clutter.

{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE StandaloneKindSignatures #-}

type Const :: Type -> Endo Type
newtype Const a b = Const a
  deriving newtype (Show)
  deriving stock (Functor, Foldable, Traversable)

type Identity :: Endo Type
newtype Identity a = Identity a
  deriving newtype (Show)
  deriving stock (Functor, Foldable, Traversable)

type Product :: Endo Type -> Endo Type -> Endo Type
data Product f g a = Pair (f a) (g a)
  deriving stock (Show, Functor, Foldable, Traversable)

type Compose :: Endo Type -> Endo Type -> Endo Type
newtype Compose f g a = Compose (f (g a))
  deriving newtype (Show)
  deriving stock (Functor, Foldable, Traversable)

I'll go ahead and redefine Void, for the same reason that it's more convenient to repeat the definition here than to make you go look at the documentation.

type Void :: Type
data Void

However, I'll assume you already know, or can reasonably guess, anything that's in Prelude, that is, anything in scope by default in Haskell.

Let's implement the instance for the empty type, Void, first! Here's the implementation we came up with in Part 2.

$$
\begin{align}

\mathbf{0} \twoheadrightarrow X & = \mathbf{1}
\\
{\left[\!\left[ \star \right]\!\right]}_{\mathbf{0}} &= \text{absurd}
\\
{\mathrm{trie}}_{\mathbf{0}}(\text{absurd}) &= \star
\\
\mathrm{adjust}_{\mathbf{0}}(a,h,\star) &= \star

\end{align}
$$

We can use the EmptyCase language extension to eliminate a Void. The trie is Const (), meaning that, no matter what, it's just () at runtime.

{-# LANGUAGE EmptyCase #-}

instance TrieKey Void where
  type (->>) Void = Const ()
  Const () ! void = case void of {}
  trie _ = Const ()
  adjust _ _ (Const ()) = Const ()

Next is the unit type.

$$
\begin{align}

\mathbf{1} \twoheadrightarrow X & = X
\\
{\left[\!\left[ x \right]\!\right]}_{\mathbf{1}} &= \lambda \star.\; x
\\
\mathrm{trie}_{\mathbf{1}}(f) &= f(\star)
\\
\mathrm{adjust}_{\mathbf{1}}(\star, h, x) &= h(x)

\end{align}
$$

The instance for () uses Identity to store exactly one output value.

instance TrieKey () where
  type (->>) () = Identity

  Identity a ! () = a
  trie f = Identity (f ())
  adjust () f (Identity a) = Identity (f a)

Tries for sum types are products of tries.

$$
\begin{align}

A + B \twoheadrightarrow X & = (A \twoheadrightarrow X) \times (B \twoheadrightarrow X)
\\
\\
{\left[\!\left[ (t_1,t_2) \right]\!\right]}_{A + B} &= \lambda
\begin{cases}
\mathrm{inj}_1(a) \mapsto {\left[\!\left[ t_1 \right]\!\right]}_A(a) \\
\mathrm{inj}_2(b) \mapsto {\left[\!\left[ t_2 \right]\!\right]}_B(b)
\end{cases}
\\
\\
{\mathrm{trie}}_{A + B}(f) &= (\mathrm{trie}_A(f \circ \mathrm{inj}_1), \mathrm{trie}_B(f \circ \mathrm{inj}_2))
\\
\\
\mathrm{adjust}_{A+B}(\mathrm{inj}_1(a), h, (t_1, t_2)) &= (\mathrm{adjust}_A(a, h, t_1),\; t_2)
\\
\mathrm{adjust}_{A+B}(\mathrm{inj}_2(b), h, (t_1, t_2)) &= (t_1,\; \mathrm{adjust}_B(b, h, t_2))

\end{align}
$$

We can use Product to pair them up. The functions port directly from the derived implementations.

instance (TrieKey a, TrieKey b) => TrieKey (Either a b) where
  type (->>) (Either a b) = Product ((->>) a) ((->>) b)

  Pair t1 _ ! Left a = t1 ! a
  Pair _ t2 ! Right b = t2 ! b

  trie f = Pair (trie (f . Left)) (trie (f . Right))

  adjust (Left a) f (Pair t1 t2) = Pair (adjust a f t1) t2
  adjust (Right b) f (Pair t1 t2) = Pair t1 (adjust b f t2)

Tries for product types are compositions of tries.

$$
\begin{align}

A \times B \twoheadrightarrow X & = A \twoheadrightarrow (B \twoheadrightarrow X)
\\
{\left[\!\left[ t \right]\!\right]}_{A \times B} &= \lambda (a,b).\; {\left[\!\left[ {\left[\!\left[ t \right]\!\right]}_{A}(a) \right]\!\right]}_{B}(b)
\\
\mathrm{trie}_{A \times B}(f) &= {\mathrm{trie}}_A(\lambda a.\; {\mathrm{trie}}_B(\lambda b.\; f(a, b)))
\\
\mathrm{adjust}_{A \times B}((a,b), h, t) &= \mathrm{adjust}_A(a, \;\lambda u.\; \mathrm{adjust}_B(b, h, u),\; t)

\end{align}
$$

We use Compose for that. Once again, the derived functions port to Haskell directly.

instance (TrieKey a, TrieKey b) => TrieKey (a, b) where
  type (->>) (a, b) = Compose ((->>) a) ((->>) b)

  Compose t ! (a, b) = t ! a ! b
  trie f = Compose (trie (\a -> trie (\b -> f (a, b))))
  adjust (a, b) f (Compose t) = Compose (adjust a (adjust b f) t)

The representation for a trie of functions is a trie of tries.

$$
\begin{align}

(A \to B) \twoheadrightarrow X & = (A \twoheadrightarrow B) \twoheadrightarrow X
\\
{\left[\!\left[ t \right]\!\right]}_{A \to B} &= \lambda f.\; {\left[\!\left[ t \right]\!\right]}_{A \twoheadrightarrow B}(\mathrm{trie}(f))
\\
\mathrm{trie}_{A \to B}(f) &= \mathrm{trie}_{A \twoheadrightarrow B}(\lambda t.\; f({\left[\!\left[ t \right]\!\right]}_A))
\\
\mathrm{adjust}_{A \to B}(f, h, t) &= \mathrm{adjust}_{A \twoheadrightarrow B}(\mathrm{trie}_A(f), h, t)

\end{align}
$$

Nesting associated type families in this way requires UndecidableInstances. The TrieKey (a ->> b) constraint requires FlexibleContexts. We also need a TrieKey a constraint in order to convert between a -> b and a ->> b.

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}

instance (TrieKey a, TrieKey (a ->> b)) => TrieKey (a -> b) where
  type (->>) (a -> b) = (->>) (a ->> b)

  t ! g = t ! trie g
  trie f = trie (f . (!))
  adjust = adjust . trie

Making Instances Automatically Derivable

We could continue manually writing instances by hand for every type we want to support, but I'd like to avoid writing so much code. Instead, let's use GHC's generics machinery to provide default definitions for everything.

A Brief Overview of Generics

Generics works by representing algebraic data types by composing various types defined in GHC.Generics, which you can write type class instances for to get instances for other types for free. The neat trick is that GHC can automatically derive these generic representations for your data types and convert back and forth between your types and the generic ones, via this type class:

class Generic a where
  type Rep a :: Type -> Type
  from :: a -> (Rep a) x
  to :: (Rep a) x -> a

The Generics language extension allows you to simply write deriving Generic to define instances of Generic for your own types. Here are the building blocks that it uses for Rep:

-- V1 is like (Const Void)
data V1 p

-- U1 is like (Const ())
data U1 p

-- (:+:) is like Data.Functor.Sum
data (f :+: g) p = L1 (f p) | R1 (g p)

-- (:*:) is like Product
data (f :*: g) p = f p :*: g p

-- (K1 i) is like Const
newtype K1 i c p = K1 c

-- (M1 i t) is like (Compose Identity)
newtype M1 i t f p = M1 (f p)

The p parameters are used when representing types of kind k -> Type, but we won't be doing that today, so when we use Rep, we'll just set p to (). The i, t, and f parameters of K1 and M1 are metadata (like field names, constructor names, strictness annotations, etc.), provided at type level, which we will also not need today.

Let's start by writing the instances. Here are the ones that closely resemble instances we've already written.

{-# LANGUAGE TypeOperators #-}

instance TrieKey (V1 p) where
  type (->>) (V1 p) = Const ()
  Const () ! v1 = case v1 of {}
  trie _ = Const ()
  adjust _ _ (Const ()) = Const ()

instance TrieKey (U1 p) where
  type (->>) (U1 p) = Identity

  Identity a ! U1 = a
  trie f = Identity (f U1)
  adjust U1 f (Identity a) = Identity (f a)

instance (TrieKey (f p), TrieKey (g p)) => TrieKey ((f :+: g) p) where
  type (->>) ((f :+: g) p) = Product ((->>) (f p)) ((->>) (g p))

  Pair t1 _ ! L1 a = t1 ! a
  Pair _ t2 ! R1 b = t2 ! b

  trie f = Pair (trie (f . L1)) (trie (f . R1))

  adjust (L1 a) f (Pair t1 t2) = Pair (adjust a f t1) t2
  adjust (R1 b) f (Pair t1 t2) = Pair t1 (adjust b f t2)

instance (TrieKey (f p), TrieKey (g p)) => TrieKey ((f :*: g) p) where
  type (->>) ((f :*: g) p) = Compose ((->>) (f p)) ((->>) (g p))

  Compose t ! (fp :*: gp) = t ! fp ! gp
  trie f = Compose (trie (\fp -> trie (\gp -> f (fp :*: gp))))
  adjust (fp :*: gp) f (Compose t) = Compose (adjust fp (adjust gp f) t)

M1 is used exclusively for annotating types with metadata which, as I said above, we don't need, so we write a simple pass-through instance.

instance (TrieKey (f p)) => TrieKey (M1 i t f p) where
  type (->>) (M1 i t f p) = (->>) (f p)

  t ! M1 fp = t ! fp
  trie f = trie (f . M1)
  adjust (M1 fp) = adjust fp

K1 is more interesting. It's used to wrap the types of constructor fields (as opposed to their generic representations) so that you can use bespoke instances for them instead of necessarily generic ones. Also, without K1, recursive data types would have infinitely large generic representations; K1 functions as a recursion breaker. This latter purpose is highly relevant to us. If we aren't careful to break recursion in (->>), recursive key types would create infinitely large types. The standard Haskell way of breaking recursion is to introduce a newtype. Unfortunately, although K1's i parameter is set to R to indicate recursion, it's never set to anything else, so we can't distinguish between recursion and no recursion. The sad consequence is that the derived (->>) type is likely to be polluted with the recursion breaker newtype even when it's not necessary. If we don't use a recursion breaker, we just won't be able to derive instances for recursive key types. I've decided not to use a recursion breaker for the blog post, to simplify presentation. That means we won't be able to derive TrieKey instances for recursive types. You can still define them by hand, at least.

instance TrieKey c => TrieKey (K1 i c p) where
  type (->>) (K1 i c p) = (->>) c

  t ! K1 c = t ! c
  trie f = trie (f . K1)
  adjust (K1 c) = adjust c

Now all we have to do is add some default definitions to TrieKey. I'm also defining some auxiliary things, including a constraint synonym GenericTrieKey that gets used a few times and toU and fromU functions to specialize to and from to Rep a (), avoiding some type annotations elsewhere.

{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DefaultSignatures #-}

toU :: (Generic a) => Rep a () -> a
toU = to

fromU :: (Generic a) => a -> Rep a ()
fromU = from

type GenericTrieKey a = (Generic a, (->>) a ~ (->>) (Rep a ()), TrieKey (Rep a ()))

class TrieKey a where
  type (->>) a :: Endo Type

  (!) :: (a ->> b) -> a -> b
  trie :: (a -> b) -> (a ->> b)
  adjust :: a -> (b -> b) -> (a ->> b) -> (a ->> b)

  type (->>) a = (->>) (Rep a ())

  default trie :: (GenericTrieKey a) => (a -> b) -> (a ->> b)
  trie f = trie (f . toU)

  default (!) :: (GenericTrieKey a) => (a ->> b) -> a -> b
  t ! a = t ! fromU a

  default adjust ::
    GenericTrieKey a => a -> (b -> b) -> (a ->> b) -> (a ->> b)
  adjust = adjust . fromU

The constraints factored out by GenericTrieKey are worth highlighting. Generic a just requires that a have a generic representation. TrieKey (Rep a ()) requires that the generic representation has the required instances. (->>) a ~ (->>) (Rep a ()) means that (->>) a must be the same as (->>) (Rep a ()), which should always be true unless you override the default definition of (->>) a.

Now the payoff! We can define instances much more easily now. Here are some examples:

{-# LANGUAGE DeriveGeneric #-}

instance TrieKey Bool
instance TrieKey a => TrieKey (Maybe a)
instance (TrieKey a, TrieKey b, TrieKey c) => TrieKey (a, b, c)

data Vec2 a = Vec2 a a deriving stock Generic
instance TrieKey a => TrieKey (Vec2 a)

Demo

Derived Trie Types

But what do the instances even look like? I've expanded the derived trie types from above to demonstrate. I've written some type constructors using infix syntax where I figured it would help with readability.

type (->>) Bool = Identity `Product` Identity
type (->>) (Maybe a) = Identity `Product` (->>) a
type (->>) (a, b, c) = (->>) a `Compose` ((->>) b `Compose` (->>) c)
type (->>) (Vec2 a) = (->>) a `Compose` (->>) a

Effectively:

Bool ->> x = (x, x)
Maybe a ->> x = (x, a ->> x)
(a, b, c) ->> x = a ->> b ->> c ->> x
Vec2 a ->> x = a ->> a ->> x

Let's have a look at a non-trivial, concrete example, (->>) (Bool, Bool, Bool). Here's what it effectively expands to:

  (Bool, Bool, Bool) ->> x
= Bool ->> Bool ->> Bool ->> x
= Bool ->> Bool ->> (x, x)
= Bool ->> ((x, x), (x, x))
= (((x, x), (x, x)), ((x, x), (x, x)))

It's a complete binary tree of depth 3. Each (->>) Bool forms a branch. Let's visualize the value trie id :: (Bool, Bool, Bool) ->> (Bool, Bool, Bool). First, here it is straight from GHCi, using the unconventional Show instances I defined earlier:

> trie id :: (Bool, Bool, Bool) ->> (Bool, Bool, Bool)
Pair
  ( Pair
      ( Pair
          (False, False, False)
          (False, False, True)
      )
      ( Pair
          (False, True, False)
          (False, True, True)
      )
  )
  ( Pair
      ( Pair
          (True, False, False)
          (True, False, True)
      )
      ( Pair
          (True, True, False)
          (True, True, True)
      )
  )

Here it is in graphical form, with some labels indicating what each path means:

Some Bespoke Tries

Now that we have this generic machinery, we can use it to implement instances that would be more of a pain otherwise. For example, we can derive an instance for 8-tuples and then define type (->>) Word8 = (->>) (Bool, Bool, Bool, Bool, Bool, Bool, Bool, Bool), using the isomorphism between Word8 and 8-tuples of Bools to implement the functions.

In my actual implementation, I decided to define a Nibble type (4-bit words) and manually implement TrieKey Nibble so that a nibble trie is a flat array of 16 elements.

data Nibble
  = N0 | N1 | N2 | N3 | N4 | N5 | N6 | N7
  | N8 | N9 | NA | NB | NC | ND | NE | NF

data Vec16 a = Vec16
  { x0, x1, x2, x3, x4, x5, x6, x7
  , x8, x9, xA, xB, xC, xD, xE, xF :: a}

instance TrieKey Nibble where
  type (->>) Nibble = Vec16

Using this as a building block, I'm able to define tries for the Word types (I'm only showing the types here, but of course there must be a proper instance for each one):

type (->>) Word8 = (->>) (Nibble, Nibble)   -- depth  2 instead of  8
type (->>) Word16 = (->>) (Word8, Word8)    -- depth  4 instead of 16
type (->>) Word32 = (->>) (Word16, Word16)  -- depth  8 instead of 32
type (->>) Word64 = (->>) (Word32, Word32)  -- depth 16 instead of 64

Functionality For Free

Unlike plain functions, tries are foldable and traversable! For example, the Traversable instances of the building blocks mean that tries themselves are traversable, allowing us to map over all of the values in a trie with IO.

> trie (id :: Bool -> Bool)
Pair False True
> traverse (\x -> print x >> randomIO @Int) it
False
True
Pair (-5863997299141589325) (-5801038465154072764)

Basic Trie Operations

It would be a bit silly of me not to at least demonstrate lookup and replacement. Let's first define a trie:

> let t = trie id :: Maybe Bool ->> Maybe Bool
> t
Pair Nothing (Pair (Just False) (Just True))

Let's demonstrate indexing:

> t ! Nothing
Nothing
> t ! Just False
Just False
> t ! Just True
Just True

Here's replacing, with another indexing demo just to prove that new value is mapped from the original key:

> let t' = replace (Just False) Nothing t
> t'
Pair Nothing (Pair Nothing (Just True))

> t' ! Just False
Nothing

Finally, a simple demo of adjust that isn't replacing the value with a mere constant like replace does:

> adjust (Just True) (fmap not) t'
Pair Nothing (Pair Nothing (Just False))

Laziness

The fact that these tries are lazy is very important when it comes to large key types, like Word64. If Word64 ->> a was evaluated strictly, there would be \(2^{64}\) values of type a manifested in memory, each a leaf in a huge tree, definitely too much for any real-world computer to handle. And yet, I can easily do this:

> let big = trie (show :: Word64 -> String)
> let [k1,k2] :: [Word64] = [42, 1337]
> big ! k1
"42"
> big ! k2
"1337"
> let big' = replace k1 "K1 is no more" big
> big' ! k1
"K1 is no more"
> big' ! k2
"1337"

To demonstrate what big looks like in memory, I ran :print big. However, the output was very large. I did some manual cleanup to make it more terse, and here's what I ended up with:

big =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16
    { x0 = Vec16 { x2 = Vec16 { xA = "42" } }
    , x5 = Vec16 { x3 = Vec16 { x9 = "1337" } }
    } } } } } } } } } } } } } }

I'm using record syntax to render them as partially-defined values. You can see that only the paths to the keys 42 (hexadecimal 0x000000000000002A, which is exactly the path to "42") and 1337 (hexadecimal 0x0000000000000539, also exactly the path to "1337") are evaluated. And here is the similarly cleaned-up output of :print big':

big' =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16 { x0 =
  Vec16
    { x0 = Vec16 { x2 = Vec16 { xA = "K1 is no more" } }
    , x5 = Vec16 { x3 = Vec16 { x9 = "1337" } }
    } } } } } } } } } } } } } }

Pretty slick.

Conclusion

It was quite a journey, but we're done! We went from a survey of data structures, through a formal derivation, to a working Haskell implementation with generic instance deriving. The tries we've built are total, purely functional, foldable, and traversable (and a whole bunch more in the real implementation, since the standard individual building blocks offer a lot more than that), and they're generated incrementally, on demand rather than all at once, which is what made this whole approach viable in the first place.