cellular automata

CellularUnrolled.hs

{-# language Strict, TypeApplications, ScopedTypeVariables,
    PartialTypeSignatures, ViewPatterns #-}
{-# options_ghc -Wno-partial-type-signatures #-}

import qualified Data.Primitive.PrimArray as A
import GHC.Exts (IsList(..))
import Data.Bits
import Data.Word
import Control.Monad.ST.Strict

type Rule = A.PrimArray Word8

mkRule :: Word8 -> Rule
mkRule w = fromList (map (fromIntegral . fromEnum . testBit w) [0..7])

step :: Rule -> A.PrimArray Word8 -> A.PrimArray Word8
step rule v = runST $ do
  let ix i = v `A.indexPrimArray` i
      {-# inline ix #-}

  let len = A.sizeofPrimArray v
      l   = ix 0
      r   = ix (len - 1)

  (v' :: _ s Word8) <- A.newPrimArray len

  let upd :: Word8 -> Word8 -> Word8 -> Int -> ST s ()
      upd (fromIntegral -> l) (fromIntegral -> c) (fromIntegral -> r) i = do
        let ruleIx = unsafeShiftL l 2 .|. unsafeShiftL c 1 .|. r
        A.writePrimArray @_ @(ST s) v' i (rule `A.indexPrimArray` ruleIx)
      {-# inline upd #-}

  upd r l (ix 1) 0
  upd (ix (len - 2)) r l (len - 1)

  let go i | i < len - 2 = do
        let a = ix (i - 1)
            b = ix i
            c = ix (i + 1)
            d = ix (i + 2)
        upd a b c i
        upd b c d (i + 1)
        go (i + 2)
      go i | i == len - 2 = do
        let a = ix (i - 1)
            b = ix i
            c = ix (i + 1)
        upd a b c i
      go i = pure ()

  go 1
  A.unsafeFreezePrimArray v'

runIterations :: Rule -> Int -> A.PrimArray Word8 -> A.PrimArray Word8
runIterations r n v = go n v where
  go 0 v = v
  go n v = go (n - 1) (step r v)

main :: IO ()
main =  do
  let v = fromList (1:replicate 1000 0)
  print $ runIterations (mkRule 110) 1000000 v