As an exercise on GADTs I wrote a type-safe implementation of AA trees. I'm quite happy that the AANode
data type correctly grasps the properties of the tree. Just from the type it can be concluded that the trees has the balancing property. But there I things I'd like to improve, but I didn't find a nice solution for it:
- I need separate data type for the result of
insert'
. I'm afraid nothing can be done about that, as the algorithm really needs to finely distinguish the possible outcomes. - For inserting a new value into subtrees with a black root, I need a separate function
insertBlack
. The reason is that in this case, the tree level cannot rise, and it is necessary to capture this property. I'd be nicer to have it all ininsert'
, but I'm afraid this would require dependent types (as the result type would depend on the color of the input). - I need a separate data type for the result of
insertBlack
. A solution would be to write the type of the result as an existential, like... -> (exists c . AANode c (Succ n) a)
, but AFAIK GHC doesn't support that. - Even though the AA trees are quite simple, I feel that I have a lot of pattern matching for all the different case. At the end, I have 8 possibilities. Maybe they really need to be there, as for each node, we have to distinguish if we insert into the left or the right subtree, we need to distinguish its color, and also we need to match on the recursive result to correctly update the tree.
- Perhaps, there is a better way how to capture the required properties of the trees? To construct the type of
AANode
differently?
Any other comments on how to make the code shorter, more readable or more beautiful are welcomed.
The code is below:
{-# LANGUAGE GADTs #-}
{-# OPTIONS_GHC -fwarn-incomplete-patterns #-}
-- Based on https://en.wikipedia.org/wiki/AA_tree
-- If it happens that anyone ever wants to use this code for something,
-- it's licensed under The BSD 3-Clause License.
import qualified Data.Foldable as F
-- Types representing natural numbers:
data Zero
data Succ a
-- Phantom types for possible colors of nodes.
data Red
data Black
{- | Every child node has a level -1 its parent node, except
the right node of a red node. Nodes have 3 type parameters:
Color (phantom), level (phantom), values.
-}
data AANode c n a where
-- Leaf nodes are black and have no value (their color isn't really important).
Leaf :: AANode Black Zero a
-- Red nodes can have the right child with the same level, but then
-- the child must be black. The left child must have -1 of its level.
Red :: !(AANode c n a) -> a -> !(AANode Black (Succ n) a) -> AANode Red (Succ n) a
-- Both children of a black node must have -1 of its level.
-- They can have arbitrary colors.
Black :: !(AANode c1 n a) -> a -> !(AANode c2 n a) -> AANode Black (Succ n) a
-- Two operations to fix improper tree shapes (see the Wikipedia article).
skew :: AANode Black (Succ n) a -> a -> AANode c n a -> AANode Red (Succ n) a
skew (Black a x b) y r = Red a x (Black b y r)
split :: AANode c n a -> a -> AANode Red (Succ n) a -> AANode Black (Succ (Succ n)) a
split a t (Red b r x) = Black (Black a t b) r x
-- Insert operation ..................................................
-- A possible outcome of an insert operation: Either the level
-- remains unchanged, or it's increased and the resulting node is black.
data Outcome n a where
Inc :: !(AANode Black (Succ n) a) -> Outcome n a
Same :: !(AANode c n a) -> Outcome n a
insert' :: (Ord a) => a -> AANode c n a -> Outcome n a
insert' x Leaf = Inc $ Black Leaf x Leaf
insert' x (Red l y r)
| x < y = case insert' x l of
Same l' -> Same $ Red l' y r
Inc l' -> Inc $ Black l' y r
| otherwise = case insertBlack x r of
AnyColor r'@(Black _ _ _) -> Same $ Red l y r'
AnyColor r'@(Red _ _ _) -> Inc $ split l y r'
insert' x t@(Black _ _ _)
= case insertBlack x t of AnyColor t -> Same t
-- Inserting a node into a black subtree never increases height.
-- To capture this, we have a separate function for inserting into
-- black nodes. We need this property in insert` when we're inserting
-- into a red node's right subtree.
data AnyColor n a where
AnyColor :: !(AANode c n a) -> AnyColor n a
insertBlack :: (Ord a) => a -> AANode Black (Succ n) a -> AnyColor (Succ n) a
insertBlack x (Black l y r)
| x < y = case insert' x l of
Same l' -> AnyColor $ Black l' y r
Inc l' -> AnyColor $ skew l' y r
| otherwise = case insert' x r of
Same r' -> AnyColor $ Black l y r'
Inc r' -> AnyColor $ Red l y r'
-- ---------------------------------------------------------------------
-- | The actual tree type hiding colors/levels.
data AATree a where
AATree :: !(AANode c n a) -> AATree a
instance F.Foldable (AANode c n) where
foldr _ x Leaf = x
foldr f x (Red l y r) = F.foldr f (f y (F.foldr f x r)) l
foldr f x (Black l y r) = F.foldr f (f y (F.foldr f x r)) l
instance F.Foldable AATree where
foldr f x (AATree t) = F.foldr f x t
insert :: (Ord a) => a -> AATree a -> AATree a
insert x (AATree t) = case insert' x t of
Same t -> AATree t
Inc t -> AATree t
empty :: AATree a
empty = AATree Leaf
-- TESTING CODE --------------------------------------------------
main :: IO ()
main = do
let fromList :: Ord a => [a] -> AATree a
fromList = foldr insert empty
print $ F.foldr (:) [] $ fromList "ABCD Kocka prede."