never executed always true always false
1 {-# LANGUAGE CPP #-}
2 {-# LANGUAGE MagicHash #-}
3 {-# LANGUAGE UnboxedTuples #-}
4 -- | Minimal bit vector implementation.
5 module Data.BloomFilter.BitVec64 (
6 BitVec64 (..),
7 unsafeIndex,
8 prefetchIndex,
9 MBitVec64 (..),
10 new,
11 unsafeWrite,
12 unsafeRead,
13 freeze,
14 unsafeFreeze,
15 thaw,
16 unsafeRemWord64,
17 ) where
18
19 import Control.Monad.ST (ST)
20 import Data.Bits
21 import Data.Primitive.ByteArray (ByteArray (ByteArray),
22 newPinnedByteArray, setByteArray)
23 import qualified Data.Vector.Primitive as P
24 import qualified Data.Vector.Primitive.Mutable as MP
25 import Data.Word (Word64, Word8)
26
27 import GHC.Exts (Int (I#), prefetchByteArray0#, uncheckedIShiftRL#,
28 (+#))
29 import GHC.ST (ST (ST))
30 import GHC.Word (Word64 (W64#))
31
32 #if MIN_VERSION_base(4,17,0)
33 import GHC.Exts (remWord64#)
34 #else
35 import GHC.Exts (remWord#)
36 #endif
37
38 -- | Bit vector backed up by an array of Word64
39 --
40 -- This vector's offset and length are multiples of 64
41 newtype BitVec64 = BV64 (P.Vector Word64)
42 deriving (Eq, Show)
43
44 {-# INLINE unsafeIndex #-}
45 unsafeIndex :: BitVec64 -> Int -> Bool
46 unsafeIndex (BV64 bv) i =
47 unsafeTestBit (P.unsafeIndex bv j) k
48 where
49 !j = unsafeShiftR i 6 -- `div` 64, bit index to Word64 index.
50 !k = i .&. 63 -- `mod` 64, bit within Word64
51
52 {-# INLINE unsafeTestBit #-}
53 -- like testBit but using unsafeShiftL instead of shiftL
54 unsafeTestBit :: Word64 -> Int -> Bool
55 unsafeTestBit w k = w .&. (1 `unsafeShiftL` k) /= 0
56
57 {-# INLINE prefetchIndex #-}
58 prefetchIndex :: BitVec64 -> Int -> ST s ()
59 prefetchIndex (BV64 (P.Vector (I# off#) _ (ByteArray ba#))) (I# i#) =
60 ST (\s -> case prefetchByteArray0# ba# (off# +# uncheckedIShiftRL# i# 3#) s of
61 s' -> (# s', () #))
62 -- We only need to shiftR 3 here, not 6, because we're going from a bit
63 -- offset to a byte offset for prefetch. Whereas in unsafeIndex, we go from
64 -- a bit offset to a Word64 offset, so an extra shiftR 3, for 6 total.
65
66 newtype MBitVec64 s = MBV64 (P.MVector s Word64)
67
68 -- | Will create an explicitly pinned byte array if it is larger than 1 kB.
69 -- This is done because pinned byte arrays allow for more efficient
70 -- serialisation, but the definition of 'isByteArrayPinned' changed in GHC 9.6,
71 -- see <https://gitlab.haskell.org/ghc/ghc/-/issues/22255>.
72 --
73 -- TODO: remove this workaround once a solution exists, e.g. a new primop that
74 -- allows checking for implicit pinning.
75 new :: Word64 -> ST s (MBitVec64 s)
76 new s
77 | numWords >= 128 = do
78 mba <- newPinnedByteArray numBytes
79 setByteArray mba 0 numBytes (0 :: Word8)
80 return (MBV64 (P.MVector 0 numWords mba))
81 | otherwise =
82 MBV64 <$> MP.new numWords
83 where
84 !numWords = w2i (roundUpTo64 s)
85 !numBytes = unsafeShiftL numWords 3 -- * 8
86
87 unsafeWrite :: MBitVec64 s -> Word64 -> Bool -> ST s ()
88 unsafeWrite (MBV64 mbv) i x = do
89 MP.unsafeModify mbv (\w -> if x then setBit w (w2i k) else clearBit w (w2i k)) (w2i j)
90 where
91 !j = unsafeShiftR i 6 -- `div` 64
92 !k = i .&. 63 -- `mod` 64
93
94 unsafeRead :: MBitVec64 s -> Word64 -> ST s Bool
95 unsafeRead (MBV64 mbv) i = do
96 !w <- MP.unsafeRead mbv (w2i j)
97 return $! testBit w (w2i k)
98 where
99 !j = unsafeShiftR i 6 -- `div` 64
100 !k = i .&. 63 -- `mod` 64
101
102 freeze :: MBitVec64 s -> ST s BitVec64
103 freeze (MBV64 mbv) = BV64 <$> P.freeze mbv
104
105 unsafeFreeze :: MBitVec64 s -> ST s BitVec64
106 unsafeFreeze (MBV64 mbv) = BV64 <$> P.unsafeFreeze mbv
107
108 thaw :: BitVec64 -> ST s (MBitVec64 s)
109 thaw (BV64 bv) = MBV64 <$> P.thaw bv
110
111 -- this may overflow, but so be it (1^64 bits is a lot)
112 roundUpTo64 :: Word64 -> Word64
113 roundUpTo64 i = unsafeShiftR (i + 63) 6
114
115 -- | Like 'rem' but does not check for division by 0.
116 unsafeRemWord64 :: Word64 -> Word64 -> Word64
117 #if MIN_VERSION_base(4,17,0)
118 unsafeRemWord64 (W64# x#) (W64# y#) = W64# (x# `remWord64#` y#)
119 #else
120 unsafeRemWord64 (W64# x#) (W64# y#) = W64# (x# `remWord#` y#)
121 #endif
122
123 w2i :: Word64 -> Int
124 w2i = fromIntegral
125 {-# INLINE w2i #-}