{-# LANGUAGE CPP           #-}
{-# LANGUAGE MagicHash     #-}
{-# LANGUAGE UnboxedTuples #-}
-- | Minimal bit vector implementation.
module Data.BloomFilter.BitVec64 (
    BitVec64 (..),
    unsafeIndex,
    prefetchIndex,
    MBitVec64 (..),
    new,
    unsafeWrite,
    unsafeRead,
    freeze,
    unsafeFreeze,
    thaw,
    unsafeRemWord64,
) where

import           Control.Monad.ST (ST)
import           Data.Bits
import           Data.Primitive.ByteArray (ByteArray (ByteArray),
                     newPinnedByteArray, setByteArray)
import qualified Data.Vector.Primitive as P
import qualified Data.Vector.Primitive.Mutable as MP
import           Data.Word (Word64, Word8)

import           GHC.Exts (Int (I#), prefetchByteArray0#, uncheckedIShiftRL#,
                     (+#))
import           GHC.ST (ST (ST))
import           GHC.Word (Word64 (W64#))

#if MIN_VERSION_base(4,17,0)
import           GHC.Exts (remWord64#)
#else
import           GHC.Exts (remWord#)
#endif

-- | Bit vector backed up by an array of Word64
--
-- This vector's offset and length are multiples of 64
newtype BitVec64 = BV64 (P.Vector Word64)
  deriving (BitVec64 -> BitVec64 -> Bool
(BitVec64 -> BitVec64 -> Bool)
-> (BitVec64 -> BitVec64 -> Bool) -> Eq BitVec64
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BitVec64 -> BitVec64 -> Bool
== :: BitVec64 -> BitVec64 -> Bool
$c/= :: BitVec64 -> BitVec64 -> Bool
/= :: BitVec64 -> BitVec64 -> Bool
Eq, Int -> BitVec64 -> ShowS
[BitVec64] -> ShowS
BitVec64 -> String
(Int -> BitVec64 -> ShowS)
-> (BitVec64 -> String) -> ([BitVec64] -> ShowS) -> Show BitVec64
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> BitVec64 -> ShowS
showsPrec :: Int -> BitVec64 -> ShowS
$cshow :: BitVec64 -> String
show :: BitVec64 -> String
$cshowList :: [BitVec64] -> ShowS
showList :: [BitVec64] -> ShowS
Show)

{-# INLINE unsafeIndex #-}
unsafeIndex :: BitVec64 -> Int -> Bool
unsafeIndex :: BitVec64 -> Int -> Bool
unsafeIndex (BV64 Vector Word64
bv) Int
i =
    Word64 -> Int -> Bool
unsafeTestBit (Vector Word64 -> Int -> Word64
forall a. Prim a => Vector a -> Int -> a
P.unsafeIndex Vector Word64
bv Int
j) Int
k
  where
    !j :: Int
j = Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR Int
i Int
6 -- `div` 64, bit index to Word64 index.
    !k :: Int
k = Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
63         -- `mod` 64, bit within Word64

{-# INLINE unsafeTestBit #-}
-- like testBit but using unsafeShiftL instead of shiftL
unsafeTestBit :: Word64 -> Int -> Bool
unsafeTestBit :: Word64 -> Int -> Bool
unsafeTestBit Word64
w Int
k = Word64
w Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. (Word64
1 Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
k) Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word64
0

{-# INLINE prefetchIndex #-}
prefetchIndex :: BitVec64 -> Int -> ST s ()
prefetchIndex :: forall s. BitVec64 -> Int -> ST s ()
prefetchIndex (BV64 (P.Vector (I# Int#
off#) Int
_ (ByteArray ByteArray#
ba#))) (I# Int#
i#) =
    STRep s () -> ST s ()
forall s a. STRep s a -> ST s a
ST (\State# s
s -> case ByteArray# -> Int# -> State# s -> State# s
forall d. ByteArray# -> Int# -> State# d -> State# d
prefetchByteArray0# ByteArray#
ba# (Int#
off# Int# -> Int# -> Int#
+# Int# -> Int# -> Int#
uncheckedIShiftRL# Int#
i# Int#
3#) State# s
s of
                State# s
s' -> (# State# s
s', () #))
    -- We only need to shiftR 3 here, not 6, because we're going from a bit
    -- offset to a byte offset for prefetch. Whereas in unsafeIndex, we go from
    -- a bit offset to a Word64 offset, so an extra shiftR 3, for 6 total.

newtype MBitVec64 s = MBV64 (P.MVector s Word64)

-- | Will create an explicitly pinned byte array if it is larger than 1 kB.
-- This is done because pinned byte arrays allow for more efficient
-- serialisation, but the definition of 'isByteArrayPinned' changed in GHC 9.6,
-- see <https://gitlab.haskell.org/ghc/ghc/-/issues/22255>.
--
-- TODO: remove this workaround once a solution exists, e.g. a new primop that
-- allows checking for implicit pinning.
new :: Word64 -> ST s (MBitVec64 s)
new :: forall s. Word64 -> ST s (MBitVec64 s)
new Word64
s
  | Int
numWords Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
128 = do
    MutableByteArray s
mba <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newPinnedByteArray Int
numBytes
    MutableByteArray (PrimState (ST s))
-> Int -> Int -> Word8 -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> Int -> a -> m ()
setByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
mba Int
0 Int
numBytes (Word8
0 :: Word8)
    MBitVec64 s -> ST s (MBitVec64 s)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return (MVector s Word64 -> MBitVec64 s
forall s. MVector s Word64 -> MBitVec64 s
MBV64 (Int -> Int -> MutableByteArray s -> MVector s Word64
forall s a. Int -> Int -> MutableByteArray s -> MVector s a
P.MVector Int
0 Int
numWords MutableByteArray s
mba))
  | Bool
otherwise =
    MVector s Word64 -> MBitVec64 s
forall s. MVector s Word64 -> MBitVec64 s
MBV64 (MVector s Word64 -> MBitVec64 s)
-> ST s (MVector s Word64) -> ST s (MBitVec64 s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> ST s (MVector (PrimState (ST s)) Word64)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MVector (PrimState m) a)
MP.new Int
numWords
  where
    !numWords :: Int
numWords = Word64 -> Int
w2i (Word64 -> Word64
roundUpTo64 Word64
s)
    !numBytes :: Int
numBytes = Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftL Int
numWords Int
3 -- * 8

unsafeWrite :: MBitVec64 s -> Word64 -> Bool -> ST s ()
unsafeWrite :: forall s. MBitVec64 s -> Word64 -> Bool -> ST s ()
unsafeWrite (MBV64 MVector s Word64
mbv) Word64
i Bool
x = do
    MVector (PrimState (ST s)) Word64
-> (Word64 -> Word64) -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MP.unsafeModify MVector s Word64
MVector (PrimState (ST s)) Word64
mbv (\Word64
w -> if Bool
x then Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
setBit Word64
w (Word64 -> Int
w2i Word64
k) else Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
clearBit Word64
w (Word64 -> Int
w2i Word64
k)) (Word64 -> Int
w2i Word64
j)
  where
    !j :: Word64
j = Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
i Int
6 -- `div` 64
    !k :: Word64
k = Word64
i Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
63         -- `mod` 64

unsafeRead :: MBitVec64 s -> Word64 -> ST s Bool
unsafeRead :: forall s. MBitVec64 s -> Word64 -> ST s Bool
unsafeRead (MBV64 MVector s Word64
mbv) Word64
i = do
    !Word64
w <- MVector (PrimState (ST s)) Word64 -> Int -> ST s Word64
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MVector (PrimState m) a -> Int -> m a
MP.unsafeRead MVector s Word64
MVector (PrimState (ST s)) Word64
mbv (Word64 -> Int
w2i Word64
j)
    Bool -> ST s Bool
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> ST s Bool) -> Bool -> ST s Bool
forall a b. (a -> b) -> a -> b
$! Word64 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word64
w (Word64 -> Int
w2i Word64
k)
  where
    !j :: Word64
j = Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
i Int
6 -- `div` 64
    !k :: Word64
k = Word64
i Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
63         -- `mod` 64

freeze :: MBitVec64 s -> ST s BitVec64
freeze :: forall s. MBitVec64 s -> ST s BitVec64
freeze (MBV64 MVector s Word64
mbv) = Vector Word64 -> BitVec64
BV64 (Vector Word64 -> BitVec64)
-> ST s (Vector Word64) -> ST s BitVec64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Word64 -> ST s (Vector Word64)
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
P.freeze MVector s Word64
MVector (PrimState (ST s)) Word64
mbv

unsafeFreeze :: MBitVec64 s -> ST s BitVec64
unsafeFreeze :: forall s. MBitVec64 s -> ST s BitVec64
unsafeFreeze (MBV64 MVector s Word64
mbv) = Vector Word64 -> BitVec64
BV64 (Vector Word64 -> BitVec64)
-> ST s (Vector Word64) -> ST s BitVec64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Word64 -> ST s (Vector Word64)
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
P.unsafeFreeze MVector s Word64
MVector (PrimState (ST s)) Word64
mbv

thaw :: BitVec64 -> ST s (MBitVec64 s)
thaw :: forall s. BitVec64 -> ST s (MBitVec64 s)
thaw (BV64 Vector Word64
bv) = MVector s Word64 -> MBitVec64 s
forall s. MVector s Word64 -> MBitVec64 s
MBV64 (MVector s Word64 -> MBitVec64 s)
-> ST s (MVector s Word64) -> ST s (MBitVec64 s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector Word64 -> ST s (MVector (PrimState (ST s)) Word64)
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
P.thaw Vector Word64
bv

-- this may overflow, but so be it (1^64 bits is a lot)
roundUpTo64 :: Word64 -> Word64
roundUpTo64 :: Word64 -> Word64
roundUpTo64 Word64
i = Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR (Word64
i Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
63) Int
6

-- | Like 'rem' but does not check for division by 0.
unsafeRemWord64 :: Word64 -> Word64 -> Word64
#if MIN_VERSION_base(4,17,0)
unsafeRemWord64 :: Word64 -> Word64 -> Word64
unsafeRemWord64 (W64# Word64#
x#) (W64# Word64#
y#) = Word64# -> Word64
W64# (Word64#
x# Word64# -> Word64# -> Word64#
`remWord64#` Word64#
y#)
#else
unsafeRemWord64 (W64# x#) (W64# y#) = W64# (x# `remWord#` y#)
#endif

w2i :: Word64 -> Int
w2i :: Word64 -> Int
w2i = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE w2i #-}