STM で解くサンタ問題

Haskellers Meeting 2010 Springで、Simon さんに STM(Software Transaction Memory)の話をして頂くようにお願いしました。参加者の人があらかじめ予習できるように、参考書として Real World Haskell とビューティフルコードを挙げておきました。

Real World Haskell―実戦で学ぶ関数型言語プログラミング

ビューティフルコード (THEORY/IN/PRACTICE)

Real World Haskell の STM の章は読んだことがあったのですが、ビューティフルコードの方はなかったので、よい機会なので読んでみました。Real World HaskellのSTMの説明よりも、ビューティフルコードの方が分かりやすかったです。

ビューティフルコードの方は、例がいかしていて、最終的にSTMでサンタ問題を解きます。

サンタは、休暇から戻った9頭のトナカイたちが全員で起こすか、10人の小人たちのうちの3人がグループになって起こすかするまで、繰り返し眠っています。トナカイたちに起こされたら、9頭全員のトナカイにハーネスをつけてソリの準備をして、おもちゃを配りにいきます。終わったらハーネスをはずし、トナカイたちに休日をやります。小人たちのグループにおこされたら、そのグループの小人たちを書斎に連れて行き、おもちゃの調査や開発の相談をして、方針がまとまったら小人たちは仕事に戻ります。サンタは、小人たちとトナカイたちの両方が同時に起こしに来た場合は、トナカイたちを優先しなければなりません。

＃「行く」と「いく」とが揺れているなぁ。。。

理解するために、例題のコードを手で入力してみました。結局、このコードを理解するには、ゲートを2つ用意する必要性が分かるかにかかっています。

コード自体は Simon さんのから入手できるようですが、僕がレイアウトを変えながら入力し、Applicative スタイルに変更したコードを掲載しておきます。

module Main where

import Control.Applicative
import Control.Concurrent.STM
import Control.Concurrent
import System.Random

----------------------------------------------------------------

main :: IO ()
main = do
    elf_group <- newGroup 3
    sequence_ [elf elf_group n | n <- [1..10]]
    rein_group <- newGroup 9
    sequence_ [reindeer rein_group n | n <- [1..9]]
    forever $ santa elf_group rein_group
    
elf :: Group -> Int -> IO ThreadId
elf gp idnt = forkIO . forever $ do
    elf1 gp idnt
    randomDelay

reindeer :: Group -> Int -> IO ThreadId
reindeer gp idnt = forkIO . forever $ do
    reindeer1 gp idnt
    randomDelay

forever :: IO () -> IO ()
forever act = do
    act
    forever act

randomDelay :: IO ()
randomDelay = do
    waitTime <- getStdRandom . randomR $ (1, 1000000)
    threadDelay waitTime

santa :: Group -> Group -> IO ()
santa elf_gp rein_gp = do
    putStr "----------\n"
    choose [ (awaitGroup rein_gp, run "deliver toys")
           , (awaitGroup elf_gp,  run "meet in my 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 :: [(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 ]

----------------------------------------------------------------

data Gate = Gate Int (TVar Int)

newGate :: Int -> STM Gate
newGate n = Gate n <$> newTVar 0

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

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

----------------------------------------------------------------

data Group = Group Int (TVar (Int, Gate, Gate))

newGroup :: Int -> IO Group
newGroup n = atomically $ do
    g1 <- newGate n
    g2 <- newGate n
    Group n <$> newTVar (n, g1, g2)

joinGroup :: Group -> IO (Gate, Gate)
joinGroup (Group _ tv) = atomically $ do
    (n_left, g1, g2) <- readTVar tv
    check (n_left > 0)
    writeTVar tv (n_left - 1, g1, g2)
    return (g1,g2)
    
awaitGroup :: Group -> STM (Gate, Gate)
awaitGroup (Group 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 (g1, g2)

----------------------------------------------------------------

elf1 :: Group -> Int -> IO ()
elf1      gp idnt = helper1 gp (meetInStudy idnt)

reindeer1 :: Group -> Int -> IO ()
reindeer1 gp idnt = helper1 gp (deliverToys idnt)

helper1 :: Group -> IO () -> IO ()
helper1 group do_task = do
    (in_gate, out_gate) <- joinGroup group
    passGate in_gate
    do_task
    passGate out_gate

meetInStudy :: Int -> IO ()
meetInStudy idnt = putStr $ "Elf " ++ show idnt ++ " meeting in the study\n"

deliverToys :: Int -> IO ()
deliverToys idnt = putStr $ "Reindeer " ++ show idnt ++ " delivering toys\n"