{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE BinaryLiterals #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UnboxedSums #-}
{-# LANGUAGE UnboxedTuples #-}
module Data.Bytes.HashMap.Word
( Map
, lookup
, fromList
, fromTrustedList
, fromListWith
, distribution
, distinctEntropies
) where
import Prelude hiding (lookup)
import Data.Bytes.Types (Bytes(Bytes))
import Data.Int (Int32)
import Data.Primitive (ByteArray,ByteArray(..),PrimArray(..))
import GHC.Exts (Int(I#),ByteArray#,ArrayArray#,Int#,Word#)
import GHC.Word (Word(W#),Word32)
import Data.Primitive.Unlifted.Array (UnliftedArray(..))
import System.Entropy (CryptHandle)
import qualified Data.Bytes as Bytes
import qualified Data.Bytes.Hash as Hash
import qualified Data.Bytes.HashMap as Lifted
import qualified Data.Bytes.HashMap.Internal as Lifted
import qualified Data.List as List
import qualified Data.Primitive as PM
import qualified Data.Primitive.Unlifted.Array as PM
import qualified GHC.Exts as Exts
data Map = Map
!ByteArray
!(UnliftedArray ByteArray)
!(PrimArray Int32)
!(UnliftedArray ByteArray)
!(PrimArray Word)
fromLifted :: Lifted.Map Word -> Map
fromLifted (Lifted.Map a b c d e) = Map a b c d (Exts.fromList (Exts.toList e))
fromList :: CryptHandle -> [(Bytes,Word)] -> IO Map
fromList h = fmap fromLifted . Lifted.fromList h
fromListWith ::
CryptHandle
-> (Word -> Word -> Word)
-> [(Bytes,Word)]
-> IO Map
fromListWith h c xs = fmap fromLifted (Lifted.fromListWith h c xs)
fromTrustedList :: [(Bytes,Word)] -> Map
fromTrustedList = fromLifted . Lifted.fromTrustedList
lookup :: Bytes -> Map -> Maybe Word
{-# inline lookup #-}
lookup
(Bytes (ByteArray keyArr) (I# keyOff) (I# keyLen))
(Map (ByteArray entropyA) (UnliftedArray entropies) (PrimArray offsets) (UnliftedArray keys) (PrimArray vals)) =
case lookup# (# keyArr,keyOff,keyLen #) (# entropyA,entropies,offsets,keys,vals #) of
(# (# #) | #) -> Nothing
(# | v #) -> Just (W# v)
lookup# ::
(# ByteArray#, Int#, Int# #)
-> (# ByteArray#, ArrayArray#, ByteArray#, ArrayArray#, ByteArray# #)
-> (# (# #) | Word# #)
{-# noinline lookup# #-}
lookup# (# keyArr#, keyOff#, keyLen# #) (# entropyA#, entropies#, offsets#, keys#, vals# #)
| sz == 0 = (# (# #) | #)
| PM.sizeofByteArray entropyA < reqEntropy = (# (# #) | #)
| ixA <- w2i (unsafeRem (upW32 (Hash.bytes entropyA key)) (i2w sz))
, entropyB <- PM.indexUnliftedArray entropies ixA
, offset <- fromIntegral @Int32 @Int (PM.indexPrimArray offsets ixA) =
case sameByteArray entropyA entropyB of
1# | ix <- ixA
, offsetIx <- offset + ix
, bytesEqualsByteArray key (PM.indexUnliftedArray keys offsetIx)
, !(W# v) <- PM.indexPrimArray vals offsetIx -> (# | v #)
| otherwise -> (# (# #) | #)
_ | PM.sizeofByteArray entropyB >= reqEntropy
, ix <- w2i (unsafeRem (upW32 (Hash.bytes entropyB key)) (i2w sz))
, offsetIx <- offset + ix
, bytesEqualsByteArray key (PM.indexUnliftedArray keys offsetIx)
, !(W# v) <- PM.indexPrimArray vals offsetIx -> (# | v #)
| otherwise -> (# (# #) | #)
where
sz = PM.sizeofUnliftedArray entropies
reqEntropy = w2i (requiredEntropy (i2w (Bytes.length key)))
key = Bytes (ByteArray keyArr#) (I# keyOff#) (I# keyLen#)
entropyA = ByteArray entropyA#
entropies = UnliftedArray entropies# :: UnliftedArray ByteArray
keys = UnliftedArray keys# :: UnliftedArray ByteArray
vals = PrimArray vals# :: PrimArray Word
offsets = PrimArray offsets# :: PrimArray Int32
unsafeRem :: Word -> Word -> Word
unsafeRem (W# a) (W# b) = W# (Exts.remWord# a b)
i2w :: Int -> Word
i2w = fromIntegral
requiredEntropy :: Word -> Word
requiredEntropy n = 8 * n + 8
w2i :: Word -> Int
w2i = fromIntegral
bytesEqualsByteArray :: Bytes -> ByteArray -> Bool
bytesEqualsByteArray (Bytes arr1 off1 len1) arr2
| len1 /= PM.sizeofByteArray arr2 = False
| otherwise = compareByteArrays arr1 off1 arr2 0 len1 == EQ
compareByteArrays :: ByteArray -> Int -> ByteArray -> Int -> Int -> Ordering
compareByteArrays (ByteArray ba1#) (I# off1#) (ByteArray ba2#) (I# off2#) (I# n#) =
compare (I# (Exts.compareByteArrays# ba1# off1# ba2# off2# n#)) 0
upW32 :: Word32 -> Word
upW32 = fromIntegral
distribution :: Map -> [(Int,Int)]
distribution (Map entropy entropies offsets keys vals) = Lifted.distribution
(Lifted.Map entropy entropies offsets keys (Exts.fromList (Exts.toList vals)))
distinctEntropies :: Map -> Int
distinctEntropies (Map entropy entropies _ _ _) =
List.length (List.group (List.sort (entropy : Exts.toList entropies)))
sameByteArray :: ByteArray -> ByteArray -> Int#
sameByteArray (ByteArray x) (ByteArray y) =
Exts.sameMutableByteArray# (Exts.unsafeCoerce# x) (Exts.unsafeCoerce# y)