Santa problem showing STM

Makefile

santa: santa.hs
	ghc $^ -package stm -o $@

run: santa
	./santa

clean:
	rm santa.hi santa.o santa

santa.hs

-- Santa makes one “Group” for the elves and one for the reindeer. Each elf (or reindeer) 
-- tries to join its Group. If it succeeds, it gets two “Gates” in return. The first Gate 
-- allows Santa to control when the elf can enter the study, and also lets Santa know
-- when they are all inside. Similarly, the second Gate controls the elves leaving
-- the study. Santa, for his part, waits for either of his two Groups to be ready, and
-- then uses that Group’s Gates to marshal his helpers (elves or reindeer) through
-- their task. Thus the helpers spend their lives in an infinite loop: try to join
-- a group, move through the gates under Santa’s control, and then delay for a
-- random interval before trying to join a group again.

import Control.Concurrent.STM
import Control.Concurrent
import System.Random
import Control.Monad

-- The main program
main = do elf_grp <- newGroup 3
          sequence_ [elf elf_grp n | n <- [1..10] ]
          rein_group <- newGroup 9
          sequence_ [reindeer rein_group n | n <- [1..9] ]
          forever (santa elf_grp rein_group)

-------------------------------------------------------------------
--santa simple solution
santa1 :: Group -> Group -> IO ()
santa1 elf_grp rein_grp = do 
    putStr "------------------\n"
    (task, (in_gate, out_gate)) <- atomically $ orElse 
        (chooseGroup rein_grp "deliver toys")
        (chooseGroup elf_grp "meet in study")
    putStr $ "Ho! Ho! Ho! Let's " ++ task ++ "\n"
    operateGate in_gate
    operateGate out_gate
    where 
        chooseGroup :: Group -> String -> STM (String, (Gate, Gate))
        chooseGroup grp task = awaitGroup grp >>= \gates -> return (task, gates)

-- alternative solution to form a pattern
santa elf_grp rein_grp = do
    putStr "------------------\n"
    choose [(awaitGroup rein_grp, run "deliver toys"),
            (awaitGroup elf_grp, run "meet in study")]
    where 
        run :: String -> (Gate, Gate) -> IO ()
        run task (in_gate, out_gate) = do
            putStr $ "Ho! Ho! Ho! Let's " ++ task ++ "\n"
            operateGate in_gate
            operateGate out_gate 

-- choose check for first STM action, and do the corresponding function if success, otherwise, 
--  it will check for the second STM action and do its function if succeed, and so on
--  if none of the actions succeed, just retry again
choose :: [(STM a, a -> IO())] -> IO ()
choose choices = do
    act <- atomically $ foldr1 orElse actions 
    act
    where
        actions :: [STM (IO())]
        actions = [ do { val <- guard; return (rhs val); } | (guard, rhs) <- choices]
-------------------------------------------------------------------
-- helper functions
helper1 :: Group -> IO () -> IO ()
helper1 group do_task = do 
    (in_gate, out_gate) <- joinGroup group
    passGate in_gate
    do_task
    passGate out_gate

-- elf
elf1 :: Group -> Int -> IO ()
elf1 gp elfId = helper1 gp (meetInStudy elfId)
elf gp elfId = forkIO (forever $ do {elf1 gp elfId;randomDelay;})

meetInStudy :: Int -> IO ()
meetInStudy elfId = putStr ("Elf " ++ show elfId ++ " meet in study")

-- reindeer 
reindeer1 :: Group -> Int -> IO ()
reindeer1 gp reId = helper1 gp (deliverToys reId)
reindeer gp reId = forkIO (forever $ do {reindeer1 gp reId;randomDelay;})

deliverToys :: Int -> IO ()
deliverToys reId = putStr ("Reindeer " ++ show reId ++ " delivering toys")

-- random delay to suspend the thread
randomDelay :: IO ()
randomDelay = getStdRandom (randomR (1, 1000000)) >>= \t -> threadDelay t

-------------------------------------------------------------------
-- A gate has capacity and a mutable capacity state
data Gate = MakeGate Int (TVar Int)

newGate :: Int -> STM Gate
newGate n = newTVar 0 >>= \tv -> return (MakeGate n tv)

passGate :: Gate -> IO ()
passGate (MakeGate n tv) = atomically $ do
    n_left <- readTVar tv
    check (n_left > 0)
    writeTVar tv (n_left - 1)

operateGate :: Gate -> IO ()
operateGate (MakeGate n tv) = do
    atomically $ writeTVar tv n
    atomically $ do
        n_left <- readTVar tv
        check (n_left == 0)

-------------------------------------------------------------------
-- A group has a capacity and a mutable state about remaining capacity, gates
data Group = MakeGroup Int (TVar (Int, Gate, Gate))
newGroup :: Int -> IO Group
newGroup n = atomically $ do
    g1 <- newGate n
    g2 <- newGate n
    tv <- newTVar (n, g1, g2)
    return (MakeGroup n tv)

joinGroup :: Group -> IO (Gate, Gate)
joinGroup (MakeGroup n tv) = atomically $ do
    (n_left, g1, g2) <- readTVar tv
    check (n_left > 0)
    writeTVar tv (n_left - 1, g1, g2)
    return (g1, g2)

-- Wait for a group until no free slots available and create new gates
awaitGroup :: Group -> STM (Gate, Gate)
awaitGroup (MakeGroup n tv) = do
    (n_left, g1, g2) <- readTVar tv
    check (n_left == 0)
    new_g1 <- newGate n
    new_g2 <- newGate n
    writeTVar tv (n, new_g1, new_g2)
    return (new_g1, new_g2)

To run the program, execute: ghc santa.hs -package stm -o santa, and pass in appropriate RTSparameters.

This example is elaborated in beautiful concurrency