foundation-0.0.14: Foundation/Array/Bitmap.hs
-- |
-- Module : Foundation.Array.Bitmap
-- License : BSD-style
-- Maintainer : Vincent Hanquez <vincent@snarc.org>
-- Stability : experimental
-- Portability : portable
--
-- A simple abstraction to a set of Bits (Bitmap)
--
-- Largely a placeholder for a more performant implementation,
-- most operation goes through the List representation (e.g. [Bool])
-- to conduct even the most trivial operation, leading to a lots of
-- unnecessary churn.
--
{-# LANGUAGE BangPatterns #-}
module Foundation.Array.Bitmap
( Bitmap
, MutableBitmap
, empty
, append
, concat
, unsafeIndex
, index
, read
, unsafeRead
, write
, unsafeWrite
, snoc
, cons
) where
import Basement.UArray (UArray)
import qualified Basement.UArray as A
import Basement.UArray.Mutable (MUArray)
import Basement.Compat.Bifunctor (first, second, bimap)
import Basement.Exception
import Basement.Compat.Base
import Basement.Types.OffsetSize
import Basement.Monad
import qualified Foundation.Collection as C
import Foundation.Numerical
import Data.Bits
import Foundation.Bits
import GHC.ST
import qualified Data.List
data Bitmap = Bitmap (CountOf Bool) (UArray Word32)
data MutableBitmap st = MutableBitmap (CountOf Bool) (MUArray Word32 st)
bitsPerTy :: Int
bitsPerTy = 32
shiftPerTy :: Int
shiftPerTy = 5
maskPerTy :: Int
maskPerTy = 0x1f
instance Show Bitmap where
show v = show (toList v)
instance Eq Bitmap where
(==) = equal
instance Ord Bitmap where
compare = vCompare
instance Monoid Bitmap where
mempty = empty
mappend = append
mconcat = concat
type instance C.Element Bitmap = Bool
instance IsList Bitmap where
type Item Bitmap = Bool
fromList = vFromList
toList = vToList
instance C.InnerFunctor Bitmap where
imap = map
instance C.Foldable Bitmap where
foldr = foldr
foldl' = foldl'
foldr' = foldr'
instance C.Collection Bitmap where
null = null
length = length
elem e = Data.List.elem e . toList
maximum = any id . C.getNonEmpty
minimum = all id . C.getNonEmpty
all = all
any = any
instance C.Sequential Bitmap where
take = take
drop = drop
splitAt = splitAt
revTake n = unoptimised (C.revTake n)
revDrop n = unoptimised (C.revDrop n)
splitOn = splitOn
break = break
breakEnd = breakEnd
span = span
filter = filter
reverse = reverse
snoc = snoc
cons = cons
unsnoc = unsnoc
uncons = uncons
intersperse = intersperse
find = find
sortBy = sortBy
singleton = fromList . (:[])
replicate n = fromList . C.replicate n
instance C.IndexedCollection Bitmap where
(!) l n
| isOutOfBound n (length l) = Nothing
| otherwise = Just $ index l n
findIndex predicate c = loop 0
where
!len = length c
loop i
| i .==# len = Nothing
| predicate (unsafeIndex c i) = Just i
| otherwise = Nothing
instance C.MutableCollection MutableBitmap where
type MutableFreezed MutableBitmap = Bitmap
type MutableKey MutableBitmap = Offset Bool
type MutableValue MutableBitmap = Bool
thaw = thaw
freeze = freeze
unsafeThaw = unsafeThaw
unsafeFreeze = unsafeFreeze
mutNew = new
mutUnsafeWrite = unsafeWrite
mutUnsafeRead = unsafeRead
mutWrite = write
mutRead = read
bitmapIndex :: Offset Bool -> (Offset Word32, Int)
bitmapIndex (Offset !i) = (Offset (i .>>. shiftPerTy), i .&. maskPerTy)
{-# INLINE bitmapIndex #-}
-- return the index in word32 quantity and mask to a bit in a bitmap
{-
bitmapAddr :: Int -> (# Int , Word #)
bitmapAddr !i = (# idx, mask #)
where (!idx, !bitIdx) = bitmapIndex i
!mask = case bitIdx of
0 -> 0x1
1 -> 0x2
2 -> 0x4
3 -> 0x8
4 -> 0x10
5 -> 0x20
6 -> 0x40
7 -> 0x80
8 -> 0x100
9 -> 0x200
10 -> 0x400
11 -> 0x800
12 -> 0x1000
13 -> 0x2000
14 -> 0x4000
15 -> 0x8000
16 -> 0x10000
17 -> 0x20000
18 -> 0x40000
19 -> 0x80000
20 -> 0x100000
21 -> 0x200000
22 -> 0x400000
23 -> 0x800000
24 -> 0x1000000
25 -> 0x2000000
26 -> 0x4000000
27 -> 0x8000000
28 -> 0x10000000
29 -> 0x20000000
30 -> 0x40000000
_ -> 0x80000000
-}
thaw :: PrimMonad prim => Bitmap -> prim (MutableBitmap (PrimState prim))
thaw (Bitmap len ba) = MutableBitmap len `fmap` C.thaw ba
freeze :: PrimMonad prim => MutableBitmap (PrimState prim) -> prim Bitmap
freeze (MutableBitmap len mba) = Bitmap len `fmap` C.freeze mba
unsafeThaw :: PrimMonad prim => Bitmap -> prim (MutableBitmap (PrimState prim))
unsafeThaw (Bitmap len ba) = MutableBitmap len `fmap` C.unsafeThaw ba
unsafeFreeze :: PrimMonad prim => MutableBitmap (PrimState prim) -> prim Bitmap
unsafeFreeze (MutableBitmap len mba) = Bitmap len `fmap` C.unsafeFreeze mba
unsafeWrite :: PrimMonad prim => MutableBitmap (PrimState prim) -> Offset Bool -> Bool -> prim ()
unsafeWrite (MutableBitmap _ ma) i v = do
let (idx, bitIdx) = bitmapIndex i
w <- A.unsafeRead ma idx
let w' = if v then setBit w bitIdx else clearBit w bitIdx
A.unsafeWrite ma idx w'
{-# INLINE unsafeWrite #-}
unsafeRead :: PrimMonad prim => MutableBitmap (PrimState prim) -> Offset Bool -> prim Bool
unsafeRead (MutableBitmap _ ma) i = do
let (idx, bitIdx) = bitmapIndex i
flip testBit bitIdx `fmap` A.unsafeRead ma idx
{-# INLINE unsafeRead #-}
write :: PrimMonad prim => MutableBitmap (PrimState prim) -> Offset Bool -> Bool -> prim ()
write mb n val
| isOutOfBound n len = primOutOfBound OOB_Write n len
| otherwise = unsafeWrite mb n val
where
len = mutableLength mb
{-# INLINE write #-}
read :: PrimMonad prim => MutableBitmap (PrimState prim) -> Offset Bool -> prim Bool
read mb n
| isOutOfBound n len = primOutOfBound OOB_Read n len
| otherwise = unsafeRead mb n
where len = mutableLength mb
{-# INLINE read #-}
-- | Return the element at a specific index from a Bitmap.
--
-- If the index @n is out of bounds, an error is raised.
index :: Bitmap -> Offset Bool -> Bool
index bits n
| isOutOfBound n len = outOfBound OOB_Index n len
| otherwise = unsafeIndex bits n
where len = length bits
{-# INLINE index #-}
-- | Return the element at a specific index from an array without bounds checking.
--
-- Reading from invalid memory can return unpredictable and invalid values.
-- use 'index' if unsure.
unsafeIndex :: Bitmap -> Offset Bool -> Bool
unsafeIndex (Bitmap _ ba) n =
let (idx, bitIdx) = bitmapIndex n
in testBit (A.unsafeIndex ba idx) bitIdx
{-# INLINE unsafeIndex #-}
-----------------------------------------------------------------------
-- higher level collection implementation
-----------------------------------------------------------------------
length :: Bitmap -> CountOf Bool
length (Bitmap sz _) = sz
mutableLength :: MutableBitmap st -> CountOf Bool
mutableLength (MutableBitmap sz _) = sz
empty :: Bitmap
empty = Bitmap 0 mempty
new :: PrimMonad prim => CountOf Bool -> prim (MutableBitmap (PrimState prim))
new sz@(CountOf len) =
MutableBitmap sz <$> A.new nbElements
where
nbElements :: CountOf Word32
nbElements = CountOf ((len `alignRoundUp` bitsPerTy) .>>. shiftPerTy)
-- | make an array from a list of elements.
vFromList :: [Bool] -> Bitmap
vFromList allBools = runST $ do
mbitmap <- new len
loop mbitmap 0 allBools
where
loop mb _ [] = unsafeFreeze mb
loop mb i (x:xs) = unsafeWrite mb i x >> loop mb (i+1) xs
{-
runST $ do
mba <- A.new nbElements
ba <- loop mba (0 :: Int) allBools
pure (Bitmap len ba)
where
loop mba _ [] = A.unsafeFreeze mba
loop mba i l = do
let (l1, l2) = C.splitAt bitsPerTy l
w = toPacked l1
A.unsafeWrite mba i w
loop mba (i+1) l2
toPacked :: [Bool] -> Word32
toPacked l =
C.foldl' (.|.) 0 $ Prelude.zipWith (\b w -> if b then (1 `shiftL` w) else 0) l (C.reverse [0..31])
-}
len = C.length allBools
-- | transform an array to a list.
vToList :: Bitmap -> [Bool]
vToList a = loop 0
where len = length a
loop i | i .==# len = []
| otherwise = unsafeIndex a i : loop (i+1)
-- | Check if two vectors are identical
equal :: Bitmap -> Bitmap -> Bool
equal a b
| la /= lb = False
| otherwise = loop 0
where
!la = length a
!lb = length b
loop n | n .==# la = True
| otherwise = (unsafeIndex a n == unsafeIndex b n) && loop (n+1)
-- | Compare 2 vectors
vCompare :: Bitmap -> Bitmap -> Ordering
vCompare a b = loop 0
where
!la = length a
!lb = length b
loop n
| n .==# la = if la == lb then EQ else LT
| n .==# lb = GT
| otherwise =
case unsafeIndex a n `compare` unsafeIndex b n of
EQ -> loop (n+1)
r -> r
-- | Append 2 arrays together by creating a new bigger array
--
-- TODO completely non optimized
append :: Bitmap -> Bitmap -> Bitmap
append a b = fromList $ toList a `mappend` toList b
-- TODO completely non optimized
concat :: [Bitmap] -> Bitmap
concat l = fromList $ mconcat $ fmap toList l
null :: Bitmap -> Bool
null (Bitmap nbBits _) = nbBits == 0
take :: CountOf Bool -> Bitmap -> Bitmap
take nbElems bits@(Bitmap nbBits ba)
| nbElems <= 0 = empty
| nbElems >= nbBits = bits
| otherwise = Bitmap nbElems ba -- TODO : although it work right now, take on the underlaying ba too
drop :: CountOf Bool -> Bitmap -> Bitmap
drop nbElems bits@(Bitmap nbBits _)
| nbElems <= 0 = bits
| nbElems >= nbBits = empty
| otherwise = unoptimised (C.drop nbElems) bits
-- TODO: decide if we have drop easy by having a bit offset in the data structure
-- or if we need to shift stuff around making all the indexing slighlty more complicated
splitAt :: CountOf Bool -> Bitmap -> (Bitmap, Bitmap)
splitAt n v = (take n v, drop n v)
-- unoptimised
splitOn :: (Bool -> Bool) -> Bitmap -> [Bitmap]
splitOn f bits = fmap fromList $ C.splitOn f $ toList bits
-- unoptimised
break :: (Bool -> Bool) -> Bitmap -> (Bitmap, Bitmap)
break predicate v = findBreak 0
where
len = length v
findBreak i
| i .==# len = (v, empty)
| otherwise =
if predicate (unsafeIndex v i)
then splitAt (offsetAsSize i) v
else findBreak (i+1)
breakEnd :: (Bool -> Bool) -> Bitmap -> (Bitmap, Bitmap)
breakEnd predicate = bimap fromList fromList . C.breakEnd predicate . toList
span :: (Bool -> Bool) -> Bitmap -> (Bitmap, Bitmap)
span p = break (not . p)
map :: (Bool -> Bool) -> Bitmap -> Bitmap
map f bits = unoptimised (fmap f) bits
--mapIndex :: (Int -> Bool -> Bool) -> Bitmap -> Bitmap
--mapIndex f Bitmap =
cons :: Bool -> Bitmap -> Bitmap
cons v l = unoptimised (C.cons v) l
snoc :: Bitmap -> Bool -> Bitmap
snoc l v = unoptimised (flip C.snoc v) l
-- unoptimised
uncons :: Bitmap -> Maybe (Bool, Bitmap)
uncons b = fmap (second fromList) $ C.uncons $ toList b
-- unoptimised
unsnoc :: Bitmap -> Maybe (Bitmap, Bool)
unsnoc b = fmap (first fromList) $ C.unsnoc $ toList b
intersperse :: Bool -> Bitmap -> Bitmap
intersperse b = unoptimised (C.intersperse b)
find :: (Bool -> Bool) -> Bitmap -> Maybe Bool
find predicate vec = loop 0
where
!len = length vec
loop i
| i .==# len = Nothing
| otherwise =
let e = unsafeIndex vec i
in if predicate e then Just e else loop (i+1)
sortBy :: (Bool -> Bool -> Ordering) -> Bitmap -> Bitmap
sortBy by bits = unoptimised (C.sortBy by) bits
filter :: (Bool -> Bool) -> Bitmap -> Bitmap
filter predicate vec = unoptimised (Data.List.filter predicate) vec
reverse :: Bitmap -> Bitmap
reverse bits = unoptimised C.reverse bits
foldr :: (Bool -> a -> a) -> a -> Bitmap -> a
foldr f initialAcc vec = loop 0
where
len = length vec
loop i
| i .==# len = initialAcc
| otherwise = unsafeIndex vec i `f` loop (i+1)
foldr' :: (Bool -> a -> a) -> a -> Bitmap -> a
foldr' = foldr
foldl' :: (a -> Bool -> a) -> a -> Bitmap -> a
foldl' f initialAcc vec = loop 0 initialAcc
where
len = length vec
loop i !acc
| i .==# len = acc
| otherwise = loop (i+1) (f acc (unsafeIndex vec i))
all :: (Bool -> Bool) -> Bitmap -> Bool
all p bm = loop 0
where
len = length bm
loop !i
| i .==# len = True
| not $ p (unsafeIndex bm i) = False
| otherwise = loop (i + 1)
any :: (Bool -> Bool) -> Bitmap -> Bool
any p bm = loop 0
where
len = length bm
loop !i
| i .==# len = False
| p (unsafeIndex bm i) = True
| otherwise = loop (i + 1)
unoptimised :: ([Bool] -> [Bool]) -> Bitmap -> Bitmap
unoptimised f = vFromList . f . vToList