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 here on SO).
- 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.