game of life

GameOfLife.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE BangPatterns #-}

module Main where

import qualified System.Console.Terminal.Size as Term (size,width,height)
import Data.Maybe (fromMaybe)
import Data.Vector (Vector,(!))
import qualified Data.Vector as V
import Data.Text (Text)
import qualified Data.Text as T (cons,empty,append)
import qualified Data.Text.IO as T (putStr)
import Data.Foldable (traverse_)
import Control.Concurrent (threadDelay)
import qualified Data.ByteString.Char8 as B (putStr)
import System.Random (randomIO)

type Grid  = Vector (Vector State)
data State = Alive | Dead deriving Eq

main :: IO ()
main = do
    size <- Term.size
    let width  = fromMaybe 10 $ Term.width  <$> size
        height = fromMaybe 10 $ Term.height <$> size
    randomGrid (width,height) >>= traverse_ tickDisplay . iterate updateGrid
    where
        tickDisplay xss = printGrid xss >> delay >> clearScreen
        delay           = threadDelay 250000
        clearScreen     = B.putStr "\ESC[2J"

updateGrid :: Grid -> Grid
updateGrid grid = V.imap (\i1 -> V.imap $ getState i1) grid
    where
        getState x y state = aliveOrDead aliveNeighbors state
            where
                aliveNeighbors = count Alive $ neighbors x y
        count x            = length . V.filter (== x)
        aliveOrDead x state
            | state == Dead  &&  x == 3            = Alive
            | state == Alive && (x == 2 || x == 3) = Alive
            | otherwise                            = Dead
        neighbors x y      = (\(x',y') -> grid ! x' ! y') <$> correct <$>
            V.fromList
                [(x+1,y)  ,(x-1,y)  ,(x,y+1)  ,(x,y-1)
                ,(x+1,y+1),(x-1,y-1),(x-1,y+1),(x+1,y-1)]
        correct (x,y) = (axisCorrect xSize x,axisCorrect ySize y)
        axisCorrect axisSize n
            | n == -1       = axisSize-1
            | n == axisSize = 0
            | otherwise     = n
        xSize              = length grid
        ySize              = length $ V.head grid

printGrid :: Grid -> IO ()
printGrid xss = T.putStr . V.foldr1 T.append $
    ((`T.append` "\n") . vectorToText) <$> (fmap . fmap) convert transposed
    where
        convert s  = if s == Alive then 'O' else '-'
        transposed = f 0 ig V.cons xss (>= size) V.empty
        f !n g h x p z
            | p n       = z
            | otherwise = g n x `h` f (n+1) g h x p z
        ig n xs    = (! n) <$> xs
        size       = length $ V.head xss

vectorToText :: Vector Char -> Text
vectorToText = V.foldr T.cons T.empty

randomGrid :: (Int,Int) -> IO Grid
randomGrid (x,y) = V.replicateM x $ V.replicateM y randomState

randomState :: IO State
randomState = do
    r <- randomIO :: IO Float
    pure $ if r > 0.9 then Alive else Dead