{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

-- |
-- Module: Data.Cuckoo.Internal
-- Copyright: Copyright © 2019 Lars Kuhtz <lakuhtz@gmail.com>
-- License: BSD3
-- Maintainer: Lars Kuhtz <lakuhtz@gmail.com>
-- Stability: experimental
--
-- Internal Utilities. No guarantee is made about the stability of these
-- functions. Changes to these function won't be announced in the CHANGELOG and
-- are not reflected in the package version.
--
module Data.Cuckoo.Internal
( w
, int
, fit
, intFit
, nextPowerOfTwo
, intNextPowerOfTwo
, set
, get

-- * Hash functions
, sip
, fnv1a
, fnv1a_bytes
, sip_bytes
, sip2
) where

import Control.Monad.Primitive

import Data.Bits
import qualified Data.ByteArray as BA
import qualified Data.ByteArray.Hash as BA
import qualified Data.ByteArray.Pack as BA
import Data.Primitive.ByteArray

import Foreign

import GHC.Exts
import GHC.TypeLits

-- | Reify type level 'Nat' into 'Int' value.
--
w :: forall (n :: Nat) . KnownNat n => Int
w :: Int
w = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
int (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Proxy# n -> Integer
forall (n :: Nat). KnownNat n => Proxy# n -> Integer
natVal' @n Proxy# n
forall k (a :: k). Proxy# a
proxy#
{-# INLINE w #-}

-- | An shorter alias for 'fromIntegral'.
--
int :: Integral a => Num b => a -> b
int :: a -> b
int = a -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE int #-}

-- | @fit a b@ computes how many @b@s are needed to fit @a@, i.e.
-- \(\left\lceil\frac{a}{b}\right\rceil\).
--
-- For instance,
--
-- >>> fit 7 3
-- 3
--
-- >>> fit 6 3
-- 2
--
fit :: Real a => Real b => Integral c => a -> b -> c
fit :: a -> b -> c
fit a
a b
b = forall b. (RealFrac Double, Integral b) => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling @Double (Double -> c) -> Double -> c
forall a b. (a -> b) -> a -> b
$ a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac a
a Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ b -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac b
b
{-# INLINE fit #-}

-- | @fit a b@ computes how many @b@s are needed to fit @a@, i.e.
-- \(\left\lceil\frac{a}{b}\right\rceil\).
--
-- For instance,
--
-- >>> intFit 7 3
-- 3
--
-- >>> intFit 6 3
-- 2
--
intFit :: Integral a => Integral b => a -> b -> a
intFit :: a -> b -> a
intFit a
a b
b = a
1 a -> a -> a
forall a. Num a => a -> a -> a
+ (a
a a -> a -> a
forall a. Num a => a -> a -> a
- a
1) a -> a -> a
forall a. Integral a => a -> a -> a
`div` b -> a
forall a b. (Integral a, Num b) => a -> b
int b
b
{-# INLINE intFit #-}

-- | @nextPowerOfTwo a@ computes the smallest power of two that is larger or
-- equal than @a@.
--
nextPowerOfTwo :: Real a => Integral b => a -> b
nextPowerOfTwo :: a -> b
nextPowerOfTwo a
x = b
2 b -> Int -> b
forall a b. (Num a, Integral b) => a -> b -> a
^ Double -> Int
forall a b. (RealFrac a, Integral b) => a -> b
ceiling @Double @Int (Double -> Double -> Double
forall a. Floating a => a -> a -> a
logBase Double
2 (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac a
x)
{-# INLINE nextPowerOfTwo #-}

-- | @nextPowerOfTwo a@ computes the smallest power of two that is larger or
-- equal than @a@.
--
intNextPowerOfTwo :: Int -> Int
intNextPowerOfTwo :: Int -> Int
intNextPowerOfTwo Int
0 = Int
1
intNextPowerOfTwo Int
x = Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftL` (Int -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
{-# INLINE intNextPowerOfTwo #-}

-- | Computes a 64 bit Fnv1a hash for a value that has an 'Storable' instance.
--
-- The first argument is use as a salt.
--
fnv1a
    :: Storable a
    => Int
        -- ^ Salt
    -> a
        -- ^ Value that is hashes
    -> Word64
fnv1a :: Int -> a -> Word64
fnv1a Int
s a
x = Word64
r
  where
    Right (BA.FnvHash64 Word64
r) = Bytes -> FnvHash64
forall ba. ByteArrayAccess ba => ba -> FnvHash64
BA.fnv1a_64Hash
        (Bytes -> FnvHash64)
-> Either String Bytes -> Either String FnvHash64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Packer () -> Either String Bytes
forall byteArray a.
ByteArray byteArray =>
Int -> Packer a -> Either String byteArray
BA.fill @BA.Bytes (Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ a -> Int
forall a. Storable a => a -> Int
sizeOf a
x) (Int -> Packer ()
forall storable. Storable storable => storable -> Packer ()
BA.putStorable Int
s Packer () -> Packer () -> Packer ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> a -> Packer ()
forall storable. Storable storable => storable -> Packer ()
BA.putStorable a
x)
{-# INLINE fnv1a #-}

-- | Computes a 64 bit Fnv1a hash for a value that is an instance of
-- 'BA.ByteArrayAccess'.
--
-- The first argument is use as a salt.
--
fnv1a_bytes
    :: BA.ByteArrayAccess a
    => Int
        -- ^ Salt
    -> a
        -- ^ Value that is hashes
    -> Word64
fnv1a_bytes :: Int -> a -> Word64
fnv1a_bytes Int
s a
x = Word64
r
  where
    Right (BA.FnvHash64 Word64
r) = Bytes -> FnvHash64
forall ba. ByteArrayAccess ba => ba -> FnvHash64
BA.fnv1a_64Hash
        (Bytes -> FnvHash64)
-> Either String Bytes -> Either String FnvHash64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Packer () -> Either String Bytes
forall byteArray a.
ByteArray byteArray =>
Int -> Packer a -> Either String byteArray
BA.fill @BA.Bytes (Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ a -> Int
forall ba. ByteArrayAccess ba => ba -> Int
BA.length a
x) (Int -> Packer ()
forall storable. Storable storable => storable -> Packer ()
BA.putStorable Int
s Packer () -> Packer () -> Packer ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> a -> Packer ()
forall ba. ByteArrayAccess ba => ba -> Packer ()
BA.putBytes a
x)
{-# INLINE fnv1a_bytes #-}

-- | Computes a Sip hash for a value that has an 'Storable' instance.
--
-- The first argument is a salt value that is used to derive the key for the
-- hash computation.
--
sip
    :: Storable a
    => Int
        -- ^ Salt
    -> a
        -- ^ Value that is hashes
    -> Word64
sip :: Int -> a -> Word64
sip Int
s a
x = Word64
r
  where
    Right (BA.SipHash Word64
r) = SipKey -> Bytes -> SipHash
forall ba. ByteArrayAccess ba => SipKey -> ba -> SipHash
BA.sipHash (Word64 -> Word64 -> SipKey
BA.SipKey (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
int Int
s) Word64
914279)
        (Bytes -> SipHash) -> Either String Bytes -> Either String SipHash
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Packer () -> Either String Bytes
forall byteArray a.
ByteArray byteArray =>
Int -> Packer a -> Either String byteArray
BA.fill @BA.Bytes (a -> Int
forall a. Storable a => a -> Int
sizeOf a
x) (a -> Packer ()
forall storable. Storable storable => storable -> Packer ()
BA.putStorable a
x)
{-# INLINE sip #-}

-- | Computes a Sip hash for a value that is an instance of
-- 'BA.ByteArrayAccess'.
--
-- The first argument is a salt value that is used to derive the key for the
-- hash computation.
--
sip_bytes
    :: BA.ByteArrayAccess a
    => Int
        -- ^ Salt
    -> a
        -- ^ Value that is hashes
    -> Word64
sip_bytes :: Int -> a -> Word64
sip_bytes Int
s a
x = Word64
r
  where
    Right (BA.SipHash Word64
r) = SipKey -> Bytes -> SipHash
forall ba. ByteArrayAccess ba => SipKey -> ba -> SipHash
BA.sipHash (Word64 -> Word64 -> SipKey
BA.SipKey (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
int Int
s) Word64
1043639)
        (Bytes -> SipHash) -> Either String Bytes -> Either String SipHash
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Packer () -> Either String Bytes
forall byteArray a.
ByteArray byteArray =>
Int -> Packer a -> Either String byteArray
BA.fill @BA.Bytes (a -> Int
forall ba. ByteArrayAccess ba => ba -> Int
BA.length a
x) (a -> Packer ()
forall ba. ByteArrayAccess ba => ba -> Packer ()
BA.putBytes a
x)
{-# INLINE sip_bytes #-}

-- | An version of a Sip hash that is used internally. In order to avoid
-- dependencies between different hash computations, it shouldn't be used in the
-- implementation of instances of 'Data.Cuckoo.CuckooFilterHash'.
--
sip2 :: Storable a => Int -> a -> Word64
sip2 :: Int -> a -> Word64
sip2 Int
s a
x = Word64
r
  where
    Right (BA.SipHash Word64
r) = SipKey -> Bytes -> SipHash
forall ba. ByteArrayAccess ba => SipKey -> ba -> SipHash
BA.sipHash (Word64 -> Word64 -> SipKey
BA.SipKey Word64
994559 (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
int Int
s Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
713243))
        (Bytes -> SipHash) -> Either String Bytes -> Either String SipHash
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Packer () -> Either String Bytes
forall byteArray a.
ByteArray byteArray =>
Int -> Packer a -> Either String byteArray
BA.fill @BA.Bytes (a -> Int
forall a. Storable a => a -> Int
sizeOf a
x) (a -> Packer ()
forall storable. Storable storable => storable -> Packer ()
BA.putStorable a
x)
{-# INLINE sip2 #-}

-- | Write a 'Word64' value into a 'Word32' aligned 'MutableByteArray'
--
set
    :: PrimMonad m
    => MutableByteArray (PrimState m)
    -> Int
        -- ^ index in terms of 'Word32'
    -> Word64
        -- ^ 'Word64' value that is written
    -> m ()
set :: MutableByteArray (PrimState m) -> Int -> Word64 -> m ()
set MutableByteArray (PrimState m)
x Int
i Word64
c = do
    MutableByteArray (PrimState m) -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray @Word32 MutableByteArray (PrimState m)
x Int
i (Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
int Word64
c)
    MutableByteArray (PrimState m) -> Int -> Word32 -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray @Word32 MutableByteArray (PrimState m)
x (Int -> Int
forall a. Enum a => a -> a
succ Int
i) (Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
int (Word64 -> Word32) -> Word64 -> Word32
forall a b. (a -> b) -> a -> b
$ Word64
c Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
32)
{-# INLINE set #-}

-- | Get a 'Word64' from a 'Word32' aligned 'MutableByteArray'.
--
get
    :: PrimMonad m
    => MutableByteArray (PrimState m)
        -- ^ byte array
    -> Int
        -- ^ index in terms of 'Word32'
    -> m Word64
        -- ^ Word64 value that contains the result bits
get :: MutableByteArray (PrimState m) -> Int -> m Word64
get MutableByteArray (PrimState m)
x Int
i = do
    Word32
a <- MutableByteArray (PrimState m) -> Int -> m Word32
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray @Word32 MutableByteArray (PrimState m)
x Int
i
    Word32
b <- MutableByteArray (PrimState m) -> Int -> m Word32
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray @Word32 MutableByteArray (PrimState m)
x (Int -> Int
forall a. Enum a => a -> a
succ Int
i)

    -- TODO check of for host byte order
    -- Here we assume littel endian
    Word64 -> m Word64
forall (m :: * -> *) a. Monad m => a -> m a
return (Word64 -> m Word64) -> Word64 -> m Word64
forall a b. (a -> b) -> a -> b
$! Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
int Word32
a Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ (Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
int Word32
b Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
32)
{-# INLINE get #-}