Automata.hs

{-# LANGUAGE QuasiQuotes #-}

import Control.Monad
import Data.Array.Repa (Array, DIM2, All(..), Any(..), (:.)(..), U)
import qualified Data.Array.Repa as R
import Data.Array.Repa.Stencil
import Data.Array.Repa.Stencil.Dim2
import Data.Monoid (Endo(..), mappend, mconcat)
import System.Random

type Dungeon = Array U DIM2 Int

initGen :: Int -> Int -> StdGen -> Dungeon
initGen x y = R.fromListUnboxed (R.ix2 x y) . take (x * y) . randomRs (0, 1)

sten1 :: Stencil DIM2 Int
sten1 = [stencil2| 1 1 1
                   1 0 1
                   1 1 1 |]

sten2 :: Stencil DIM2 Int
sten2 = [stencil2| 1 1 1 1 1
                   1 1 1 1 1
                   1 1 0 1 1
                   1 1 1 1 1
                   1 1 1 1 1 |]

step :: (Int -> Int -> Int) -> Stencil DIM2 Int -> Dungeon -> Dungeon
{-# INLINE step #-}
step rule sten dung = R.computeUnboxedS $ R.zipWith rule dung appSten
    where {-# INLINE appSten #-}
          appSten = mapStencil2 (BoundConst 1) sten dung

iterStep :: Endo Dungeon
iterStep = Endo $ step rule sten1
    where {-# INLINE rule #-}
          rule 0 n = if n > 4 then 1 else 0
          rule _ n = if n > 3 then 1 else 0

smoothStep :: Endo Dungeon
smoothStep = Endo $ step rule sten2
    where {-# INLINE rule #-}
          rule 0 _ = 0
          rule _ 0 = 0
          rule _ _ = 1

display :: Dungeon -> IO ()
display dung = forM_ [0..nRows] $ \i ->
    print . R.toList . R.map pretty . R.slice dung $ Any :. i :. All
    where {-# INLINE pretty #-}
          pretty 0 = '.'
          pretty _ = '#'
          nRows = (R.listOfShape (R.extent dung) !! 1) - 1

main :: IO ()
main = newStdGen >>= display . steps . initGen 50 100
    where Endo steps = smoothStep `mappend` mconcat (replicate 3 iterStep)