Improving General Algebraic DataType type-safe code for AA trees

Question

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:

1. 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.
2. 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 in insert', but I'm afraid this would require dependent types (as the result type would depend on the color of the input).
3. 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.
4. 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.
5. 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."

Answer 1

As suggested, I'm doing a review myself after a few years.

First, let's use GHC's extensions to derive the phantom types from data definitions. This also has the advantage of having proper kinds for each of these types.

{-# LANGUAGE GADTs, DataKinds, KindSignatures, TypeFamilies #-}
{-# OPTIONS_GHC -fwarn-incomplete-patterns #-}

-- |
-- Module: AATree
-- Copyright: Petr Pudlak
-- License: BSD-3
--
-- Based on https://en.wikipedia.org/wiki/AA_tree
module AATree where

import qualified Data.Foldable as F
import Data.Monoid

-- * Internals of the tree with type-level colors and depths.

-- | Types representing natural numbers:
data Nat = Zero | Succ Nat

-- | Phantom types for possible colors of nodes.
data Color = Red | Black

The definition of the main data type is basically unchanged. One important change here is the color of Leaf. Originally it was black, but actually it needs to be red in order to have consistent color change in the insert operation, see SomeOrNext below.

-- | 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 :: Color) (n :: Nat) a where
    Leaf :: AANode 'Red '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.
    BlackNode :: !(AANode c1 n a) -> a -> !(AANode c2 n a) -> AANode 'Black ('Succ n) a
    -- | Both children of a black node must have -1 of its level.
    -- They can have arbitrary colors.
    RedNode :: !(AANode c1 n a) -> a -> !(AANode Black ('Succ n) a) -> AANode 'Red ('Succ n) a

-- ** Two operations to fix improper tree shapes (see the Wikipedia article).

-- | Left node grows too large.
skew :: AANode 'Black ('Succ n) a -> a -> AANode c n a -> AANode 'Red ('Succ n) a
skew (BlackNode a x b) y r = RedNode a x (BlackNode b y r)

-- | Right node grows too large.
split :: AANode c n a -> a -> AANode 'Red (Succ n) a -> AANode 'Black (Succ (Succ n)) a
split a t (RedNode b r x) = BlackNode (BlackNode a t b) r x

-- ** Insert operation

Instead of two types, Outcome and AnyColor, which were holding outputs of the different insert operations, let's have just one, parametrized by Color.

-- | Holds the outcome of an insert operation on a node of color 'c'.
-- Inserting into a black node retains its depth and possibly chcanges it
-- to red ('Same'). Inserting into a red one either keeps it the same
-- ('Same') or changes it to a black node of increased depth ('Next').
data SameOrNext (c :: Color) (n :: Nat) a where
    Same :: AANode c_any n a -> SameOrNext c n a
    Next :: AANode 'Black ('Succ n) a -> SameOrNext 'Red n a

This allows us to have just a single insert function with a simple recursive definition. There are still 8 cases to cover, but it seems this can't be simplified, at least not without a considerable sacrifice of code readability. In particular, we need to branch whether we move to the left or right, distinguish inserting into a red/black node, and act depending on the outcome of a recursive insert.

insert' :: (Ord a) => a -> AANode c n a -> SameOrNext c n a
insert' x Leaf = Next (BlackNode Leaf x Leaf)
insert' x (BlackNode l y r)
    | x < y     = case insert' x l of
                    Same l' -> Same (BlackNode l' y r)
                    Next l' -> Same (skew l' y r)
    | otherwise = case insert' x r of
                    Same r' -> Same (BlackNode l y r')
                    Next r' -> Same (RedNode l y r')
insert' x (RedNode l y r)
    | x < y     = case insert' x l of
                    Same l' -> Same (RedNode l' y r)
                    Next l' -> Next (BlackNode l' y r)
    | otherwise = case insert' x r of
                    Same r'@(BlackNode _ _ _) -> Same (RedNode l y r')
                    Same r'@(RedNode _ _ _) -> Next (split l y r')

-- * Public interface.

-- | The actual tree type hiding colors/levels.
data AATree a where
    AATree :: !(AANode c n a) -> AATree a

Minor change in the Foldable definition: Instead of using foldr, it's more natural to define it using foldMap, especially for trees. This can also have significant performance benefits in certain cases, like when using the Last monoid (untested).

instance F.Foldable (AANode c n) where
    foldMap f Leaf = mempty
    foldMap f (RedNode l y r) = F.foldMap f l <> f y <> F.foldMap f r
    foldMap f (BlackNode l y r) = F.foldMap f l <> f y <> F.foldMap f r

instance F.Foldable AATree where
    foldMap f (AATree t) = F.foldMap f t

insert :: (Ord a) => a -> AATree a -> AATree a
insert x (AATree t) = case insert' x t of
                                Same t -> AATree t
                                Next t -> AATree t

empty :: AATree a
empty = AATree Leaf


-- * Very simple test.

main :: IO ()
main = do
    let fromList :: Ord a => [a] -> AATree a
        fromList = foldr insert empty

    print . F.toList . fromList $ "ABCD Kocka prede."

I also changed comments to conform to proper the Haddock documenation format.

---

Further ideas: In addition to insertion and deletion, there are two other fundamental operations on search trees:

• Split a tree into two, one with elements less and other with elements greater than a given key.
• Merge two trees for which we know that all elements in one are greater than all elements in another one.

Not only they are very useful on their own, but they can be also used to easily implement other operations such as insert, delete, merge and intersect. The general idea is that when we want to do an operation on two trees, we can split the first one using the root key of the second one and perform the operation recursively on the subtrees of the second one.

Splitting two AATrees is more-or-less straightforward, but merging them is more problematic, as for merging we need to start from the bottom, not from the top, and we don't have information about heights of trees. We could traverse both trees until the bottom layer, unwinding the visited nodes into a heterogeneous lists that holds trees of increasing heights and then merge them backwards up. Or we could embrace this idea into the data structure itself and implement finger trees based on AANode.