packages feed

comonad-coactions-0.1.0.0: examples/Life.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE QualifiedDo #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE NoStarIsType #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver -threaded -rtsopts #-}

module Main (main) where

import Control.Comonad.Coaction
import Control.Comonad.Coaction.Right qualified as R
import Control.Comonad.Identity hiding ((<@>))
import Control.Comonad.Store hiding (pos, (<@>))
import Control.Concurrent
import Control.Monad
import Data.Bifunctor
import Data.Constraint
import Data.Distributive
import Data.Finite
import Data.Foldable
import Data.Functor.Rep
import Data.Vector.Sized qualified as V
import GHC.TypeLits
import Ki qualified
import System.Clock
import System.IO
import System.Random hiding (Finite)
import Termbox.Tea qualified as TB
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 = Bool

pattern Off :: Cell
pattern Off = False

pattern On :: Cell
pattern On = True

{-# COMPLETE On, Off #-}

cellArray :: (KnownNat m) => Array2D m n Cell -> TB.Image
cellArray =
  fold
    . imap
      ( \(i, j) c ->
          if c
            then
              TB.atRow (fromInteger $ getFinite i)
                . TB.atCol (fromInteger $ 2 * getFinite j)
                . ap mappend (TB.atCol 1)
                $ TB.char '█'
            else mempty
      )

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

{-# COMPLETE Store #-}

life :: (KnownNat m, KnownNat n, KnownNat (m * n)) => Array2D m n Cell -> Array2D m n Cell
life = rextend $ \(Store grid (i, j)) ->
  let s =
        sum . fmap (fromEnum . grid . bimap modulo modulo) $
          (,)
            <$> fmap (getFinite i +) [-1, 0, 1]
            <*> fmap (getFinite j +) [-1, 0, 1]
   in s == 3 || (s == 4 && grid (i, j))

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,
    done :: !Bool,
    delay :: !TimeSpec,
    drawing :: !(Maybe Cell),
    steps :: !Int
  }

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

initialize :: TimeSpec -> TB.Size -> LifeState
initialize time = \TB.Size {width, height} ->
  withSomeSNat (fromIntegral $ width `div` 2) $
    withSomeSNat (fromIntegral height) $
      \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 Off,
                  running = True,
                  time,
                  done = 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 :: LifeState -> TB.Event TimeSpec -> IO LifeState
handleEvent s@(LifeState {board = SomeBoard (sm :: SNat m) (sn :: SNat n) b, ..}) =
  \case
    TB.EventKey (TB.KeyChar 'r') ->
      do
        randomBoard <- sequence $ replicate randomIO
        pure $ s {board = SomeBoard sm sn randomBoard, steps = 0}
    TB.EventKey (TB.KeyChar 'c') -> pure $ s {board = SomeBoard sm sn $ replicate Off, steps = 0}
    TB.EventKey (TB.KeyChar 'q') -> pure $ s {done = 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 R.=>> \(Store grid (i, j)) -> if (i, j) == (r, c) then not $ grid (i, j) else grid (i, j)
               in pure s {board = SomeBoard sm sn b', drawing = Just $ index b' (r, c)}
            Just cell ->
              let b' = b R.=>> \(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 b, time = t, steps = steps + 1} else pure s
    _ -> pure s

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

main :: IO ()
main = do
  t0 <- getTime Monotonic
  result <-
    Ki.scoped $ \scope -> do
      timeVar <- newEmptyMVar
      Ki.fork_ scope
        . forever
        $ threadDelay 1000
          >> getTime Monotonic
          >>= putMVar timeVar
      TB.run
        TB.Program
          { initialize = initialize t0,
            pollEvent = pollEvent timeVar,
            handleEvent,
            render,
            finished = done
          }
  case result of
    Left err -> hPutStrLn stderr $ "Failed to initialize: " ++ show err
    Right LifeState {steps} -> putStrLn $ "Ran for " ++ show steps ++ " steps"