(How) should I avoid writing procedural code in Haskell?

Question

I'm trying to implement a left leaning red black tree as described here. This is their snippet for insert

private Node insert(Node h, Key key, Value value) {
    if (h == null) return new Node(key, value);

    if (isRed(h.left) && isRed(h.right)) colorFlip(h);

    int cmp = key.compareTo(h.key);
    if (cmp == 0) h.val = value;
    else if (cmp < 0) h.left = insert(h.left, key, value);
    else h.right = insert(h.right, key, value);

    if (isRed(h.right) && !isRed(h.left)) h = rotateLeft(h);
    if (isRed(h.left) && isRed(h.left.left)) h = rotateRight(h);

    return h;
}

My attempt at a port in Haskell is... quite ugly, with a lot of repetition. I think it's because I'm still thinking procedurally. Any feedback on what I should do differently? Or is there no way around having many next-state variables (x', x'', x''')? Should I be approaching this completely differently?

data Colour = Red | Black deriving (Show)

data Tree a
    = Branch (Tree a) a (Tree a) Colour
    | Leaf
    deriving (Show)

add :: (Ord a) => Tree a -> a -> Tree a
add tree val
    = let
        (Branch left' node' right' _) = fix_up $ do_add tree val
    in (Branch left' node' right' Black) -- root always black

do_add :: (Ord a) => Tree a -> a -> Tree a
do_add (Branch left node right colour) val
    | val < node = (Branch (add left val) node right colour)
    | val > node = (Branch left node (add right val) colour)
    | otherwise = (Branch left node right colour)
do_add Leaf val = (Branch Leaf val Leaf Black)

get_left_node :: Tree a -> Tree a
get_left_node (Branch left _ _ _) = left
get_left_node Leaf = Leaf

fix_up :: Tree a -> Tree a
fix_up (Branch left node right colour)
    = let
        branch' = if ((not (is_red left)) && (is_red right)) then (rotate_left (Branch left node right colour)) else (Branch left node right colour)
        (Branch left' _ right' _) = branch'
        branch'' = if ((is_red left') && (is_red (get_left_node left'))) then (rotate_right branch') else branch'
        (Branch left'' _ right'' _) = branch''
        branch''' = if ((is_red left'') && (is_red right'')) then (flip_colours branch'') else branch''
    in branch'''


rotate_left :: Tree a -> Tree a
rotate_left (Branch left node (Branch right_left right_node right_right right_colour) colour)
    = let
        left' = (Branch left node right_left Red)
        centre' = (Branch left' right_node right_right colour)
    in centre'

rotate_right :: Tree a -> Tree a
rotate_right (Branch (Branch left_left left_node left_right left_colour) node right colour)
    = let
        right' = (Branch left_right node right Red)
        centre' = (Branch left_left left_node right' colour)
    in centre'

flip_colours :: Tree a -> Tree a
flip_colours (Branch (Branch left_left left_node left_right left_colour) node (Branch right_left right_node right_right right_colour) colour) = let
        left' = (Branch left_left left_node left_right (invert_colour left_colour))
        right' = (Branch right_left right_node right_right (invert_colour right_colour))
        centre' = (Branch left' node right' (invert_colour colour))
    in centre'

is_red :: Tree a -> Bool
is_red (Branch _ _ _ Red) = True
is_red _ = False

is_black :: Tree a -> Bool
is_black node = not $ is_red node

Answer 1

The code looks much better after refactoring!

Additionally I'd strongly suggest to keep line lengths within some limit. Usual choices are something between 72 and 80. There are two reasons for it:

• People with smaller screens aren't able to see a piece code at once and have to scroll. This makes reading the code next to impossible (like when your snippet here on SO doesn't fit into its frame).
• Even for a person with a wide screen it's difficult to read text with long lines. It's hard for eyes to focus where the next line starts.

Don't be afraid to use short identifiers, if theirs scope is limited to a short function. For example, in my opinion this

rotate_left :: Tree a -> Tree a
rotate_left (Branch l v (Branch rl rv rr _rc) c)
          = (Branch (Branch l v rl Red) rv rr c)

is more readable, and it's visually easy to spot what is going on.

One thing that will also help you make the code shorter more readable is using as-patterns (see also this question):

do_add :: (Ord a) => Tree a -> a -> Tree a
do_add branch@(Branch left node right colour) val
    | val < node = (Branch (add left val) node right colour)
    | val > node = (Branch left node (add right val) colour)
    | otherwise = branch -- HERE: we don't need to recreate the node
do_add Leaf val = (Branch Leaf val Leaf Black)

It also can give a small performance boost as we don't re-create objects identical to those we pattern match on.

Concerning fix_up: Let's try to factor out common and duplicate code. The common pattern is that we check some conditions on the sub-nodes of a node and if it's true, we apply a function or it. Otherwise we keep it intact. We can split this idea into two functions - one that is general, and other that is then specialized for branches:

fix_up :: Tree a -> Tree a
fix_up =
    onBr [is_red]                         [is_red] flip_colours .
    onBr [is_red, is_red . get_left_node] []       rotate_right .
    onBr [not . is_red]                   [is_red] rotate_left
  where
    -- Apply a function on a branch, if its left and right subnodes match
    -- given predicates.
    onBr :: [Tree a -> Bool] -> [Tree a -> Bool]
         -> (Tree a -> Tree a) -> (Tree a -> Tree a)
    onBr lps rps = on (\b -> all ($ get_left_node b) lps
                          && all ($ get_right_node b) rps)
    -- Apply a function on a value, if it matches a predicate.
    on :: (a -> Bool) -> (a -> a) -> (a -> a)
    on p f x | p x       = f x
             | otherwise = x

This allows us to represent the whole operation as the composition of several functions. And it localizes bindings of subnodes (l and r) to the predicates, which again helps readability.

---

I'd also like to draw your attention to AA trees. They have very similar performance as red-black trees, but are simpler and easier to implement.