Ameisen (Haskell) (Vorbereitung)

import Data.Array
import Control.Concurrent.STM
import Control.Concurrent
import Control.Monad ( forM, forever, void )
import System.Random
import System.IO

type Pos = (Int,Int)
type Dir = Int

data Ant = Ant { position :: TVar Pos 
               , direction :: TVar Dir
               , moves :: TVar Int  
               }

data Cell = Cell { occupied :: TVar Bool
                 }
type Board = Array Pos Cell

data World = World
           { size :: Pos 
           , board :: Board
           , population :: [ Ant ]
           , generator :: TVar StdGen 
           }

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

main :: IO ()
main = do
    w <- make_world (20,20) 10
    forM ( population w ) $ \ ant -> 
        forkIO $ forever $ walk w ant
    forever $ do
        pos <- snapshot w
        threadDelay $ 10^6 -- microseconds

info :: Ant -> STM String
info ant = do
    pos <- readTVar $ position ant
    dir <- readTVar $ direction ant
    mov <- readTVar $ moves ant
    return $ unwords [ "pos", show pos, "dir", show dir, "moves", show mov ]

snapshot :: World -> IO ()
snapshot w = do
  infos <- atomically $ forM ( population w ) $ info
  putStrLn $ unlines infos
  
---------------------------------------------------------------------------

-- | verschiebe in gegebene Richtung, 
-- mit wrap-around an den Rändern (d.h. Torus)
shift :: (Int,Int) -> Pos -> Dir -> Pos
shift (width,height) (x,y) d =
    let (dx,dy) = vector d
    in  ( mod (x+dx) width, mod (y+dy) height )

vector :: Dir -> Pos
vector d = case mod d 8 of
    0 -> ( 1,0) ; 1 -> ( 1, 1) ; 2 -> (0, 1) ; 3 -> (-1, 1)
    4 -> (-1,0) ; 5 -> (-1,-1) ; 6 -> (0,-1) ; 7 -> ( 1,-1)

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

randomRT :: Random a => TVar StdGen -> (a,a) -> STM a
randomRT ref bnd = do
    g <- readTVar ref
    let (x, g') = randomR bnd g
    writeTVar ref g'    
    return x

random_selection :: TVar StdGen -> Int -> [a] -> STM [a]
random_selection ref 0 xs = return []
random_selection ref k xs = do
    ( pre, y : post ) <- random_split ref xs
    ys <- random_selection ref (k-1) ( pre ++ post )
    return $ y : ys

random_split :: TVar StdGen -> [a] -> STM ( [a], [a] )
random_split ref xs = do
    k <- randomRT ref ( 0, length xs - 1 )
    return $ splitAt k xs

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

make_world :: (Int,Int) -> Int -> IO World
make_world (width,height) num_ants = do
    b <- make_board (width, height)        
    gen <- newStdGen
    ref <- atomically $ newTVar gen    
    ants <- make_ants ref b num_ants
    return $ World 
        { size = (width, height), board = b, population = ants 
        , generator = ref
        }

make_board :: (Int,Int) -> IO Board
make_board (width,height) = do    
    let bnd = ((0,0),(width-1,height-1))
    cells <- forM ( range bnd ) $ \ xy -> do
          occ <- atomically $ newTVar False
          return (xy, Cell { occupied = occ } )
    return $ array bnd cells
    
make_ants :: TVar StdGen -> Board -> Int -> IO [ Ant ]    
make_ants ref b num_ants = atomically $ do
    sel <- random_selection ref num_ants $ indices b
    forM sel $ \ pos -> do
        p <- newTVar pos
        enter $ b ! pos
        dir <- randomRT ref ( 0, 7 )
        d <- newTVar dir
        m <- newTVar 0
        return $ Ant { position = p, direction = d, moves = m }