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 #-}