{-# LANGUAGE FlexibleContexts, StandaloneDeriving, TypeFamilies #-}
module Game.LambdaHack.Common.PointArray
( UnboxRepClass(..), Array(..), pindex, punindex
, (!), accessI, (//), unsafeUpdateA, unsafeWriteA, unsafeWriteManyA
, replicateA, replicateMA, generateA, generateMA, unfoldrNA, sizeA
, foldrA, foldrA', foldlA', ifoldrA, ifoldrA', ifoldlA', foldMA', ifoldMA'
, mapA, imapA, imapMA_, safeSetA, unsafeSetA
, minIndexA, minLastIndexA, minIndexesA, maxIndexA, maxLastIndexA, forceA
, fromListA, toListA
#ifdef EXPOSE_INTERNAL
, toUnboxRep
#endif
) where
import Prelude ()
import Game.LambdaHack.Common.Prelude
import Control.Monad.ST.Strict
import Data.Binary
import Data.Vector.Binary ()
import qualified Data.Vector.Fusion.Bundle as Bundle
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
class ( Ord c, Eq (UnboxRep c), Ord (UnboxRep c), Bounded (UnboxRep c)
, Binary (UnboxRep c), U.Unbox (UnboxRep c) )
=> UnboxRepClass c where
type UnboxRep c
type instance UnboxRep c = c
toUnboxRepUnsafe :: c -> UnboxRep c
fromUnboxRep :: UnboxRep c -> c
instance UnboxRepClass Bool where
toUnboxRepUnsafe c = c
fromUnboxRep c = c
instance UnboxRepClass Word8 where
toUnboxRepUnsafe c = c
fromUnboxRep c = c
data Array c = Array
{ axsize :: X
, aysize :: Y
, avector :: U.Vector (UnboxRep c)
}
deriving instance UnboxRepClass c => Eq (Array c)
instance Show (Array c) where
show a = "PointArray.Array with size " ++ show (axsize a, aysize a)
instance UnboxRepClass c => Binary (Array c) where
put Array{..} = do
put axsize
put aysize
put avector
get = do
axsize <- get
aysize <- get
avector <- get
return $! Array{..}
toUnboxRep :: UnboxRepClass c => c -> UnboxRep c
{-# INLINE toUnboxRep #-}
toUnboxRep c =
#ifdef WITH_EXPENSIVE_ASSERTIONS
assert (c <= fromUnboxRep maxBound) $
#endif
toUnboxRepUnsafe c
pindex :: X -> Point -> Int
{-# INLINE pindex #-}
pindex xsize (Point x y) = x + y * xsize
punindex :: X -> Int -> Point
{-# INLINE punindex #-}
punindex xsize n = let (y, x) = n `quotRem` xsize
in Point x y
(!) :: UnboxRepClass c => Array c -> Point -> c
{-# INLINE (!) #-}
(!) Array{..} p = fromUnboxRep $ avector U.! pindex axsize p
accessI :: UnboxRepClass c => Array c -> Int -> UnboxRep c
{-# INLINE accessI #-}
accessI Array{..} p = avector `U.unsafeIndex` p
(//) :: UnboxRepClass c => Array c -> [(Point, c)] -> Array c
{-# INLINE (//) #-}
(//) Array{..} l = let v = avector U.// map (pindex axsize *** toUnboxRep) l
in Array{avector = v, ..}
unsafeUpdateA :: UnboxRepClass c => Array c -> [(Point, c)] -> ()
{-# INLINE unsafeUpdateA #-}
unsafeUpdateA Array{..} l = runST $ do
vThawed <- U.unsafeThaw avector
mapM_ (\(p, c) -> VM.write vThawed (pindex axsize p) (toUnboxRep c)) l
void $ U.unsafeFreeze vThawed
unsafeWriteA :: UnboxRepClass c => Array c -> Point -> c -> ()
{-# INLINE unsafeWriteA #-}
unsafeWriteA Array{..} p c = runST $ do
vThawed <- U.unsafeThaw avector
VM.write vThawed (pindex axsize p) (toUnboxRep c)
void $ U.unsafeFreeze vThawed
unsafeWriteManyA :: UnboxRepClass c => Array c -> [Point] -> c -> ()
{-# INLINE unsafeWriteManyA #-}
unsafeWriteManyA Array{..} l c = runST $ do
vThawed <- U.unsafeThaw avector
let d = toUnboxRep c
mapM_ (\p -> VM.write vThawed (pindex axsize p) d) l
void $ U.unsafeFreeze vThawed
replicateA :: UnboxRepClass c => X -> Y -> c -> Array c
{-# INLINE replicateA #-}
replicateA axsize aysize c =
Array{avector = U.replicate (axsize * aysize) $ toUnboxRep c, ..}
replicateMA :: (Monad m, UnboxRepClass c) => X -> Y -> m c -> m (Array c)
{-# INLINE replicateMA #-}
replicateMA axsize aysize m = do
v <- U.replicateM (axsize * aysize) $ liftM toUnboxRep m
return $! Array{avector = v, ..}
generateA :: UnboxRepClass c => X -> Y -> (Point -> c) -> Array c
{-# INLINE generateA #-}
generateA axsize aysize f =
let g n = toUnboxRep $ f $ punindex axsize n
in Array{avector = U.generate (axsize * aysize) g, ..}
generateMA :: (Monad m, UnboxRepClass c)
=> X -> Y -> (Point -> m c) -> m (Array c)
{-# INLINE generateMA #-}
generateMA axsize aysize fm = do
let gm n = liftM toUnboxRep $ fm $ punindex axsize n
v <- U.generateM (axsize * aysize) gm
return $! Array{avector = v, ..}
unfoldrNA :: UnboxRepClass c => X -> Y -> (b -> (c, b)) -> b -> Array c
{-# INLINE unfoldrNA #-}
unfoldrNA axsize aysize fm b =
let gm = Just . first toUnboxRep . fm
v = U.unfoldrN (axsize * aysize) gm b
in Array {avector = v, ..}
sizeA :: Array c -> (X, Y)
{-# INLINE sizeA #-}
sizeA Array{..} = (axsize, aysize)
foldrA :: UnboxRepClass c => (c -> a -> a) -> a -> Array c -> a
{-# INLINE foldrA #-}
foldrA f z0 Array{..} =
U.foldr (\c a-> f (fromUnboxRep c) a) z0 avector
foldrA' :: UnboxRepClass c => (c -> a -> a) -> a -> Array c -> a
{-# INLINE foldrA' #-}
foldrA' f z0 Array{..} =
U.foldr' (\c a-> f (fromUnboxRep c) a) z0 avector
foldlA' :: UnboxRepClass c => (a -> c -> a) -> a -> Array c -> a
{-# INLINE foldlA' #-}
foldlA' f z0 Array{..} =
U.foldl' (\a c -> f a (fromUnboxRep c)) z0 avector
ifoldlA' :: UnboxRepClass c => (a -> Point -> c -> a) -> a -> Array c -> a
{-# INLINE ifoldlA' #-}
ifoldlA' f z0 Array{..} =
U.ifoldl' (\a n c -> f a (punindex axsize n) (fromUnboxRep c)) z0 avector
ifoldrA :: UnboxRepClass c => (Point -> c -> a -> a) -> a -> Array c -> a
{-# INLINE ifoldrA #-}
ifoldrA f z0 Array{..} =
U.ifoldr (\n c a -> f (punindex axsize n) (fromUnboxRep c) a) z0 avector
ifoldrA' :: UnboxRepClass c => (Point -> c -> a -> a) -> a -> Array c -> a
{-# INLINE ifoldrA' #-}
ifoldrA' f z0 Array{..} =
U.ifoldr' (\n c a -> f (punindex axsize n) (fromUnboxRep c) a) z0 avector
foldMA' :: (Monad m, UnboxRepClass c) => (a -> c -> m a) -> a -> Array c -> m a
{-# INLINE foldMA' #-}
foldMA' f z0 Array{..} =
U.foldM' (\a c -> f a (fromUnboxRep c)) z0 avector
ifoldMA' :: (Monad m, UnboxRepClass c)
=> (a -> Point -> c -> m a) -> a -> Array c -> m a
{-# INLINE ifoldMA' #-}
ifoldMA' f z0 Array{..} =
U.ifoldM' (\a n c -> f a (punindex axsize n) (fromUnboxRep c)) z0 avector
mapA :: (UnboxRepClass c, UnboxRepClass d) => (c -> d) -> Array c -> Array d
{-# INLINE mapA #-}
mapA f Array{..} =
Array{avector = U.map (toUnboxRep . f . fromUnboxRep) avector, ..}
imapA :: (UnboxRepClass c, UnboxRepClass d)
=> (Point -> c -> d) -> Array c -> Array d
{-# INLINE imapA #-}
imapA f Array{..} =
let v = U.imap (\n c ->
toUnboxRep $ f (punindex axsize n) (fromUnboxRep c)) avector
in Array{avector = v, ..}
imapMA_ :: (Monad m, UnboxRepClass c) => (Point -> c -> m ()) -> Array c -> m ()
{-# INLINE imapMA_ #-}
imapMA_ f Array{..} =
U.imapM_ (\n c -> f (punindex axsize n) (fromUnboxRep c)) avector
unsafeSetA :: UnboxRepClass c => c -> Array c -> Array c
{-# INLINE unsafeSetA #-}
unsafeSetA c Array{..} = runST $ do
vThawed <- U.unsafeThaw avector
VM.set vThawed (toUnboxRep c)
vFrozen <- U.unsafeFreeze vThawed
return $! Array{avector = vFrozen, ..}
safeSetA :: UnboxRepClass c => c -> Array c -> Array c
{-# INLINE safeSetA #-}
safeSetA c Array{..} =
Array{avector = U.modify (\v -> VM.set v (toUnboxRep c)) avector, ..}
minIndexA :: UnboxRepClass c => Array c -> Point
{-# INLINE minIndexA #-}
minIndexA Array{..} = punindex axsize $ U.minIndex avector
minLastIndexA :: UnboxRepClass c => Array c -> Point
{-# INLINE minLastIndexA #-}
minLastIndexA Array{..} =
punindex axsize
$ fst . Bundle.foldl1' imin . Bundle.indexed . G.stream
$ avector
where
imin (i, x) (j, y) = i `seq` j `seq` if x >= y then (j, y) else (i, x)
minIndexesA :: UnboxRepClass c => Array c -> [Point]
{-# INLINE minIndexesA #-}
minIndexesA Array{..} =
Bundle.foldr imin [] . Bundle.indexed . G.stream $ avector
where
imin (i, x) acc = if x == minE
then let !j = punindex axsize i
in j : acc
else acc
!minE = U.minimum avector
maxIndexA :: UnboxRepClass c => Array c -> Point
{-# INLINE maxIndexA #-}
maxIndexA Array{..} = punindex axsize $ U.maxIndex avector
maxLastIndexA :: UnboxRepClass c => Array c -> Point
{-# INLINE maxLastIndexA #-}
maxLastIndexA Array{..} =
punindex axsize
$ fst . Bundle.foldl1' imax . Bundle.indexed . G.stream
$ avector
where
imax (i, x) (j, y) = i `seq` j `seq` if x <= y then (j, y) else (i, x)
forceA :: UnboxRepClass c => Array c -> Array c
{-# INLINE forceA #-}
forceA Array{..} = Array{avector = U.force avector, ..}
fromListA :: UnboxRepClass c => X -> Y -> [c] -> Array c
{-# INLINE fromListA #-}
fromListA axsize aysize l =
Array{avector = U.fromListN (axsize * aysize) $ map toUnboxRep l, ..}
toListA :: UnboxRepClass c => Array c -> [c]
{-# INLINE toListA #-}
toListA Array{..} = map fromUnboxRep $ U.toList avector