From map[T]struct{} to Red-Black Trees: A Deep Dive into go-set (and a bug I fixed)

I was reading through hashicorp/go-set recently, which is a small Go library used by the Nomad team and spotted a // TODO optimize comment sitting in a method called EqualSliceSet. That comment led me down a rabbit hole of Red-Black trees, lock-step iteration. This post documents the whole journey.

1. What is go-set and why does it exist?

If you have written Go for any length of time, you have almost certainly written this:

It works. It is also noisy, repetitive, and easy to get wrong. Every project that needs a set ends up with a slightly different version of that pattern. go-set exists to give you a clean, type-safe, generic set API so you stop writing boilerplate and start writing with intent.

Much better. But the library does not stop at one implementation, it ships three, each with a different trade-off.

The three implementations

• Set[T] is the basic one. It's backed by a Go map and gives you O(1) membership checks. You can use it when your type is comparable, that is with, strings, ints, simple structs and you don't care about order.
• HashSet[T] is for when your type isn't comparable but you can define a hash function for it. Still backed by a map under the hood, still unordered, but works with complex structs.
• TreeSet[T] is where it gets interesting. It's backed by a Red-Black Binary Search Tree and keeps elements in sorted order. You'd reach for it when you need ordered iteration, range queries, or operations like Min(), Max(), TopK().

Set and HashSet are O(1) average for insert/contains/delete. TreeSet is O(log n) for those same operations, but it buys you something the other two cannot give: the elements always come out in a defined order.

2. Zooming into `TreeSet` and what makes it different?

The headline feature is CompareFunc[T]:

You pass this in when you create the set. It must return negative if a < b, zero if equal, positive if a > b. For built-in ordered types this is just cmp.Compare:

But because it is generic, you can sort by anything:

What you get on top of the basic set operations

Because the tree is always sorted, the library can give you operations that are meaningless on an unordered set:

All of these are efficient because the data is already in a sorted structure and you are not scanning a hash map.

3. What is a Red-Black Tree and why does `TreeSet` use one?

Start with a plain Binary Search Tree

A Binary Search Tree (BST) has one rule: for any node, everything in its left subtree is smaller and everything in its right subtree is larger.

BST Example

Searching, inserting, and deleting are all O(h) where h is the height of the tree. In the tree above, h = 2, and n = 7, so h = log₂(n) (approx.). That is the ideal case.

The problem is that a plain BST can degenerate. If you insert elements in sorted order:

Degenerate BST

Now it is a linked list. h = n, search is O(n). Terrible.

Red-Black Trees; stay balanced always

A Red-Black Tree is a BST with four extra invariants enforced on every insert and delete:

1. Every node is either red or black.
2. The root is always black.
3. Red nodes may not have red children (no two reds in a row).
4. Every path from root to a null leaf contains the same number of black nodes.

These rules together guarantee that the longest path (alternating red-black) is at most twice the shortest path (all black). That caps the height at 2 log₂(n), which means O(log n) is guaranteed and not just average-case.

Let's build one from the same insertions to see what the tree actually does:

Insert 1, 2, 3:

After inserting 1 (root, turns black), 2 (red, right child), 3 (red triggers rotation):

Red Black Tree

The tree rotated left and recoloured. It will never become a linked list no matter what order you insert in.

A fuller example; insert 1 through 7:

Insertion in Red Black Tree

Seven elements, maximum height 3. That is the balance guarantee in action.

Why `TreeSet` specifically uses it

Hash maps give O(1) but no order. Sorted slices give order but O(n) insertions. A Red-Black Tree gives O(log n) for everything -> insert, delete, search, min, max, range queries, while keeping elements in sorted order at all times. For a data structure that needs to be simultaneously a set and a sorted collection, it is the right choice.

4. Walking through the `TreeSet` code

Insertion

Every new node starts red. That is by design adding a red node does not change the black-height of any path (invariant 4 is untouched). It might violate invariant 3 (two reds in a row), which is then fixed by rebalanceInsertion.

rebalanceInsertion works bottom-up, checking the uncle node's colour to decide between a simple recolour or a rotation:

• uncle is red -> recolour parent, grandparent, uncle; recurse up
• uncle is black, zig-zag shape -> rotate to straighten; then fall through
• uncle is black, straight line -> rotate at grandparent; recolour

All rotations are O(1) pointer swaps. The total work per insert is O(log n), which is, one walk down to find the spot, at most O(log n) rebalance steps walking back up.

How `iterate()` gives you elements in ascending order

In-order traversal (left -> node -> right) of any BST yields elements in ascending order. TreeSet does this lazily using an explicit stack:

Visualising with the tree {1, 2, 3, 4, 5, 6, 7}:

Visualization

Each call to the returned closure costs O(1) amortised. The total iteration of n elements is O(n).

How `CompareFunc` ties everything together

The same CompareFunc[T] stored in s.comparison is used in every operation: insertion, deletion, search, and the iterate() result order. This is what makes the lock-step comparison in the optimization work. Both the tree traversal and the sorted slice use the identical ordering function, so they produce the same sequence.

5. The bug I found -> EqualSliceSet

While reading through treeset.go I found this:

The // TODO optimize is a hint from the original authors that they knew something better was possible but left it for later. Let's follow the call chain to see why.

EqualSlice

TreeSetFrom

What actually happens for `EqualSliceSet([]int{3, 1, 2})`

EqualSliceSet Example

For a slice of n elements:

• n heap allocations -> one node[T] struct per element, each with 5 fields (element, color, parent, left, right).
• O(n log n) comparisons just to build the tree.
• Up to O(n log n) rotation/recolour operations to maintain Red-Black invariants.
• The entire tree structure is thrown away as soon as Equal returns.
• GC pressure -> n newly allocated objects that immediately become garbage.

And all of this is happening inside a method whose only job is to answer: "do these two things contain the same values?" There has to be a better way.

6. The fix -> sort and walk in lock-step

The key insight

TreeSet always iterates in ascending order via s.iterate(). If we sort items using the same CompareFunc[T], it will also be in ascending order. Two sorted sequences represent the same set if and only if they are element-wise identical. So we can compare them in a single O(n) pass, no tree needed.

The implementation

Why the defensive copy matters

slices.SortFunc sorts in place. If we sorted items directly, the caller's slice would come back in a different order. A silent mutation they never asked for. make + copy gives us our own backing array. The test "does not mutate caller slice" enforces this:

Visualising the lock-step comparison

For set {1, 2, 3, 4} and items [4, 2, 1, 3]:

Visualising lock-step

For set {1, 2, 3, 4} and items [4, 2, 9, 3]:

Visualising lock-step 2

Handling duplicates; the subtle case

Obvious case: items = [1, 1, 2, 3], set = {1, 2, 3}. Lengths differ (4 vs 3) -> the size check catches it before we even sort.

Subtle case: items = [1, 1, 2, 4], set = {1, 2, 3, 4}. Lengths are both 4. The size check passes. After sort:

After Sort

The duplicate 1 pushed everything right in the sorted slice. The tree's second element is 2, the slice's second element is still 1, results in mismatch caught.

Complexity comparison

Heap allocations

• Before: O(n) node[T] structs, each heap-allocated
• After: one flat slice copy

Tree rotations

• Before: O(n log n)
• After: 0

Comparisons

• Before: O(n log n) build + O(n) compare
• After: O(n log n) sort + O(n) scan

GC objects

• Before: n pointer-linked nodes
• After: 1 contiguous array

Total time

• Before: O(n log n)
• After: O(n log n) -> same class, much smaller constant

Wrapping up

A small // TODO optimize comment turned into a tour through generic set design, Red-Black tree mechanics, lazy iterator implementation, and the surprisingly subtle semantics of duplicate detection. The final change is about 15 lines, but those 15 lines only make sense if you understand the 1000 lines underneath them.

That is most of what I have come to enjoy about reading library code: the interesting problems are rarely in the obvious places.