Binary trees with type-enforced balance

A Red-black tree is a binary tree which has some properties which assure that it is balanced.

The properties of the RB-Tree can be described using Haskell's type system, thus creating trees that must be balanced.

The compiler will catch any bug in the code handling them which causes the tree to not be balanced - it will not type-check.

The properties of Red-Black trees:

• Each node in the tree has a "color", either Red or Black.
• All different paths from the root of the tree all the way down have the same number of Black nodes.
• No path in the tree has two consecutive Red nodes.

This assures that the ratio between the maximum and minimum depth will be no more than two. The minimum depth will be achieved by a path of Black nodes with no Red nodes between them, while the maximum depth would have as many Red nodes possible crammed between the Black nodes - at most one Red node between each two Black ones.

In code:

{-# LANGUAGE EmptyDataDecls, ExistentialQuantification, GADTs #-}
data Zero
data Succ n

-- We use trees whose root is Black. Simply for implementation convinience.
data Tree a = forall n. Tree (BlackNode n a)
-- Node types are tagged by their subtree's Black-degree (number of Black nodes per path).
data BlackNode n a where
    __-- A degree 0 Black node must be an empty tree.
__Nil :: BlackNode Zero a
__BlackNode :: RBNode n a -> a -> RBNode n a -> BlackNode (Succ n) a
data RBNode n a
    __= ItsRed (RedNode n a)
    __| ItsBlack (BlackNode n a)
-- A Red node's children must be Black.
data RedNode n a = RedNode (BlackNode n a) a (BlackNode n a)

The types for the tree nodes are all indexed by the number of Black nodes per path / "RB-Depth".
The outward facing type, "Tree", wraps a node type using an Existential for the depth, so the depth could be any value and the user of the type would not need to care.

While the type captures the structural rules of the tree, it does not capture the conditions required by binary search trees. I suppose that might be possible in fancier languages such as Agda, and perhaps I'll give it a try :)

Complete working code for the tree, including the insert function, at github.com/yairchu/red-black-tree/blob/master/RedBlackTree.hs

Charlemagne, Disraeli, and Jefferson didn't have undo.

James McNeill translated the ancient computer game "Hamurabi" from BASIC to Haskell.

This article will discuss the ease of adding a simple "undo" feature to the modular Haskell implementation with the List monad transformer.

The version in BASIC was 121 lines long, while the Haskell version is 261 lines. At first this seems bad for Haskell. Let's analyze it a bit more:

Without imports, type definitions, and blank lines, the Haskell version weights at 200 lines. The BASIC program's lines are more dense - line numbers are used as labels, while in Haskell one tends to give a full line for the name of a do-block function.

All in all the difference doesn't seem very big.

The Haskell code has one big advantage: it is modular. I'll demonstrate by showing the diff to adding a simple "undo" mechanism using backtracking with the list monad transformer from my generator (list monad and related utilities) package.

The program (with annotated diffs in big font):

{-# LANGUAGE FlexibleContexts, TemplateHaskell #-}
-- to compile: ghc --make hammurabi.hs -package transformers -hide-package mtl

-- Converted from the original FOCAL program and modified for Edusystem 70 by David Ahl, Digital
-- Modified for 8K Microsoft BASIC by Peter Turnbull
-- Ported to Haskell by James McNeill
-- Refactored to use MaybeT and StateT by Yair Chuchem

import Control.Monad (forM, forever, join, liftM, liftM2, liftM3, mzero, when)
import Control.Monad.ListT
import Control.Monad.Trans
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.State (StateT, runStateT)
import Data.Accessor.MonadState (get, modify, set)
import Data.Accessor.Template
import Data.Function
import Data.List.Class (execute, repeat, takeWhile)
import Data.Ratio
import Random
import IO
import Text.Printf

import Prelude hiding (repeat, takeWhile)

data GameState = GameState {
 year_         :: Integer,
 people_       :: Integer,
 food_         :: Integer, -- bushels
 land_         :: Integer, -- acres
 landPrice_    :: Integer, -- bushels per acre
 totalDeaths_  :: Integer,
 cumDeathRate_ :: Rational,
 rng_          :: StdGen
} deriving (Show)
$(deriveAccessors ''GameState)

data Orders = Orders {
 acresToPlant   :: Integer,
 bushelsForFood :: Integer
} deriving (Show)

data Results = Results {
 peopleStarved      :: Integer,
 peopleDiedOfPlague :: Integer,
 peopleBorn         :: Integer,
 bushelsEatenByRats :: Integer,
 bushelsPerAcre     :: Integer
} deriving (Show)

main :: IO ()
main = do
 putStrLn "Try your hand at governing ancient Sumeria successfully for a 10-year term of\noffice."
 start <- liftM initialState getStdGen
-- Before Undo:
 runMaybeT . (`runStateT` start) $ do
-- With Undo:
 runMaybeT . execute . (`runStateT` start) $ do
-- Same:
    showResults initialResults
   forever mainLoopIter
 putStrLn "So long for now."

-- Before Undo:
type Hammurabi = StateT GameState (MaybeT IO)

quit :: Hammurabi a
quit = mzero
-- With Undo:
type Hammurabi = StateT GameState (ListT (MaybeT IO))

-- ListT is an instance of mtl's MonadTrans.
-- here we have to use transformers' because of using Data.Accessor
instance MonadTrans ListT where
 lift = ListT . liftM (`Cons` mzero)
instance MonadIO m => MonadIO (ListT m) where
 liftIO = lift . liftIO

quit :: Hammurabi a
quit = lift $ lift mzero
-- Same:

mainLoopIter :: Hammurabi ()
mainLoopIter = do
 set landPrice =<< useRandom (randomR (17, 26))
 peopleStart <- get people
 resultsOut <- applyOrders =<< readOrders
 showResults resultsOut
 let numStarved = peopleStarved resultsOut
 when (numStarved % peopleStart > 0.45) $ do
   liftIO $ printf "You starved %d people in one year!\n%s" numStarved finkMessage
   quit
 join . liftM2 (when . (== 10)) (get year) . return $ do
   finalReport
   quit

useRandom :: (StdGen -> (a, StdGen)) -> Hammurabi a
useRandom func = do
 (res, rngOut) <- fmap func $ get rng
 set rng rngOut
 return res

showResults :: Results -> Hammurabi ()
showResults results = do
 liftIO $ printf "\nHamurabi:  I beg to report to you,\n"
 y <- get year
 liftIO $ printf "In year %d, %d people starved, %d came to the city.\n" y (peopleStarved results) (peopleBorn results)
 when (peopleDiedOfPlague results > 0) . liftIO . putStrLn $ "A horrible plague struck!  Half the people died."
 liftIO . printf "Population is now %d.\n" =<< get people
 liftIO . printf "The city now owns %d acres.\n" =<< get land
 liftIO $ printf "You harvested %d bushels per acre.\n" (bushelsPerAcre results)
 liftIO $ printf "Rats ate %d bushels.\n" (bushelsEatenByRats results)
 liftIO . printf "You now have %d bushels in store.\n" =<< get food

-- :(
killPeople :: Integer -> Hammurabi ()
killPeople n = do
 modify people (+ negate n)
 modify totalDeaths (+ n)

applyOrders :: Orders -> Hammurabi Results
applyOrders orders = do
 modify year (+ 1)
 peopleInit <- get people
 harvestYield <- useRandom $ randomR (1, 6)
 let
   starved = max 0 $ peopleInit - (bushelsForFood orders `div` 20)
   calcEatenByRats r f
     | ((r `mod` 2) == 1) = 0
     | otherwise = f `div` r
   calcDiedOfPlague :: Double -> Integer -> Integer
   calcDiedOfPlague r p = if r >= 0.15 then 0 else p `div` 2
   calcBorn r l b = 1 + ((r * (20 * l + b)) `div` (peopleInit * 100))
   bushelsHarvested = harvestYield * acresToPlant orders
 killPeople starved
 modify cumDeathRate (+ starved % peopleInit)
 modify food (+ negate (acresToPlant orders `div` 2))
 modify food (+ bushelsHarvested)
 eatenByRats <- liftM2 calcEatenByRats (useRandom (randomR (1, 6))) (get food)
 modify food (+ negate eatenByRats)
 born <- liftM3 calcBorn (useRandom (randomR (1, 6))) (get land) (get food)
 modify people (+ born)
 diedOfPlague <- liftM2 calcDiedOfPlague (useRandom random) (get people)
 killPeople diedOfPlague
 return Results {
   peopleStarved = starved,
   peopleDiedOfPlague = diedOfPlague,
   peopleBorn = born,
   bushelsEatenByRats = eatenByRats,
   bushelsPerAcre = harvestYield
   }

readNum ::
 String -> String ->
 [(Hammurabi Integer, Hammurabi ())] ->
 Maybe Integer -> Hammurabi Integer
readNum units purpose limits idealN = do
 lims <-
   forM limits $ \(alim, amsg) -> do
     r <- alim
     return (r, amsg)
 let
   maxN = minimum $ map fst lims
   defaultN = maybe maxN (min maxN) idealN
 case maxN of
   0 -> return 0
   _ ->
     fix $ \resume -> do
-- Before Undo:
        liftIO . putStr $
         "How many " ++ units ++
         " do you wish to " ++ purpose ++
         " (0-" ++ show maxN ++ ")? [" ++ show defaultN ++ "] "
       liftIO $ hFlush stdout
       line <- liftIO getLine
-- With Undo:
        line <-
         lift . takeWhile (/= "undo") $ do
           repeat () -- backtracking would be to this point
           liftIO . putStr $
             "How many " ++ units ++
             " do you wish to " ++ purpose ++
             " (0-" ++ show maxN ++ ")? [" ++ show defaultN ++ "] "
           liftIO $ hFlush stdout
           liftIO getLine
-- Same:
        case maybeRead line of
         Nothing -> return defaultN
         Just n
           | n < 0 -> do
             liftIO $ putStrLn abortMessage
             quit
           | n < maxN -> do
             snd . head $ filter ((< n) . fst) lims
             resume
           | otherwise -> return n

maybeRead :: Read a => String -> Maybe a
maybeRead s =
 case reads s of
   [(x, str)] | all (== ' ') str -> Just x
   _ -> Nothing

readOrders :: Hammurabi Orders
readOrders = do
 let
   thinkAgain item suf = do
     x <- get item
     liftIO $ printf
       "Hammurabi: Think again. You have only %d %s.  Now then,\n"
       x suf
   thinkAgainFood = thinkAgain food "bushels of grain"
   tendFieldsMsg =
     liftIO . printf
       "But you have only %d people to tend the fields.  Now then,\n"
        =<< get people
 liftIO . printf "Land is trading at %d bushels per acre\n" =<< get landPrice
 buyLand <-
   readNum "acres" "buy"
   [(liftM2 div (get food) (get landPrice), thinkAgainFood)]
   (Just 0)
 sellLand <-
   readNum "acres" "sell"
   [(if buyLand > 0 then return 0 else get land, thinkAgain land "acres")]
   (Just 0)
 let landDiff = buyLand - sellLand
 modify land (+ landDiff)
 modify food . (+) . negate . (* landDiff) =<< get landPrice
 feed <-
   readNum "bushels" "feed your people"
   [(get food, thinkAgainFood)]
   . Just . (* 20) =<< get people
 modify food (+ negate feed)
 plant <-
   readNum "acres" "plant with seed"
   [(get land, thinkAgain land "acres")
   ,(fmap (* 2) (get food), thinkAgainFood)
   ,(fmap (* 10) (get people), tendFieldsMsg)]
   Nothing
 return $ Orders plant feed

finalReport :: Hammurabi ()
finalReport = do
 numYears <- get year
 avgDeathRate <- fmap (/ (numYears % 1)) $ get cumDeathRate
 numPeople <- get people
 acresPerPerson <- fmap (% numPeople) $ get land
 numHaters <- useRandom $ randomR (0, (numPeople * 4) `div` 5)
 let
   comments
     | avgDeathRate > 0.33 || acresPerPerson < 7 = finkMessage
     | avgDeathRate > 0.1 || acresPerPerson < 9 =
       "Your heavy-handed performance smacks of Nero and Ivan IV.\n" ++
       "The people (remaining) find you an unpleasant ruler, and,\n" ++
       "frankly, hate your guts!\n"
     | avgDeathRate > 0.03 || acresPerPerson < 10 =
       "Your performance could have been somewhat better, but\n" ++
       "really wasn't too bad at all. " ++
       show numHaters ++ " people would\n" ++
       "dearly like to see you assassinated but we all have our\n" ++
       "trivial problems.\n"
     | otherwise =
       "A fantastic performance!!!  Charlemagne, Disraeli, and\n" ++
       "Jefferson combined could not have done better!\n"
 td <- get totalDeaths
 liftIO . putStrLn $
   "In your " ++ show numYears ++ "-year term of office, " ++
   show (round ((100%1) * avgDeathRate) :: Integer) ++ " percent of the\n" ++
   "population starved per year on average, i.e., " ++
   "a total of " ++ show td ++ " people died!!\n" ++
   "You started with 10 acres per person and ended with " ++
   show (round acresPerPerson :: Integer) ++ " acres per person.\n" ++
   comments

initialState :: StdGen -> GameState
initialState r = GameState {
 year_ = 0,
 people_ = 100,
 food_ = 2800,
 land_ = 1000,
 landPrice_ = 0,
 totalDeaths_ = 0,
 cumDeathRate_ = 0,
 rng_ = r }

initialResults :: Results
initialResults = Results 0 0 5 200 3

abortMessage :: String
abortMessage = "Hammurabi: I cannot do what you wish!\nGet yourself another steward!!!!!"

finkMessage :: String
finkMessage =
 "Due to this extreme mismanagement you have not only\n" ++
 "been impeached and thrown out of office but you have\n" ++
 "also been declared 'National Fink' !!\n"

It's a nice game. If you play it - only use undo to correct typos!
This is a very basic undo. It doesn't support redo etc. I'm not endorsing this way of implementing undo. The purpose was just to show an example of what can be done with the list monad transformer.
There is a small problem with Haskell libraries. I had to named my module Control.Monad.ListT instead of Control.Monad.List because that one is taken by mtl (and imho doesn't provide what a list monad should). There are two monad-transformers packages (mtl and transformers) and both are being used. I had to use both because I chose to support mtl but then Data.Accessor uses transformers. Perhaps the solution is to break libraries to smaller parts. Then it would be easy for the community to agree on those parts. That way I could have also used the Control.Monad.List name. Perhaps it means I should break my library to parts as well.