Merkle Tries

I recently took a take-home programming test with a potential employer. While no job offer materialized from the experience, since I enjoyed working through the problems so much, I decided to push myself further and see what kind of solution I could come up with, when no restrictions were imposed.

The current blog post is my commentary on how I completed that exercise, but if you want, you can skip all of commentary and get straight to the dang code by clicking here.

So what is this all about? The take-home test in question was split into a few separate parts, one of which was to develop a Merkle trie library from scratch (in Rust, with LLM assistance permitted). While I had used/seen a few trie libraries previously, I had never written one on my own. As I embarked on this new adventure, I found there were a couple different design decisions that I had to hold in tension. I’ve listed a few of the more interesting ones below:

  1. How can we make APIs that are efficient and respect Rust’s owernship discipline?
  2. How can we write that safely and generically handles operations over Merkle tries as well as their witnesses?
  3. Should we optimize for dense tries or sparse tries?

In the remainder of this post, I’ll tackle each of these questions in turn.

NB: Since a trie is a kind of tree, I prefer to use term ’tree’ to talk about considerations that apply to all trees and ’trie’ for trie-specific issues.

Interestingly, I found that the first two questions required more careful thinking than the last one, so let’s start with it:

  1. Should we optimize for dense tries or sparse tries?

In my estimation, the above question is basically the same as the following question: should we optimize for a higher or lower branching factor? To see why, first note that:

Now, for sparse trees, we know most branches will be empty. This implies, probabilistically, that the tree will be shallower. This means that, even with a lower branching factor, our tree will likely not grow too large. Altogether, this means that a low branching factor gives us free memory and and likely better memory locality, as more nodes can fit in any given cahce level, page, etc…

For dense trees, the opposite is true: most branches will be full. In that case, a lower branching factor will cause the tree depth to balloon, leading to redundant allocations (e.g. a single complete 4-ary node is equivalent to three complete binary nodes) and pointer chasing. Altogether, this means that a higher branch factor saves us time and also will likely save memory because branch node metadata can be shared amongst multiple child nodes.

In any event, I punted on figuring out if some branching factor is better in general (whatever that would mean) and used generics to let the user pick from a few different branching factors (either 2, 4, 16, or 256). The decision was made to restrict to powers of two, since:

Put altogether, this means a child node lookup from a parent branch node equates to exactly the following four steps:

  1. Grab the appropriate selector byte from the target key;
  2. Apply a branch-stored bitmask to the selector byte to remove all bits that don’t correspond to the current branch node’s child nodes;
  3. Downshift the masked byte to remove all trailing zeroes;
  4. Use the resulting 1-byte number as an index into the child node array.

Now, let’s come back to the first question:

  1. How can we make APIs that are efficient and respect Rust’s owernship discipline?

Originally, to answer this question, I looked to the std::collections BTreeMap and HashMap entry-style APIs. These APIs are nice in that they essentially return a wrapped pointer to a map’s key-value pair. This pointer lets you perform all of the operations you want at once.

This API can optimize common cases where, for example, you need to set the value of a key only if it hasn’t been set before. With a standard map API, you need to split this operation into two parts:

  1. Check if the key is in the map (requires either O(log(N)) or amortized O(1), depending on the implementation);
  2. If the key isn’t in the map, insert it (requires another O(log(N)) or amortized O(1)).

The key point (pun intended) is that, in the base case, the map structure has to be examined twice, identically, to perform this operation when, clearly, only one map traversal is actually required which just needs to return a pointer to where the key should be inserted. From this information alone, the key’s presence in the map can be checked and a new value can be inserted, if required.

However, the standard entry-style APIs require non-trivial work to be transferred to merkleized data structures for a simple reason: mutation in a merkleized data structure is inherently non-local! In particular, any mutation at some node N invalidates any cached node digests for all of N’s ancestor nodes; this leaves two options:

  1. make tree mutations and ancestor digest updates atomic;
  2. make ancestor digest updates lazy.

In approach (1), an API similar to the entry-style API can be supported—but it requires an inversion of control flow. Instead of returning a wrapped pointer entry-object, the update method must take an updater object (basically, a closure) which contains all logic needed to update the map. It then invokes the closure to update the map locally and walks back up the map structure to ensure that all ancestor node hashes are recalculated.

In approach (2), the entry API can be directly supported. However, the cost is that all cached ancestor node digests must be invalidated. Then, when a digest is requested, the digest function checks whether the current node’s cached digest has been invalidated. If so, it walks down the trie structure and back up, looking for invalided child node digests, recalculating them, and then recalculating their ancestors’ digests.

In this implementation, I chose approach (1). My reasoning was two-fold:

In actuality, the updater object is not a closure, but a trait, to avoid wasted memory; see the code for full details.

Finally, we’ll tackle the second question:

  1. How can we write that safely and generically handles operations over Merkle tries as well as their witnesses?

This one was a bit tricky because I wanted two things that were at odds:

  1. (essentially) the same API and same code for both complete Merkle tries and partial Merkle tries (that is, witnesses);
  2. a way to programmatically forbid a Merkle trie desiginated as “complete” to contain pruned sub-tries.

Point (2) prevents stupid logic errors and allows a clean separation of concerns that immediately gives the following safety property: when working with a complete trie, the entire structure is available and all opertions are total.

To accomplish this, I did the following:

  1. Used an extra generic trait parameter on my Trie struct that attached to potentially pruned nodes;
  2. Instantiated this trait parameter with either the unit type (a.k.a. the 0-tuple () in Rust) or the std::convert::Infallible type, i.e., an alias for the empty type (called never in Rust parlance and written !).

In complete mode, the trait parameter is the empty type, which implies that pruned nodes cannot be constructed (as no valid term has this type).

The one last trick is that: since both the unit and empty types are zero-sized types with identical alignment, if we are careful, we can safely bitcast from a “complete” trie to a “partial” trie, as every “complete” trie is a “partial” trie where no nodes have been pruned yet.

Aside from these main questions, before closing this post out, I’ll detail one other issue that I encountered.

The issue: how to best DRY my code (does anyone use this as a verb?).

In particular, Rust’s standout feature, the borrow checker which tracks mutable (&mut T) and immutable (&T) reference type usage, can sometimes fight against DRY. This can happen due when some bit of functionality f is generic over reference mutability, but some other functions use_mut and use both want to use f in different mutability contexts.

A very good example of this data structure traversal: whether one is reading or updating a key in a tree, one will first need to traverse the tree to find the appropriate node. If we had something like built-in mutability-generics, then tree traversal could be written as a single function that:

However, as far as I can tell, there is:

  1. no way to write this directly;
  2. no canonical trait that can capture this pattern.

As I was building this merkle trie library, for the specific example above, i.e. tree traversal, I eventually gave up and just duplicated the code (for mutable/immutable cases respesctively).

For a deeper discussion of this implementation detail and others, see the implementation details section of the library README.

In addition to being an interesting problem in its own right, this was also my most long-running experiment to-date in using LLMs for code assistance.

This came naturally: on the initial take-home test, their use was encouraged and, furthermore, the development of a merkle trie library was merely a building block to be used in a later test question.

When I decided to continue the project on my own, I also decided I would continue my LLM experimentation (but in a more relaxed sort of way) as I no longer had any time constraints.

This post is too long to have a complete discussion of my thoughts on LLMs, but I hope to return to this topic in a future post.

Finally, if you’ve made it this far, I hope you found the above discussion interesting! As always, the code is freely available, and you can read more and/or use it for yourself by visiting https://github.com/sskeirik/merkle-trie.