1-- |
2-- Module      : Foundation.Array.Bitmap
3-- License     : BSD-style
4-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
5-- Stability   : experimental
6-- Portability : portable
7--
8-- A simple abstraction to a set of Bits (Bitmap)
9--
10-- Largely a placeholder for a more performant implementation,
11-- most operation goes through the List representation (e.g. [Bool])
12-- to conduct even the most trivial operation, leading to a lots of
13-- unnecessary churn.
14--
15{-# LANGUAGE BangPatterns #-}
16{-# LANGUAGE DeriveDataTypeable #-}
17module Foundation.Array.Bitmap
18    ( Bitmap
19    , MutableBitmap
20    , empty
21    , append
22    , concat
23    , unsafeIndex
24    , index
25    , read
26    , unsafeRead
27    , write
28    , unsafeWrite
29    , snoc
30    , cons
31    ) where
32
33import           Basement.UArray (UArray)
34import qualified Basement.UArray as A
35import           Basement.UArray.Mutable (MUArray)
36import           Basement.Compat.Bifunctor (first, second, bimap)
37import           Basement.Compat.Semigroup
38import           Basement.Exception
39import           Basement.Compat.Base
40import           Basement.Types.OffsetSize
41import           Basement.Monad
42
43import qualified Foundation.Collection as C
44import           Foundation.Numerical
45import           Data.Bits
46import           Foundation.Bits
47import           GHC.ST
48import qualified Data.List
49
50data Bitmap = Bitmap (CountOf Bool) (UArray Word32)
51    deriving (Typeable)
52
53data MutableBitmap st = MutableBitmap (CountOf Bool) (MUArray Word32 st)
54
55bitsPerTy :: Int
56bitsPerTy = 32
57
58shiftPerTy :: Int
59shiftPerTy = 5
60
61maskPerTy :: Int
62maskPerTy = 0x1f
63
64instance Show Bitmap where
65    show v = show (toList v)
66instance Eq Bitmap where
67    (==) = equal
68instance Ord Bitmap where
69    compare = vCompare
70instance Semigroup Bitmap where
71    (<>) = append
72instance Monoid Bitmap where
73    mempty  = empty
74    mappend = append
75    mconcat = concat
76
77type instance C.Element Bitmap = Bool
78
79instance IsList Bitmap where
80    type Item Bitmap = Bool
81    fromList = vFromList
82    toList = vToList
83
84instance C.InnerFunctor Bitmap where
85    imap = map
86
87instance C.Foldable Bitmap where
88    foldr = foldr
89    foldl' = foldl'
90    foldr' = foldr'
91
92instance C.Collection Bitmap where
93    null = null
94    length = length
95    elem e = Data.List.elem e . toList
96    maximum = any id . C.getNonEmpty
97    minimum = all id . C.getNonEmpty
98    all = all
99    any = any
100
101instance C.Sequential Bitmap where
102    take = take
103    drop = drop
104    splitAt = splitAt
105    revTake n = unoptimised (C.revTake n)
106    revDrop n = unoptimised (C.revDrop n)
107    splitOn = splitOn
108    break = break
109    breakEnd = breakEnd
110    span = span
111    filter = filter
112    reverse = reverse
113    snoc = snoc
114    cons = cons
115    unsnoc = unsnoc
116    uncons = uncons
117    intersperse = intersperse
118    find = find
119    sortBy = sortBy
120    singleton = fromList . (:[])
121    replicate n = fromList . C.replicate n
122
123instance C.IndexedCollection Bitmap where
124    (!) l n
125        | isOutOfBound n (length l) = Nothing
126        | otherwise                     = Just $ index l n
127    findIndex predicate c = loop 0
128      where
129        !len = length c
130        loop i
131            | i .==# len                  = Nothing
132            | predicate (unsafeIndex c i) = Just i
133            | otherwise                   = Nothing
134
135instance C.MutableCollection MutableBitmap where
136    type MutableFreezed MutableBitmap = Bitmap
137    type MutableKey MutableBitmap = Offset Bool
138    type MutableValue MutableBitmap = Bool
139
140    thaw = thaw
141    freeze = freeze
142    unsafeThaw = unsafeThaw
143    unsafeFreeze = unsafeFreeze
144
145    mutNew = new
146    mutUnsafeWrite = unsafeWrite
147    mutUnsafeRead = unsafeRead
148    mutWrite = write
149    mutRead = read
150
151
152
153bitmapIndex :: Offset Bool -> (Offset Word32, Int)
154bitmapIndex (Offset !i) = (Offset (i .>>. shiftPerTy), i .&. maskPerTy)
155{-# INLINE bitmapIndex #-}
156
157-- return the index in word32 quantity and mask to a bit in a bitmap
158{-
159bitmapAddr :: Int -> (# Int , Word #)
160bitmapAddr !i = (# idx, mask #)
161  where (!idx, !bitIdx) = bitmapIndex i
162        !mask = case bitIdx of
163                    0  -> 0x1
164                    1  -> 0x2
165                    2  -> 0x4
166                    3  -> 0x8
167                    4  -> 0x10
168                    5  -> 0x20
169                    6  -> 0x40
170                    7  -> 0x80
171                    8  -> 0x100
172                    9  -> 0x200
173                    10 -> 0x400
174                    11 -> 0x800
175                    12 -> 0x1000
176                    13 -> 0x2000
177                    14 -> 0x4000
178                    15 -> 0x8000
179                    16 -> 0x10000
180                    17 -> 0x20000
181                    18 -> 0x40000
182                    19 -> 0x80000
183                    20 -> 0x100000
184                    21 -> 0x200000
185                    22 -> 0x400000
186                    23 -> 0x800000
187                    24 -> 0x1000000
188                    25 -> 0x2000000
189                    26 -> 0x4000000
190                    27 -> 0x8000000
191                    28 -> 0x10000000
192                    29 -> 0x20000000
193                    30 -> 0x40000000
194                    _  -> 0x80000000
195-}
196
197thaw :: PrimMonad prim => Bitmap -> prim (MutableBitmap (PrimState prim))
198thaw (Bitmap len ba) = MutableBitmap len `fmap` C.thaw ba
199
200freeze :: PrimMonad prim => MutableBitmap (PrimState prim) -> prim Bitmap
201freeze (MutableBitmap len mba) = Bitmap len `fmap` C.freeze mba
202
203unsafeThaw :: PrimMonad prim => Bitmap -> prim (MutableBitmap (PrimState prim))
204unsafeThaw (Bitmap len ba) = MutableBitmap len `fmap` C.unsafeThaw ba
205
206unsafeFreeze :: PrimMonad prim => MutableBitmap (PrimState prim) -> prim Bitmap
207unsafeFreeze (MutableBitmap len mba) = Bitmap len `fmap` C.unsafeFreeze mba
208
209unsafeWrite :: PrimMonad prim => MutableBitmap (PrimState prim) -> Offset Bool -> Bool -> prim ()
210unsafeWrite (MutableBitmap _ ma) i v = do
211    let (idx, bitIdx) = bitmapIndex i
212    w <- A.unsafeRead ma idx
213    let w' = if v then setBit w bitIdx else clearBit w bitIdx
214    A.unsafeWrite ma idx w'
215{-# INLINE unsafeWrite #-}
216
217unsafeRead :: PrimMonad prim => MutableBitmap (PrimState prim) -> Offset Bool -> prim Bool
218unsafeRead (MutableBitmap _ ma) i = do
219    let (idx, bitIdx) = bitmapIndex i
220    flip testBit bitIdx `fmap` A.unsafeRead ma idx
221{-# INLINE unsafeRead #-}
222
223write :: PrimMonad prim => MutableBitmap (PrimState prim) -> Offset Bool -> Bool -> prim ()
224write mb n val
225    | isOutOfBound n len = primOutOfBound OOB_Write n len
226    | otherwise          = unsafeWrite mb n val
227  where
228    len = mutableLength mb
229{-# INLINE write #-}
230
231read :: PrimMonad prim => MutableBitmap (PrimState prim) -> Offset Bool -> prim Bool
232read mb n
233    | isOutOfBound n len = primOutOfBound OOB_Read n len
234    | otherwise        = unsafeRead mb n
235  where len = mutableLength mb
236{-# INLINE read #-}
237
238-- | Return the element at a specific index from a Bitmap.
239--
240-- If the index @n is out of bounds, an error is raised.
241index :: Bitmap -> Offset Bool -> Bool
242index bits n
243    | isOutOfBound n len = outOfBound OOB_Index n len
244    | otherwise          = unsafeIndex bits n
245  where len = length bits
246{-# INLINE index #-}
247
248-- | Return the element at a specific index from an array without bounds checking.
249--
250-- Reading from invalid memory can return unpredictable and invalid values.
251-- use 'index' if unsure.
252unsafeIndex :: Bitmap -> Offset Bool -> Bool
253unsafeIndex (Bitmap _ ba) n =
254    let (idx, bitIdx) = bitmapIndex n
255     in testBit (A.unsafeIndex ba idx) bitIdx
256
257{-# INLINE unsafeIndex #-}
258
259-----------------------------------------------------------------------
260-- higher level collection implementation
261-----------------------------------------------------------------------
262length :: Bitmap -> CountOf Bool
263length (Bitmap sz _) = sz
264
265mutableLength :: MutableBitmap st -> CountOf Bool
266mutableLength (MutableBitmap sz _) = sz
267
268empty :: Bitmap
269empty = Bitmap 0 mempty
270
271new :: PrimMonad prim => CountOf Bool -> prim (MutableBitmap (PrimState prim))
272new sz@(CountOf len) =
273    MutableBitmap sz <$> A.new nbElements
274  where
275    nbElements :: CountOf Word32
276    nbElements = CountOf ((len `alignRoundUp` bitsPerTy) .>>. shiftPerTy)
277
278-- | make an array from a list of elements.
279vFromList :: [Bool] -> Bitmap
280vFromList allBools = runST $ do
281    mbitmap <- new len
282    loop mbitmap 0 allBools
283  where
284    loop mb _ []     = unsafeFreeze mb
285    loop mb i (x:xs) = unsafeWrite mb i x >> loop mb (i+1) xs
286
287{-
288    runST $ do
289    mba <- A.new nbElements
290    ba  <- loop mba (0 :: Int) allBools
291    pure (Bitmap len ba)
292  where
293    loop mba _ [] = A.unsafeFreeze mba
294    loop mba i l  = do
295        let (l1, l2) = C.splitAt bitsPerTy l
296            w = toPacked l1
297        A.unsafeWrite mba i w
298        loop mba (i+1) l2
299
300    toPacked :: [Bool] -> Word32
301    toPacked l =
302        C.foldl' (.|.) 0 $ Prelude.zipWith (\b w -> if b then (1 `shiftL` w) else 0) l (C.reverse [0..31])
303-}
304    len        = C.length allBools
305
306-- | transform an array to a list.
307vToList :: Bitmap -> [Bool]
308vToList a = loop 0
309  where len = length a
310        loop i | i .==# len  = []
311               | otherwise = unsafeIndex a i : loop (i+1)
312
313-- | Check if two vectors are identical
314equal :: Bitmap -> Bitmap -> Bool
315equal a b
316    | la /= lb  = False
317    | otherwise = loop 0
318  where
319    !la = length a
320    !lb = length b
321    loop n | n .==# la = True
322           | otherwise = (unsafeIndex a n == unsafeIndex b n) && loop (n+1)
323
324-- | Compare 2 vectors
325vCompare :: Bitmap -> Bitmap -> Ordering
326vCompare a b = loop 0
327  where
328    !la = length a
329    !lb = length b
330    loop n
331        | n .==# la = if la == lb then EQ else LT
332        | n .==# lb = GT
333        | otherwise =
334            case unsafeIndex a n `compare` unsafeIndex b n of
335                EQ -> loop (n+1)
336                r  -> r
337
338-- | Append 2 arrays together by creating a new bigger array
339--
340-- TODO completely non optimized
341append :: Bitmap -> Bitmap -> Bitmap
342append a b = fromList $ toList a `mappend` toList b
343
344-- TODO completely non optimized
345concat :: [Bitmap] -> Bitmap
346concat l = fromList $ mconcat $ fmap toList l
347
348null :: Bitmap -> Bool
349null (Bitmap nbBits _) = nbBits == 0
350
351take :: CountOf Bool -> Bitmap -> Bitmap
352take nbElems bits@(Bitmap nbBits ba)
353    | nbElems <= 0      = empty
354    | nbElems >= nbBits = bits
355    | otherwise         = Bitmap nbElems ba -- TODO : although it work right now, take on the underlaying ba too
356
357drop :: CountOf Bool -> Bitmap -> Bitmap
358drop nbElems bits@(Bitmap nbBits _)
359    | nbElems <= 0      = bits
360    | nbElems >= nbBits = empty
361    | otherwise         = unoptimised (C.drop nbElems) bits
362        -- TODO: decide if we have drop easy by having a bit offset in the data structure
363        -- or if we need to shift stuff around making all the indexing slighlty more complicated
364
365splitAt :: CountOf Bool -> Bitmap -> (Bitmap, Bitmap)
366splitAt n v = (take n v, drop n v)
367
368-- unoptimised
369splitOn :: (Bool -> Bool) -> Bitmap -> [Bitmap]
370splitOn f bits = fmap fromList $ C.splitOn f $ toList bits
371
372-- unoptimised
373break :: (Bool -> Bool) -> Bitmap -> (Bitmap, Bitmap)
374break predicate v = findBreak 0
375  where
376    len = length v
377    findBreak i
378        | i .==# len = (v, empty)
379        | otherwise  =
380            if predicate (unsafeIndex v i)
381                then splitAt (offsetAsSize i) v
382                else findBreak (i+1)
383
384breakEnd :: (Bool -> Bool) -> Bitmap -> (Bitmap, Bitmap)
385breakEnd predicate = bimap fromList fromList . C.breakEnd predicate . toList
386
387span :: (Bool -> Bool) -> Bitmap -> (Bitmap, Bitmap)
388span p = break (not . p)
389
390map :: (Bool -> Bool) -> Bitmap -> Bitmap
391map f bits = unoptimised (fmap f) bits
392
393--mapIndex :: (Int -> Bool -> Bool) -> Bitmap -> Bitmap
394--mapIndex f Bitmap =
395
396cons :: Bool -> Bitmap -> Bitmap
397cons v l = unoptimised (C.cons v) l
398
399snoc :: Bitmap -> Bool -> Bitmap
400snoc l v = unoptimised (flip C.snoc v) l
401
402-- unoptimised
403uncons :: Bitmap -> Maybe (Bool, Bitmap)
404uncons b = fmap (second fromList) $ C.uncons $ toList b
405
406-- unoptimised
407unsnoc :: Bitmap -> Maybe (Bitmap, Bool)
408unsnoc b = fmap (first fromList) $ C.unsnoc $ toList b
409
410intersperse :: Bool -> Bitmap -> Bitmap
411intersperse b = unoptimised (C.intersperse b)
412
413find :: (Bool -> Bool) -> Bitmap -> Maybe Bool
414find predicate vec = loop 0
415  where
416    !len = length vec
417    loop i
418        | i .==# len = Nothing
419        | otherwise  =
420            let e = unsafeIndex vec i
421             in if predicate e then Just e else loop (i+1)
422
423sortBy :: (Bool -> Bool -> Ordering) -> Bitmap -> Bitmap
424sortBy by bits = unoptimised (C.sortBy by) bits
425
426filter :: (Bool -> Bool) -> Bitmap -> Bitmap
427filter predicate vec = unoptimised (Data.List.filter predicate) vec
428
429reverse :: Bitmap -> Bitmap
430reverse bits = unoptimised C.reverse bits
431
432foldr :: (Bool -> a -> a) -> a -> Bitmap -> a
433foldr f initialAcc vec = loop 0
434  where
435    len = length vec
436    loop i
437        | i .==# len = initialAcc
438        | otherwise  = unsafeIndex vec i `f` loop (i+1)
439
440foldr' :: (Bool -> a -> a) -> a -> Bitmap -> a
441foldr' = foldr
442
443foldl' :: (a -> Bool -> a) -> a -> Bitmap -> a
444foldl' f initialAcc vec = loop 0 initialAcc
445  where
446    len = length vec
447    loop i !acc
448        | i .==# len = acc
449        | otherwise  = loop (i+1) (f acc (unsafeIndex vec i))
450
451all :: (Bool -> Bool) -> Bitmap -> Bool
452all p bm = loop 0
453  where
454    len = length bm
455    loop !i
456      | i .==# len = True
457      | not $ p (unsafeIndex bm i) = False
458      | otherwise = loop (i + 1)
459
460any :: (Bool -> Bool) -> Bitmap -> Bool
461any p bm = loop 0
462  where
463    len = length bm
464    loop !i
465      | i .==# len = False
466      | p (unsafeIndex bm i) = True
467      | otherwise = loop (i + 1)
468
469unoptimised :: ([Bool] -> [Bool]) -> Bitmap -> Bitmap
470unoptimised f = vFromList . f . vToList
471