module Game.LambdaHack.Common.PointArray
( Array
, (!), (//), replicateA, replicateMA, generateA, generateMA, sizeA
, foldlA, ifoldlA, mapA, imapA, mapWithKeyMA
, safeSetA, unsafeSetA, unsafeUpdateA
, minIndexA, minLastIndexA, minIndexesA, maxIndexA, maxLastIndexA, forceA
) where
import Control.Arrow ((***))
import Control.Monad
import Control.Monad.ST.Strict
import Data.Binary
import Data.Vector.Binary ()
import qualified Data.Vector.Fusion.Stream as Stream
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as VM
import Game.LambdaHack.Common.Point
data Array c = Array
{ axsize :: !X
, aysize :: !Y
, avector :: !(U.Vector Word8)
}
deriving Eq
instance Show (Array c) where
show a = "PointArray.Array with size " ++ show (sizeA a)
cnv :: (Enum a, Enum b) => a -> b
cnv = toEnum . fromEnum
pindex :: X -> Point -> Int
pindex xsize (Point x y) = x + y * xsize
punindex :: X -> Int -> Point
punindex xsize n = let (y, x) = n `quotRem` xsize
in Point x y
(!) :: Enum c => Array c -> Point -> c
(!) Array{..} p = cnv $ avector U.! pindex axsize p
(//) :: Enum c => Array c -> [(Point, c)] -> Array c
(//) Array{..} l = let v = avector U.// map (pindex axsize *** cnv) l
in Array{avector = v, ..}
unsafeUpdateA :: Enum c => Array c -> [(Point, c)] -> Array c
unsafeUpdateA Array{..} l = runST $ do
vThawed <- U.unsafeThaw avector
mapM_ (\(p, c) -> VM.write vThawed (pindex axsize p) (cnv c)) l
vFrozen <- U.unsafeFreeze vThawed
return $! Array{avector = vFrozen, ..}
replicateA :: Enum c => X -> Y -> c -> Array c
replicateA axsize aysize c =
Array{avector = U.replicate (axsize * aysize) $ cnv c, ..}
replicateMA :: Enum c => Monad m => X -> Y -> m c -> m (Array c)
replicateMA axsize aysize m = do
v <- U.replicateM (axsize * aysize) $ liftM cnv m
return $! Array{avector = v, ..}
generateA :: Enum c => X -> Y -> (Point -> c) -> Array c
generateA axsize aysize f =
let g n = cnv $ f $ punindex axsize n
in Array{avector = U.generate (axsize * aysize) g, ..}
generateMA :: Enum c => Monad m => X -> Y -> (Point -> m c) -> m (Array c)
generateMA axsize aysize fm = do
let gm n = liftM cnv $ fm $ punindex axsize n
v <- U.generateM (axsize * aysize) gm
return $! Array{avector = v, ..}
sizeA :: Array c -> (X, Y)
sizeA Array{..} = (axsize, aysize)
foldlA :: Enum c => (a -> c -> a) -> a -> Array c -> a
foldlA f z0 Array{..} =
U.foldl' (\a c -> f a (cnv c)) z0 avector
ifoldlA :: Enum c => (a -> Point -> c -> a) -> a -> Array c -> a
ifoldlA f z0 Array{..} =
U.ifoldl' (\a n c -> f a (punindex axsize n) (cnv c)) z0 avector
mapA :: (Enum c, Enum d) => (c -> d) -> Array c -> Array d
mapA f Array{..} = Array{avector = U.map (cnv . f . cnv) avector, ..}
imapA :: (Enum c, Enum d) => (Point -> c -> d) -> Array c -> Array d
imapA f Array{..} =
let v = U.imap (\n c -> cnv $ f (punindex axsize n) (cnv c)) avector
in Array{avector = v, ..}
unsafeSetA :: Enum c => c -> Array c -> Array c
unsafeSetA c Array{..} = runST $ do
vThawed <- U.unsafeThaw avector
VM.set vThawed (cnv c)
vFrozen <- U.unsafeFreeze vThawed
return $! Array{avector = vFrozen, ..}
safeSetA :: Enum c => c -> Array c -> Array c
safeSetA c Array{..} =
Array{avector = U.modify (\v -> VM.set v (cnv c)) avector, ..}
mapWithKeyMA :: Enum c => Monad m
=> (Point -> c -> m ()) -> Array c -> m ()
mapWithKeyMA f Array{..} =
U.ifoldl' (\a n c -> a >> f (punindex axsize n) (cnv c))
(return ())
avector
minIndexA :: Enum c => Array c -> Point
minIndexA Array{..} = punindex axsize $ U.minIndex avector
minLastIndexA :: Enum c => Array c -> Point
minLastIndexA Array{..} =
punindex axsize
$ fst . Stream.foldl1' imin . Stream.indexed . G.stream
$ avector
where
imin (i, x) (j, y) = i `seq` j `seq` if x >= y then (j, y) else (i, x)
minIndexesA :: Enum c => Array c -> [Point]
minIndexesA Array{..} =
map (punindex axsize)
$ Stream.foldl' imin [] . Stream.indexed . G.stream
$ avector
where
imin acc (i, x) = i `seq` if x == minE then i : acc else acc
minE = cnv $ U.minimum avector
maxIndexA :: Enum c => Array c -> Point
maxIndexA Array{..} = punindex axsize $ U.maxIndex avector
maxLastIndexA :: Enum c => Array c -> Point
maxLastIndexA Array{..} =
punindex axsize
$ fst . Stream.foldl1' imax . Stream.indexed . G.stream
$ avector
where
imax (i, x) (j, y) = i `seq` j `seq` if x <= y then (j, y) else (i, x)
forceA :: Enum c => Array c -> Array c
forceA Array{..} = Array{avector = U.force avector, ..}
instance Binary (Array c) where
put Array{..} = do
put axsize
put aysize
put avector
get = do
axsize <- get
aysize <- get
avector <- get
return $! Array{..}