Compile time invariants in Haskell

One of the goals of a type system is to ensure that certain invariants are held when the program runs. Those invariants usually have to do with the state of the program. For example, a type system can tell you that a function was declared (or inferred) to receive a parameter of a type \(\tau\), and yet an argument of type \(\rho\) was passed, and there’s no way to unify those two types.

A type system can check these properties either at compile time, or at runtime. These are called static and dynamic type systems respectively. I prefer having my compiler catch as many errors as it can before I run my program, so I prefer, in general, static typing.

To that end, I wanted to write a small implementation of AVL trees which ensured, at compile time, that the AVL height invariants held for all trees I could ever generate with that implementation. I’ve uploaded the resulting code to GitHub, as well as packaged it into a Cabal library over at Hackage, under the name avl-static.

I’ll explain a bit how this code ensures statically that the invariants hold.

AVL trees

Height invariant

To refresh your memory on AVL trees, they are binary search trees that fulfill the AVL height invariant. The height invariant says that if I have a node which has two subtrees \(l\) and \(r\), then their heights, \(h(l)\) and \(h(r)\) respectively, differ by no more than 1. That is, \(|h(l) - h(r)| \le 1\).

Using this invariant, one can prove the following theorem:

Theorem: The height of an AVL tree of \(n\) nodes is never more than \(\frac{\log_2(n + 2)}{\log_2(\varphi)} - 2 + \log_\varphi(\sqrt{5})\), where \(\varphi\) is the golden ratio.

Inductive construction

One way we can describe trees that follow the AVL height invariants is by considering “balance factors”.

Definition. The balance factor of a node \(v\) with left subtree \(l\) and right subtree \(r\) is \(h(r) - h(l)\).

It should be clear, then, that an AVL tree is one where each node has a balance factor of \(0\), \(1\), or \(-1\). Furthermore, if we allow our heights to increase by \(1\), letting the empty tree have height \(0\), we can state this as a way to construct every AVL tree:

• The empty tree is an AVL tree of \(0\) nodes and height \(0\).
• If \(l\) and \(r\) are AVL trees, then let \(T =\) .
  • If \(h(l) = h\), and \(h(r) = h + 1\), then \(T\) is an AVL tree of balance factor \(1\) and height \(h + 2\).
  • If \(h(l) = h + 1\), and \(h(r) = h\), then \(T\) is an AVL tree of balance factor \(-1\) and height \(h + 2\).
  • If \(h(l) = h(r) = h\), then \(T\) is an AVL tree of balance factor \(0\) and height \(h + 1\).
• Nothing else is an AVL tree.

This, then, allows us to state explicitly every possible way to construct an AVL tree. If we could encode the needed information as a type, and these constructions as type constructors of this type, then the type checker could tell us, at compile time, whether or not our trees are AVL trees.

Type-encoding height invariants

DataKinds

In our construction of this type, we see that we need to mention the height of a given AVL tree in order to tell if a given tree is AVL. That is, we would want to write code like this:

data Tree a (h + 1) = EmptyTree
                    | BalanceFactor0 a (Tree a h) (Tree a h)
                    | BalanceFactor1 a (Tree a (h - 1)) (Tree a h)
                    | BalanceFactor-1 a (Tree a h) (Tree a (h - 1))

This seems a bit strange, because in ordinary Haskell, types don’t have values as part of them. For example, you can’t express the types “Lists of 45 Bools”, or “Strings containing the substring ‘potato’”. This is because, in Haskell, types can only depend on other types. You can say “Lists of Ints” ([Int]), or “Sets of Strings” ( Set String).

To be able to say these things, we would need dependent types. These are types that can depend on values. Haskell doesn’t currently have that (as Agda or Idris do), but we can fake just enough of it for our purposes, with the DataKinds extension. This extension allows us to take a given type, and lift its values to be types. As an example:

data Nat = Zero | Succ Nat

This declares a datatype Nat, which encodes the usual natural numbers. Without DataKinds, the following are values: Zero, Succ Zero, Succ (Succ Zero), … . With DataKinds, these can be used as values, but they also get promoted to full blown types. That is, Succ (Succ Zero) can be used as a type (whose kind is Nat, as opposed to the usual kind of types like Int or [Bool], which is *). With this, we can now say “A tree of height 13”, by encoding the number 13 as a series of 13 Succ type constructors on a Zero. This means we can now talk about natural numbers in our types, which is enough of dependent types to do what we want - to keep track of tree heights.

To fully use the power of DataKinds, however, we’ll need one more extension.

GADTs

GADTs stands for generalized abstract datatypes. If you look carefully at the datatype we wanted to define above, there’s a bit of a mistake: An EmptyTree can have height \(h + 1\) for any \(h\). We don’t want this, since we want the \(h\) parameter to measure the tree’s height. We would like to have a bunch of constructors for our tree, and each constructor spits out a slightly different type (because they will have different heights, depending on which constructor you use).

GADTs let us do this. That is, we can have the following:

data List :: * -> Nat -> * where
  Nil :: List a Zero
  Cons :: a -> List a n -> List a (Succ n)

Notice that the return type of the constructors for the type is not necessarily the same for all of them. Nil constructs the type List a Zero, which is different from List a (Succ n). We want these types to be different, because we are, presumably, interested in knowing the length of this list at compile time. Without GADTs, we’re forced to have the return type of all constructors be the same.

The signature to the right of List means that it takes some type (*), a Nat, and returns some other type (*). * is one of Haskell’s kinds. You can think of kinds as “types of types”. For example, Integer has kind *, meaning “it’s a type”. Maybe, on the other hand, has kind * -> *, meaning when you give it a type, it’ll give you a type. That way, Maybe Integer is a type, but Maybe, by itself, is not a type. Similarly, Either has kind * -> * -> *.

When we use DataKinds, we can not only say “a type”, but “a type that is lifted from the following user type”. Meaning, we can request that a type is either Zero, Succ Zero, Succ (Succ Zero), and so on, by saying it has kind Nat.

AVL datatype

With these two extensions, we can finally express the above rules as one Haskell datatype:

data AVLNode :: Nat -> * -> * where
  Nil :: AVLNode Zero a
  Leftie :: a -> AVLNode (Succ n) a -> AVLNode n a -> AVLNode (Succ (Succ n)) a
  Rightie :: a -> AVLNode n a -> AVLNode (Succ n) a -> AVLNode (Succ (Succ n)) a
  Balanced :: a -> AVLNode n a -> AVLNode n a -> AVLNode (Succ n) a

Here we name Leftie an AVL tree with a balance factor of \(-1\), and Rightie an AVL tree with balance factor \(1\). It is easy to see, then, by induction, that any AVLNode h a has height \(h\). Furthermore, because of the type of each of the constructors, in order to create a value of the type AVLNode, the compiler must be able to prove that the resulting tree is indeed balanced. If you look at the type of these constructors as implications, this is one way where the Howard-Curry isomorphism shows up: our constructors are proofs that the resulting tree is balanced.

As an example, the following are all AVL nodes of varying heights:

t1 :: AVLNode (Succ Zero) Int
t1 = Balanced 4 Nil Nil

t2 :: AVLNode (Succ (Succ Zero)) Char
t2 = Rightie Nil (Balanced 'X' Nil Nil)

t3 :: AVLNode (Succ (Succ Zero)) String
t3 = Balanced ", " (Balanced "Hello" Nil Nil) (Balanced "World" Nil Nil)

While the following are both compile time errors (can you see why?):

t4 = Balanced "Hi" Nil (Balanced "Bye" Nil Nil)
t5 = Rightie "Aloha" Nil (Rightie "Hawaii" Nil (Balanced "!" Nil Nil))

Operations

Navigating through the tree

As you can imagine, in order to implement operations on this tree, we’ll often be “navigating” it. We’ll sometimes want to go to the left subtree of a node, sometimes to the right, and sometimes to the parent (even though, if you notice, an AVLNode does not keep track of its parent). To be able to do this, we’ll incorporate a classic functional programming technique: the Zipper.

Instead of explaining zippers poorly, I refer to a couple of good sources on it:

• The original Zipper paper
• Edward Yang’s You Could Have Invented Zippers
• Learn You A Haskell For Great Good’s chapter on Zippers

So to be able to navigate around the tree, reconstructing a zipper to a subtree \(T\) from a zipper to a sub-subtree \(T'\), we need to know five things:

• Whether \(T'\) was the left or right child of \(T\).
• Who was the other child of \(T\)?
• What was \(T\)’s balance factor?
• What value (of type a) was in the root of the subtree \(T\)?
• How does one get from \(T\) back to the root?

For example, if you tell me that you have a zipper to a subtree T, and that

• T has a left subtree T'.
• T' has a sibling subtree S.
• T, had a balance factor of \(1\).
• T had a value v, of type a, in its root.
• To get from the root of the tree to T, you follow a path ctx.

Then I can give you a zipper to T'. The subtree the zipper points to, is T', and the path from the root to T', is ctx followed by “go left”.

Similarly, given a zipper to T', I can give you a zipper to T: I reconstruct T because I know T = Rightie v T' S, and I “pop” one instruction from the path to T', yielding ctx again.

To see actual code, here is how one could implement a zipper for this type of AVL tree.

-- A Context n a means a traversal from a root of an AVL Tree of height n,
-- whose non-nil nodes have values of type a, to some subtree.
data Context :: Nat -> * -> * where
  -- A balanced context. BC True x y means the traversal went left,
  --  on a node Balanced x, where y is the subtree not taken in the
  --  traversal.
  BC :: Bool -> a -> AVLNode n a -> Context (Succ n) a -> Context n a
  -- A leftie context, where we've taken the right branch of the subtree.
  LRC :: a -> AVLNode (Succ n) a -> Context (Succ (Succ n)) a -> Context n a
  -- A leftie context, where we've taken the left branch of the subtree.
  LLC :: a -> AVLNode n a -> Context (Succ (Succ n)) a -> Context (Succ n) a
  -- A rightie context, where we've taken the left branch of the subtree.
  RLC :: a -> AVLNode (Succ n) a -> Context (Succ (Succ n)) a -> Context n a
  -- A rightie context, where we've taken the right branch of the subtree.
  RRC :: a -> AVLNode n a -> Context (Succ (Succ n)) a -> Context (Succ n) a
  -- The root context, where every traversal of an AVL tree starts.
  Root :: Context n a

data Zipper a = forall n. Zipper (AVLNode n a) (Context n a)

For example, a context representing “Start at the root of a tree of height 3, then go left on a Rightie node whose value is x and whose right subtree is s, then go Right on a Balanced node whose root value is y and whose left subtree is t”, we would use the following context:

ctx = BC False y t (RLC x s Root)

Notice that a context has, as part of its type, a height parameter, indicating the height of the subtree we traversed to. The type of Root above is Context (Succ (Succ (Succ Zero))) a.

Thus, if we see a context RLC x s ctx, we know that we took a left turn in a Rightie parent node. Let’s call this parent node p. Since we are the left node of a Rightie parent p, s was the right subtree of a Rightie parent, and so if s has height Succ n, then p had height Succ (Succ n), and the subtree we are standing in has height n. Then if you give me a traversal from the root to p, which would have type Context (Succ (Succ n)) a, a value x :: a, and you give me s :: AVLNode (Succ n) a, I can give you a traversal RLC x s ctx :: Context n a, to p’s left subtree, which has height n. Thus, RLC :: a -> AVLNode (Succ n) a -> Context (Succ (Succ n)) a -> Context n a.

Lastly, we see that a Zipper is an existential type, which means that if you are given a Zipper a, you know that it is of the form Zipper (AVLNode n a) (Context n a), for some n, but you don’t know what n is. You only know that it exists. That is why Zipper is called an existential type, and why the extension that allows us to type that forall n. is called ExistentialQuantification. The reason I chose to do this was that I didn’t want users of a Zipper to need to know the height of the subtree the Zipper was pointing to, it’s more of an implementation detail.

So to cement this a bit more, let’s see how one can traverse the tree using this Zipper.

left :: Zipper a -> Zipper a
left (Zipper (Balanced x l r) ctx) = Zipper l (BC True x r ctx)
left (Zipper (Leftie x l r) ctx) = Zipper l (LLC x r ctx)
left (Zipper (Rightie x l r) ctx) = Zipper l (RLC x r ctx)
left z = z

right :: Zipper a -> Zipper a
right (Zipper (Rightie x l r) ctx) = Zipper r (RRC x l ctx)
right (Zipper (Leftie x l r) ctx) = Zipper r (LRC x l ctx)
right (Zipper (Balanced x l r) ctx) = Zipper r (BC False x l ctx)
right z = z

up :: Zipper a -> Zipper a
up (Zipper x (LRC v y ctx)) = Zipper (Leftie v y x) ctx
up (Zipper x (LLC v y ctx)) = Zipper (Leftie v x y) ctx
up (Zipper x (RLC v y ctx)) = Zipper (Rightie v x y) ctx
up (Zipper x (RRC v y ctx)) = Zipper (Rightie v y x) ctx
up (Zipper x (BC True v y ctx)) = Zipper (Balanced v x y) ctx
up (Zipper x (BC False v y ctx)) = Zipper (Balanced v y x) ctx
up z@(Zipper _ (Root _)) = z

With these 3 functions, we can now traverse an AVL tree arbitrarily. We’ll use these to implement insertion and deletion.

Inserting

Such a tree would be pretty useless if we had to manually construct its instances by applying constructors ourselves. After all, an AVL tree is supposed to balance itself, and its user is not supposed to know the internal tree structure. So the function we need to implement is something like this:

insert :: Ord a => a -> AVLTree a -> AVLTree a

where AVLTree a means an AVL tree whose values are of type a. Notice that I explicitly don’t want to care about the height parameter of my trees. This is a design decision, the API shoud not expose the user to the height of the tree. This is, in fact, a necessity: It’s not possible to know statically, given an element \(x\) of type \(a\) and an AVL tree \(t\), if inserting \(x\) into \(t\) will increase \(t\)’s height, so neither of these two signatures would be valid (that is, we could not implement them):

insert' :: Ord a => a -> AVLNode n a -> AVLNode n a
insert'' :: Ord a => a -> AVLNode n a -> AVLNode (Succ n) a

We can use existential types again for this task:

data AVLTree a = forall n. T (AVLNode n a)

Note this doesn’t mean that an AVLTree a doesn’t know the height of the tree it holds - it definitely knows it, it just doesn’t show it. Again looking at this constructor as an implication via Howard-Curry, this can be read as “for all types \(a\) and \(n \in \mathbb{N}\), given an AVLNode \(n\) \(a\), you can construct an AVLTree \(a\).” Thus if you have an AVLTree \(a\), you know there existed some \(n \in \mathbb{N}\) from which the tree was made. These types are thus called “existential”, even when the quantifier is universal, not existential. It’s simply a matter of reversing the implication.

One way we can implement insertion of \(x\) into \(t\) is by navigating (via a zipper) to the site in \(t\) where \(x\) should be found, placing a single node containing \(x\) there, and then fixing \(t\) as we go up the zipper. If when finding the insertion site we see that we’ve zipped to a location where there’s already an element (which then will be equal to \(x\)), we needn’t do anything, since \(t\) already contains \(x\). In code, this means something like this:

insert :: Ord a => a -> AVLTree a -> AVLTree a
insert x t = case zipTo x (unZip t) of
  Zipper Nil ctx -> insertAndFix (Balanced x Nil Nil) ctx
    _ -> t

zipTo x t is a function that finds where x should be, if it was in t. It either returns a zipper focused on x, or focused on the place where x should be (that is, focused on a Nil where we should find x, if x were in t). It can be implemented as such:

value :: Zipper a -> a
value (Zipper (Balanced x _ _) _) = x
value (Zipper (Leftie x _ _) _) = x
value (Zipper (Rightie x _ _) _) = x
value (Zipper Nil _) = error "Zipper points to Nil."

zipTo :: Ord a => a -> Zipper a -> Zipper a
zipTo _ z@(Zipper Nil _) = z
zipTo x z = let v = value z
            in case compare x v of
                 EQ -> z
                 LT -> zipTo x $ left z
                 GT -> zipTo x $ right z

Creating a zipper from a tree is easy, it’s simply a pointer to the root node, with a context of Root:

unZip :: AVLTree a -> Zipper a
unZip (T r) = Zipper r Root

Now we need to write insertAndFix. The type of this function is quite interesting. Notice that the type of the context we are giving it must be Context Zero a, because the zipper must be pointing to a Nil node. Since we will make insertAndFix recursive, its type will be as follows:

insertAndFix :: AVLNode (Succ n) a -> Context n a -> AVLTree a

This means we have a Context n a, which means “a hole in an AVLTree \(a\) where a subtree of height \(n\) fits”, and an AVLNode (Succ n) a, which means “a subtree of an AVLTree \(a\) of height \(n + 1\)”, and we want to return an AVLTree \(a\). So the general notion is that we have a subtree of height \(n + 1\), and we want to insert it in a place where a subtree of height \(n\) fits. This will guide the function’s implementation, and will be dual to our deletion routine.

Let’s do an easy case of insertAndFix first. Say we want to insert at the root. This means the tree was empty, and we’re inserting our first node. We simply create a new tree with that sole node as its root:

insertAndFix t Root = T t

For a less trivial case, suppose we want to insert a tree \(t\) of height \(h + 1\), and our context is LRC v y ctx. This means it’s a hole in the right-child-side of a Leftie subtree, whose root is v, and whose left child is y. So we have the following picture:

We know the height of \(y\) because its parent is a Leftie, and the hole on the right is for a subtree of height \(h\), so the left subtree must have height \(h + 1\), and the entire subtree must have height \(h + 2\). This case is indeed very easy! We simply insert our tree \(t\) in that hole, and we now have the following picture:

This means we turned a Leftie into a Balanced, because both children now have the same height, \(h + 1\). In code, this translates to:

insertAndFix t (LRC v y ctx) = zipUp $ Zipper (Balanced v y t) ctx

where zipUp will create a tree from this zipper:

zipUp :: Zipper a -> AVLTree a
zipUp (Zipper t Root) = T t
zipUp z = zipUp (up z)

The same goes for its symmetric case:

insertAndFix t (RLC v y ctx) = zipUp $ Zipper (Balanced v t y) ctx

Now for a harder case. Suppose we are inserting in the left side of a Leftie subtree. That is, we have a tree \(t\) of height \(h + 1\), and we want to insert it in the following context:

We can’t simply insert \(t\) in there, because the difference between the two braches would be \(2\), which violates the AVL invariants. So in order to determine what we need to do, we’ll split into three cases for \(t\) (note \(t\) can’t be Nil because its height is always of the form \(h + 1\) for some \(h\)):

• \(t\) is a Leftie, of the form We can then perform a left rotation, in order for the resulting subtree to be as follows: This leaves us with a balanced tree. Written in code, the above procedure is

insertAndFix (Leftie b g p) (LLC a d ctx) = zipUp z
  where
    z = Zipper (Balanced b g (Balanced a p d)) ctx

• \(t\) is Balanced. This case is recursive. We have the same context as before, but now \(t\) looks like this: We can perform another rotation to get this to be locally balanced: Now the subtree is locally balanced… but it has increased in height! The place where this subtree lived had room here for a height \(h+1\) subtree. We’ve now increased its height. This is still fine, however, because we already have a routine for inserting a subtree of a height \(h + 2\) in a place (a context) that has room for a subtree of height \(h + 1\), that’s precisely what insertAndFix does! So we need to call insertAndFix on this new rotated subtree. In code, this works out to be this:

insertAndFix (Balanced b g p) (LLC a d ctx) = goUp
  where
    goUp = insertAndFix (Rightie b g (Leftie a p d)) ctx

Notice that the context we use in the recursive call is our contexts’ context. That is, the context we were passed is LLC a d ctx, and we are recursing on ctx.

• \(t\) is a Rightie. This case is the trickiest. We need to split into 3 cases, based on what the right child of \(t\) is. Since \(t\) is a Rightie, it’s of the following form: Just to refresh our memory, we’re trying to insert this subtree \(t\), which has height \(h + 1\), in the following context (marked in blue), which has room for a subtree of height \(h\): Onto the splitting, then. Since \(t\) is a Rightie, its right child \(q\) can’t be Nil. So we have \(3\) cases for \(q\):

  • \(q\) is a Leftie, of the form We can rebalance the whole thing by rotating twice, leaving us with this: This is balanced, and has height \(h + 1\), which means we can now zipUp it with ctx, which has type Context (Suc n) a (remembering that our insertion context, which had height \(n\), was LLC a d ctx, so ctx is the parent context, which has a height of \(n + 1\).

  • \(q\) is Balanced. The same rotations as above apply, only the heights change (\(t2\) has height \(h - 1\)).

  • \(q\) is a Rightie. Again the same rotations as above leave the entire subtree, inside its insertion context, balanced. The heights of the subtrees change, namely, \(t2\) has height \(h - 1\), and \(t1\) has height \(h - 2\).

  Translated into code, we have the following:

  insertAndFix (Rightie b g q) (LLC a d ctx) = zipUp z
  where
        z = Zipper t ctx
        t = case q of
              Rightie p t1 t2 -> Balanced p (Leftie b g t1) (Balanced a t2 d)
              Leftie p t1 t2 -> Balanced p (Balanced b g t1) (Rightie a t2 d)
              Balanced p t1 t2 -> Balanced p (Balanced b g t1) (Balanced a t2 d)

This is how we deal with contexts of the form LRC, RLC, and LLC. RRC is just like LLC, it’s its mirror image. I’ll leave the balanced case, BC, as an exercise for the interested reader (hint: it’s short, and always recursive).

The cool part about this is that there’s no way we could’ve made a mistake and left the tree unbalanced. The type system takes care of that part for us, because of the constructors’ types. We can make mistakes that violate the binary search tree invariants (but not the “binary” part, just the “search” part ;)), but never the AVL balance invariants.

Deleting

One interesting fact about deletion is that it’s essentially a dual operation to insertion. In insertion, we created a tree of height \(h + 1\), and we wished to place it in a context of height \(h\). In deletion, we will have a tree of height \(h\), and we wish to place it in a context of height \(h + 1\). Here again, the type system will take care that we cannot violate the height invariants. To put it another way, our types will serve as proofs that our height invariants hold.

BST deletion

Recall that non-leaf deletion in a binary search tree (in particular, in an AVL tree) happens by first replacing a node’s value with the value of its predecessor (or successor), and then deleting the other node whose value we took. Deleting a leaf is straightforward removal of the leaf. In addition, an AVL tree requires rebalancing after the usual binary search tree deletion, to make sure the height invariants hold.

To implement deletion, then, we first find the desired node, then delete it a-la BST, and then rebalance. Let’s do this step by step.

delete :: Ord a => a -> AVLTree a -> AVLTree a
delete x t = case zipTo x (unZip t) of
  Zipper Nil _ -> t
  z -> deleteBST z

deleteBST will perform a BST removal, followed by rebalancing. If the node is not in the tree (as evidenced by the zipTo returning a Zipper to Nil), we need not do anything. The rebalancing will be dual to the insertion case, as the type suggests.

The type of deleteBST should be straightforward:

deleteBST :: Eq a => Zipper a -> AVLTree a
rebalance :: AVLNode n a -> Context (Succ n) a -> AVLTree a

We don’t require the Ord typeclass anymore, since we already found the deletion site. Now all our actions are guided by the tree’s structure, not its values. We require Eq solely for an efficiency concern we’ll see later, but we could implement deleteBST :: Zipper a -> AVLTree a.

First the easy cases: deletion of a leaf.

deleteBST (Zipper (Balanced _ Nil Nil) ctx) = rebalance Nil ctx
deleteBST (Zipper (Rightie _ Nil r) ctx) = rebalance r ctx
deleteBST (Zipper (Leftie _ l Nil) ctx) = rebalance l ctx

I’ll only write one nontrivial case of deleteBST here, since the other two are extremely similar:

-- |Given a 'Zipper' to a node @X@ in the tree, returns a 'Zipper' to the
-- smallest node in the subtree rooted at @X@.
zipToSmallest :: Zipper a -> Zipper a
zipToSmallest z | canGoLeft z = zipToSmallest (left z)
                | otherwise = z

zipToSuccessor :: Zipper a -> Maybe (Zipper a)
zipToSuccessor z | canGoRight z = Just . zipToSmallest $ right z
                 | otherwise = let parent = zipToFirstLeftChild z
                               in  up <$> parent

-- |Given a 'Zipper' @Z@, which points to a subtree @S@, returns a 'Zipper' to
-- the first ancestor of @S@ which is a left child of its parent. If such an
-- ancestor does not exist, returns 'Nothing'.
zipToFirstLeftChild :: Zipper a -> Maybe (Zipper a)
zipToFirstLeftChild z | isLeft z = Just z
zipToFirstLeftChild z | canGoUp z = zipToFirstLeftChild (up z)
                      | otherwise = Nothing

deleteBST z@(Zipper (Rightie k _ _) _) =
  let Just s = zipToSuccessor z
  in case s of
       Zipper (Balanced k' Nil Nil) ctx' -> rebalance Nil (fixContext k k' ctx')
       Zipper (Rightie k' Nil r) ctx' -> rebalance r (fixContext k k' ctx')
       _ -> error "The impossible has happened, bad successor found."

Notice the call to fixContext. This is where we use Eq. The issue is that if we were to simply use ctx', the context of the successor of z (s), this context would eventually have z in it, as we climbed up. We actually want to change z’s value, from k to k', the value of s. This function climbs up the contexts, and when it finds a node with a value of k, it replaces it by k' (remembering that z is s’s ancestor). It’s not a very pretty solution, but it is simple :)

-- |Replaces a given value by another, in the 'AVLTree' represented by a
-- 'Context'.
fixContext :: forall a n. Eq a => a -> a -> Context n a -> Context n a
fixContext k k' = go
  where
    z x = if x == k then k' else x
    go :: Context n' a -> Context n' a
    go r@(Root _) = r
    go (BC goLeft x y ctx) = BC goLeft (z x) y (go ctx)
    go (LLC x y ctx) = LLC (z x) y (go ctx)
    go (LRC x y ctx) = LRC (z x) y (go ctx)
    go (RRC x y ctx) = RRC (z x) y (go ctx)
    go (RLC x y ctx) = RLC (z x) y (go ctx)

Not too high-tech here. Now onto the rebalancing.

Rebalancing

First, an easy case: We are trying to place a subtree of height \(n\) in a Root context. This just means we decremented the height of the tree, but since we don’t need to return a tree of the same height we started with (remember, we only return an AVLTree a, not an AVLNode n a for some specific n), we can simply use this tree as our new root.

rebalance t Root = T t

Two easy cases follow. In these, we have some “slack” because of the height invariants. That is, it’s OK if the tree we are changing decreases its height by 1, we don’t violate the invariants. The cases are as such:

rebalance t (BC True a d ctx) = zipUp $ Zipper (Rightie a t d) ctx
rebalance t (BC False a d ctx) = zipUp $ Zipper (Leftie a d t) ctx

So if we are decreasing the height of the left child of a balanced node, we now simply have made this node become a Rightie.

Now two recursive cases. Here we can easily rebalance locally, but we decrement the height of the parent of this subtree. To fix this, we simply need to call rebalance on the resulting tree.

rebalance t (LLC a d ctx) = rebalance (Balanced a t d) ctx
rebalance t (RRC a d ctx) = rebalance (Balanced a d t) ctx

To illustrate the first line (the second one is just a mirror image), we have the following context:

If we now insert a tree of height \(h\) in the blue context, we transform the parent (\(a\)) into a Balanced node, but we’ve decremented its height, from \(h + 2\) to \(h + 1\). We need to rebalance it, then, in it’s context. This is why we recursively call rebalance on it.

Two context types are left: going left on a Rightie subtree (RLC), and going right on a Leftie context (LRC). We’ll do the first one, since the second one is, as expected, its mirror image.

We have three cases, based on whether our sibling is Balanced, a Leftie, or a Rightie.

• Our sibling is Balanced. Then the context where we need to delete from looks like this (recall we wish to place a tree \(t\) of height \(h\) in the blue circle, where a tree of height \(h + 1\) used to be): We can rebalance this situation by rotating as such: We can write this down as

rebalance t (RLC a (Balanced d t1 t2) ctx) = zipUp z
  where
    z = Zipper (Leftie d (Rightie a t t1) t2) ctx

This finishes the rebalancing, since we did not change the height of the root of this subtree (\(d\)). * Our sibling is a Rightie. We have the following situation: We can rebalance as before, and obtain the following: But note that we have now changed the height of the root of the subtree, so we need to keep rebalancing up the context chain. This translates to the following:

rebalance t (RLC a (Rightie d t1 t2) ctx) = rebalance z ctx
  where
    z = Balanced d (Balanced a t t1) t2

• Our sibling is a Leftie. Here we need to split by cases, depending on the type of its left child. Pictorially, the context we are inserting into is as such: In all cases, our resulting balanced tree will look like this: We just need to be careful with the heights of the subtrees \(p1\) and \(p2\). The code looks like this:

rebalance t (RLC a (Leftie d t1 t2) ctx) = rebalance z ctx
  where
        z = case t1 of
          Leftie q p1 p2 -> Balanced q (Balanced a t p1) (Rightie d p2 t2)
          Rightie q p1 p2 ->  Balanced q (Leftie a t p1) (Balanced d p2 t2)
          Balanced q p1 p2 -> Balanced q (Balanced a t p1) (Balanced d p2 t2)

The case of LRC is a mirror image of this one. Once more the type system will make sure we can’t actually compile code that will create unbalanced AVL trees, so all we need to be careful with is the binary search tree property.

Conclusion

We’ve seen how we can use a few modifications to Haskell’s type system to encode a nonontrivial invariant for our data structures. This is already hitting Haskell’s limitations, however - it cannot, in general, express types that depend on values. More powerful type systems can do this. For instance, in Agda we could verify both the height invariants of this tree, and the fact that it is ordered.

Writing all the rebalancing code, while complicated, it only complicated at compile time. There is no need to run tests to check if the tree is indeed balanced. This saves a lot of time (and thus money) when one can stop thinking about some property failing in production, at 3am. While strong type systems are sometimes seen as a burden to programmer productivity (after all, we sometimes want to say “I don’t know how to explain this to you, type system, just believe me, it works!”), this example makes a case for the opposite: The development of the actual functions was completely guided by the types, and the types also ensured we could not make a lot of mistakes which we otherwise might if we had used a weaker, or less expressive, type system.