packages feed

comonad-coactions-0.1.0.1: examples/Life.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE QualifiedDo #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE NoStarIsType #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver -threaded -with-rtsopts=-N -O2 #-}

module Main (main) where

import Control.Applicative
import Control.Comonad.Coaction
import Control.Comonad.Coaction.Right qualified as RC
import Control.Comonad.Identity hiding ((<@>))
import Control.Comonad.Store hiding (pos, (<@>))
import Control.Concurrent
import Control.Monad
import Control.Monad.Action.Left qualified as LA
import Control.Monad.Action.Right qualified as RA
import Control.Monad.State
import Data.Bits
import Data.Char
import Data.Constraint
import Data.Distributive
import Data.Finite
import Data.Foldable
import Data.Functor ((<&>))
import Data.Functor.Rep
import Data.Maybe
import Data.Proxy
import Data.Text qualified as T
import Data.Vector.Sized qualified as V
import Data.Word
import GHC.TypeLits
import Ki qualified
import Options.Applicative qualified as Opt
import System.Clock
import System.IO
import System.Random hiding (Finite)
import Termbox.Tea qualified as TB
import Text.Read hiding (get)
import Prelude hiding (replicate)

newtype Array2D (m :: Nat) (n :: Nat) a = Array2D {getArray2D :: V.Vector (m * n) a}
  deriving
    ( Eq,
      Ord,
      Show,
      Functor,
      Foldable,
      Traversable
    )

instance (KnownNat (m * n)) => Distributive (Array2D m n) where
  distribute = Array2D . distribute . fmap getArray2D

instance (KnownNat m, KnownNat n) => Representable (Array2D m n) where
  type Rep (Array2D m n) = (Finite m, Finite n)
  index Array2D{getArray2D} = index getArray2D . combineProduct
  tabulate f = Array2D . tabulate $ f . separateProduct

replicate :: (KnownNat (m * n)) => a -> Array2D m n a
replicate = Array2D . V.replicate

imap :: (KnownNat m) => ((Finite m, Finite n) -> a -> b) -> Array2D m n a -> Array2D m n b
imap f = Array2D . V.imap (f . separateProduct) . getArray2D

type Cell = Word8

cellArray :: (KnownNat m) => Settings -> Array2D m n Cell -> TB.Image
cellArray Settings{} =
  fold
    . imap
      \(i, j) c ->
        case c of
          0 -> mempty
          _ ->
            TB.fg (TB.color . fromIntegral . (`mod` 216) . (+ 100) . (* 2) $ c)
              . TB.atRow (fromInteger $ getFinite i)
              . TB.atCol (fromInteger $ 2 * getFinite j)
              . ap mappend (TB.atCol 1)
              $ TB.char '█'

pattern Store :: (s -> a) -> s -> Store s a
pattern Store a b = StoreT (Identity a) b

{-# COMPLETE Store #-}

wrapBoundary :: forall m n. (KnownNat m, KnownNat n) => Surface -> Integer -> Integer -> Maybe (Finite m, Finite n)
wrapBoundary Rectangle a b = (,) <$> packFinite a <*> packFinite b
wrapBoundary Torus a b = Just (modulo a, modulo b)
wrapBoundary Cylinder a b = packFinite a <&> (,modulo b)
wrapBoundary Moebius a b = do
  a' <- packFinite a
  let b' = modulo b
  Just (case packFinite @n b of Nothing -> -a'; _ -> a', b')
wrapBoundary Klein a b =
  let a' = modulo a
      b' = modulo b
   in Just (case packFinite @n b of Nothing -> -a'; _ -> a', b')
wrapBoundary Projective a b =
  let a' = modulo a
      b' = modulo b
   in Just (case packFinite @n b of Nothing -> -a'; _ -> a', case packFinite @m a of Nothing -> -b'; _ -> b')
wrapBoundary Sphere a b =
  let m = natVal $ Proxy @m
      a' = a `mod` (2 * m)
      b' = b `mod` (2 * m)
   in if
        | a' < m && b' < m -> Just (modulo a', modulo b')
        | a' >= m && b' < m -> Just (modulo b', modulo $ -a' - 1)
        | a' < m && b' >= m -> Just (modulo $ -b' - 1, modulo a')
        | otherwise -> Just (modulo $ -b' - 1, modulo $ -a' - 1)

life :: (KnownNat m, KnownNat n) => Settings -> Array2D m n Cell -> Array2D m n Cell
life Settings{rule = Rule{..}, ..} = rextend $ \s@(Store grid (i, j)) ->
  let t = sum
        $ LA.do
          a <- [-radius .. radius]
          b <-
            let r = case neighbourhood of
                  Moore -> radius
                  VonNeumann -> radius - abs a
                  Circular -> ceiling @Double . sqrt . fromIntegral $ radius * radius - a * a
                  Cross -> if a == 0 then radius else 0
                  Custom _ -> radius
             in [-r .. r]
          (a', b') <- wrapBoundary surface (getFinite i + fromIntegral a) (getFinite j + fromIntegral b)
          let weight = case neighbourhood of
                Custom c ->
                  let ix = (((radius - b) * (2 * radius + 1)) + radius - a)
                      centre = radius * (2 * radius + 1) + radius
                   in fromEnum $ ix /= centre && if ix < centre then c `testBit` ix else c `testBit` (ix - 1)
                _ -> fromEnum $ (a, b) /= (0, 0)
          pure @[] . (* weight) . fromEnum . (== 1) $ grid (a', b')
   in case extract s of
        0 -> if birth `testBit` t then 1 else 0
        1 -> if survival `testBit` t then 1 else 2 `mod` nStates
        x -> succ x `mod` nStates

data SomeBoard where SomeBoard :: (KnownNat m, KnownNat n) => SNat m -> SNat n -> Array2D m n Cell -> SomeBoard

data LifeState = LifeState
  { board :: !SomeBoard,
    running :: !Bool,
    time :: !TimeSpec,
    finished :: !Bool,
    delay :: !TimeSpec,
    drawing :: !(Maybe Cell),
    steps :: !Int
  }

data Neighbourhood = Moore | VonNeumann | Circular | Cross | Custom Integer

data Surface = Torus | Cylinder | Rectangle | Moebius | Klein | Projective | Sphere

data Rule = Rule
  { birth :: !Integer,
    survival :: !Integer,
    nStates :: !Word8,
    neighbourhood :: !Neighbourhood,
    radius :: !Int
  }

data Settings = Settings
  { rule :: Rule,
    surface :: !Surface
  }

snatDict :: SNat n -> Dict (KnownNat n)
snatDict sn = withKnownNat sn Dict

initialize :: Settings -> TimeSpec -> TB.Size -> LifeState
initialize Settings{..} time TB.Size{width, height} =
  let h = fromIntegral height
      w = fromIntegral $ width `div` 2
      (h', w') = case surface of
        Sphere -> (min h w, min h w)
        _ -> (h, w)
   in withSomeSNat w'
        $ withSomeSNat h'
        $ \case
          Nothing -> error "Unknown nat"
          Just sm -> \case
            Nothing -> error "Unknown nat"
            Just sn -> case (snatDict sm, snatDict sn) of
              (Dict, Dict) ->
                LifeState
                  { board = SomeBoard sm sn $ replicate 0,
                    running = False,
                    time,
                    finished = False,
                    delay = TimeSpec{sec = 0, nsec = 100_000_000},
                    drawing = Nothing,
                    steps = 0
                  }

pollEvent :: MVar TimeSpec -> Maybe (IO TimeSpec)
pollEvent m = Just $ takeMVar m

handleEvent :: Settings -> LifeState -> TB.Event TimeSpec -> IO LifeState
handleEvent settings@Settings{rule = Rule{..}} s@(LifeState{board = SomeBoard (sm :: SNat m) (sn :: SNat n) b, ..}) =
  \case
    TB.EventKey (TB.KeyChar 'r') ->
      do
        randomBoard <- sequence . replicate $ fmap (`mod` 2) randomIO
        pure $ s{board = SomeBoard sm sn randomBoard, steps = 0}
    TB.EventKey (TB.KeyChar 'c') -> pure s{board = SomeBoard sm sn $ replicate 0, running = False, steps = 0}
    TB.EventKey (TB.KeyChar 'q') -> pure s{finished = True}
    TB.EventKey (TB.KeyChar '+') -> pure s{delay = max 0 $ delay - TimeSpec{sec = 0, nsec = 20_000_000}}
    TB.EventKey (TB.KeyChar '-') -> pure s{delay = delay + TimeSpec{sec = 0, nsec = 20_000_000}}
    TB.EventKey TB.KeySpace -> pure s{running = not running}
    TB.EventMouse TB.Mouse{button = TB.LeftClick, pos = TB.Pos{..}} ->
      let r = modulo $ fromIntegral row
          c = modulo $ fromIntegral $ col `div` 2
       in case drawing of
            Nothing ->
              let b' = b RC.=>> \(Store grid (i, j)) -> if (i, j) == (r, c) then (grid (i, j) + 1) `mod` nStates else grid (i, j)
               in pure s{board = SomeBoard sm sn b', drawing = Just $ index b' (r, c)}
            Just cell ->
              let b' = b RC.=>> \(Store grid (i, j)) -> if (i, j) == (r, c) then cell else grid (i, j)
               in pure s{board = SomeBoard sm sn b'}
    TB.EventMouse TB.Mouse{button = TB.ReleaseClick} -> pure s{drawing = Nothing}
    TB.EventUser t -> if t - time >= delay && running then pure s{board = SomeBoard sm sn $ life settings b, time = t, steps = steps + 1} else pure s
    _ -> pure s

render :: Settings -> LifeState -> TB.Scene
render settings = TB.image . (\(SomeBoard _ _ b) -> cellArray settings b) . board

type Parser = StateT T.Text Maybe

getT :: Parser T.Text
getT = get

putT :: T.Text -> Parser ()
putT = put

satisfy :: (Char -> Bool) -> Parser Char
satisfy p = LA.do
  t <- getT
  (c, t') <- T.uncons t
  putT t'
  if p c then pure c else empty

parseChar :: Char -> Parser ()
parseChar = void . satisfy . (==)

parseNat :: (Integral a, Read a) => Parser a
parseNat = LA.do
  d <- some $ satisfy isDigit
  n <- readMaybe d
  pure $ fromInteger n

parseRange :: (Integral a, Read a) => Parser [a]
parseRange = LA.do
  m <- some (satisfy isDigit) RA.>>= readMaybe
  parseChar '-'
  n <- some (satisfy isDigit) RA.>>= readMaybe
  pure [fromInteger m .. fromInteger n]

eof :: Parser ()
eof = do
  t <- getT
  unless (T.null t) empty

sepBy :: Parser a -> Parser b -> Parser [b]
sepBy sep p = liftM2 (:) p (many (sep *> p)) <|> pure []

parseNbhd :: Parser Neighbourhood
parseNbhd =
  (Moore <$ parseChar 'M')
    <|> (VonNeumann <$ parseChar 'N')
    <|> (Circular <$ parseChar 'C')
    <|> (Cross <$ parseChar '+')
    <|> ( LA.do
            parseChar '@'
            s <- ("0x" ++) <$> many (satisfy isHexDigit)
            n <- readMaybe s
            pure $ Custom n
        )

-- | Higher-range outer totalistic notation for larger than life rules.
parseHROT :: Parser Rule
parseHROT = do
  parseChar 'R'
  radius <- parseNat
  parseChar ','
  parseChar 'C'
  nStates <- parseNat
  parseChar ','
  parseChar 'S'
  survival <- fmap (sum . fmap bit . join) . sepBy (parseChar ',') $ parseRange <|> fmap pure parseNat
  parseChar ','
  parseChar 'B'
  birth <- fmap (sum . fmap bit . join) . sepBy (parseChar ',') $ parseRange <|> fmap pure parseNat
  neighbourhood <- optional $ parseChar ',' *> parseChar 'N' *> parseNbhd
  eof
  pure Rule{neighbourhood = fromMaybe Moore neighbourhood, ..}

-- Birth/survival/states for generations rules.
parseBSC :: Parser Rule
parseBSC = do
  parseChar 'B'
  birth <- fmap (sum . fmap (bit . subtract (fromEnum '0') . fromEnum)) $ many $ satisfy isDigit
  parseChar '/'
  parseChar 'S'
  survival <- fmap (sum . fmap (bit . subtract (fromEnum '0') . fromEnum)) $ many $ satisfy isDigit
  nStates <- fmap (fromMaybe 2) . optional $ parseChar '/' >> optional (parseChar 'C') >> parseNat
  eof
  pure Rule{neighbourhood = Moore, radius = 1, ..}

-- Survival/birth/states for generations rules.
parseSBC :: Parser Rule
parseSBC = do
  survival :: Integer <- fmap (sum . fmap (bit . subtract (fromEnum '0') . fromEnum)) $ many $ satisfy isDigit
  parseChar '/'
  birth <- fmap (sum . fmap (bit . subtract (fromEnum '0') . fromEnum)) $ many $ satisfy isDigit
  nStates <- fmap (fromMaybe 2) . optional $ parseChar '/' >> parseNat
  eof
  pure Rule{neighbourhood = Moore, radius = 1, ..}

parseRulestring :: Parser Rule
parseRulestring = parseHROT <|> parseBSC <|> parseSBC

parseSettings :: Opt.ParserInfo (Maybe T.Text, Maybe Surface)
parseSettings =
  let parser =
        (,)
          <$> ( optional . Opt.strOption
                  $ Opt.short 'r' <> Opt.long "rule" <> Opt.metavar "RULESTRING" <> Opt.help "Rule string"
              )
          <*> Opt.optional
            ( Opt.flag' Rectangle (Opt.long "rectangle" <> Opt.help "Run cellular automaton in a rectangle (topologically a disk)")
                <|> Opt.flag' Torus (Opt.long "torus" <> Opt.help "Run cellular automaton in a torus")
                <|> Opt.flag' Cylinder (Opt.long "cylinder" <> Opt.help "Run cellular automaton in a cylinder")
                <|> Opt.flag' Moebius (Opt.long "moebius" <> Opt.help "Run cellular automaton in a Moebius strip")
                <|> Opt.flag' Klein (Opt.long "klein" <> Opt.help "Run cellular automaton in a Klein bottle")
                <|> Opt.flag' Projective (Opt.long "projective" <> Opt.help "Run cellular automaton in a real projective plane with singular points at the corners (orbifold symbol 22×)")
                <|> Opt.flag' Sphere (Opt.long "sphere" <> Opt.help "Run cellular automaton in a sphere with singular points at the corners (orbifold symbol 442)")
            )
   in Opt.info (parser Opt.<**> Opt.helper) (Opt.fullDesc <> Opt.progDesc "Larger than life cellular automaton")

main :: IO ()
main = do
  (mRulestring, mSurface) <- Opt.execParser parseSettings
  let rule = case mRulestring of
        Nothing -> Rule{radius = 1, neighbourhood = Moore, birth = 8, survival = 12, nStates = 2} -- Default rule: Conway's life
        Just rulestring -> fromMaybe (error "Failed to parse rulestring") $ evalStateT parseRulestring rulestring
  let surface = fromMaybe Rectangle mSurface
  t0 <- getTime Monotonic
  let settings =
        Settings{..}
  result <-
    Ki.scoped $ \scope -> do
      timeVar <- newEmptyMVar
      Ki.fork_ scope
        . forever
        $ threadDelay 1000
          >> getTime Monotonic
          >>= putMVar timeVar
      TB.run
        TB.Program
          { initialize = initialize settings t0,
            pollEvent = pollEvent timeVar,
            handleEvent = handleEvent settings,
            render = render settings,
            finished
          }
  case result of
    Left err -> hPutStrLn stderr $ "Failed to initialize: " ++ show err
    Right LifeState{steps} -> putStrLn $ "Ran for " ++ show steps ++ " steps"