gistfile1.hs

{-# LANGUAGE BangPatterns #-}

{-
Row/Column storage grid where each node is an 8x8 chunk of bytes.
-}

module ByteGrid 
( ByteGrid
, make
, readUnsafe
, writeUnsafe
, readOrDefault
, readMaybe
, write
, modify
, fromFunction
) where

import Data.Word (Word8)
import qualified Data.Vector as V
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Primitive.ByteArray as B
import Control.Monad.ST.Safe (runST)

data ByteGrid = ByteGrid 
    !Int -- Width
    !Int -- Height
    !(V.Vector (V.Vector B.ByteArray))

-- O(w*h)
make :: (Int,Int) -> Word8 -> ByteGrid
make (!w,!h) !def = ByteGrid w h vec
    where
    vec = runST $ do
        mg <- B.newByteArray 64
        B.fillByteArray mg 0 64 def
        ig <- B.unsafeFreezeByteArray mg
        return $ V.replicate (wd8 + 1) (V.replicate (hd8 + 1) ig)
    wd8 = w `div` 8
    hd8 = h `div` 8

-- O(1)
readUnsafe :: ByteGrid -> (Int,Int) -> Word8
readUnsafe (ByteGrid _ _ vec) (!x,!y) = 
    vec `V.unsafeIndex` 
    xd8 `V.unsafeIndex` 
    yd8 `B.indexByteArray` (ym8 * 8 + xm8)
    where
    xd8 = x `div` 8
    yd8 = y `div` 8
    xm8 = x `mod` 8
    ym8 = y `mod` 8

-- O(w+h)
writeUnsafe :: ByteGrid -> (Int,Int) -> Word8 -> ByteGrid
writeUnsafe (ByteGrid w h vecX) (!x,!y) !val = ByteGrid w h vec
    where
    vec = V.modify 
            (\vx -> M.unsafeWrite vx xd8 $! V.modify 
                (\vy -> M.unsafeWrite vy yd8 $! runST $ do
                    mg <- B.newByteArray 64
                    B.copyByteArray mg 0 ba 0 64
                    B.writeByteArray mg (ym8 * 8 + xm8) val
                    B.unsafeFreezeByteArray mg)
                vecY) 
            vecX 
    vecY = vecX V.! xd8
    ba = vecY V.! yd8
    xd8 = x `div` 8
    yd8 = y `div` 8
    xm8 = x `mod` 8
    ym8 = y `mod` 8

-- O(1)
readOrDefault :: ByteGrid -> Word8 -> (Int,Int) -> Word8
readOrDefault bg !def xy = 
    if isDefinedOn bg xy
    then readUnsafe bg xy
    else def

-- O(1)
readMaybe :: ByteGrid -> (Int,Int) -> Maybe Word8
readMaybe bg xy = 
    if isDefinedOn bg xy
    then Just $! readUnsafe bg xy
    else Nothing

-- O(w+h)
write :: ByteGrid -> (Int,Int) -> Word8 -> ByteGrid
write bg xy !val = 
    if isDefinedOn bg xy
    then writeUnsafe bg xy val
    else error $ "Bad write to grid at " ++ show xy

-- O(w+h)
modify :: ByteGrid -> (Int,Int) -> (Word8 -> Word8) -> ByteGrid
modify bg xy f = 
    if isDefinedOn bg xy
    then writeUnsafe bg xy (f (readUnsafe bg xy))
    else error $ "Bad modify to grid at " ++ show xy

isDefinedOn :: ByteGrid -> (Int,Int) -> Bool
isDefinedOn (ByteGrid w h vec) (!x,!y) = 
    x < w && y < h && x >= 0 && y >= 0

-- O(w*h)
fromFunction :: (Int,Int) -> ((Int,Int) -> Word8) -> ByteGrid
fromFunction wh f = foldl (\g xy -> write g xy $ f xy) 
                          (make wh 0) 
                          [(x,y) | x <- [0..fst wh - 1], y <- [0..snd wh - 1]]