gistfile1.hs

{-# LANGUAGE BangPatterns #-}

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

module ByteGrid 
( ByteGrid
, make
, readUnsafe
, 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)
import Control.DeepSeq (deepseq,NFData(..))

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

-- O(w*h)
-- {-# INLINE make #-}
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)
-- {-# INLINE readUnsafe #-}
readUnsafe :: ByteGrid -> (Int,Int) -> Word8
readUnsafe (ByteGrid !w !h !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(1)
-- {-# INLINE readOrDefault #-}
readOrDefault :: ByteGrid -> Word8 -> (Int,Int) -> Word8
readOrDefault (ByteGrid !w !h !vec) !def (!x,!y) = 
    if x < w && y < h && x >= 0 && y >= 0
    then vec `V.unsafeIndex` 
         xd8 `V.unsafeIndex` 
         yd8 `B.indexByteArray` (ym8 * 8 + xm8)
    else def
    where
    xd8 = x `div` 8
    yd8 = y `div` 8
    xm8 = x `mod` 8
    ym8 = y `mod` 8

-- O(1)
-- {-# INLINE readMaybe #-}
readMaybe :: ByteGrid -> (Int,Int) -> Maybe Word8
readMaybe (ByteGrid !w !h !vec) (!x,!y) = 
    if x < w && y < h && x >= 0 && y >= 0
    then Just $ vec `V.unsafeIndex` 
                xd8 `V.unsafeIndex` 
                yd8 `B.indexByteArray` (ym8 * 8 + xm8)
    else Nothing
    where
    xd8 = x `div` 8
    yd8 = y `div` 8
    xm8 = x `mod` 8
    ym8 = y `mod` 8

-- O(w+h)
-- {-# INLINE write #-}
write :: ByteGrid -> (Int,Int) -> Word8 -> ByteGrid
write bg@(ByteGrid !w !h !vecX) (!x,!y) !val = 
    if x < w && y < h && x >= 0 && y >= 0
    then ByteGrid w h vec
    else bg
    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(w+h)
-- {-# INLINE modify #-}
modify :: ByteGrid -> (Int,Int) -> (Word8 -> Word8) -> ByteGrid
modify !bg@(ByteGrid !w !h !vecX) (!x,!y) !f = 
    if x < w && y < h && x >= 0 && y >= 0
    then ByteGrid w h vec
    else bg
    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) (f $ readUnsafe bg (x,y))
                    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

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]]