{-# LANGUAGE BangPatterns, UnboxedTuples #-}

module CityHash
    (
      hash
    ) where

import Control.Applicative
import Data.Bits
import Data.Word (Word64)
import Foreign.Storable (peek)
import Foreign.Ptr
import Prelude hiding (fst, snd)

data Pair64 = (:*) {
      fst :: {-# UNPACK #-} !Word64
    , snd :: {-# UNPACK #-} !Word64
    }

hash :: Ptr a -> Int -> IO Word64
hash s len0 = do
  let s_len = s `plusPtr` len0
      len64 = fromIntegral len0
  x0 <- peek64 s
  y0 <- (`xor` k1) <$> peek64 (s_len `plusPtr` (-16))
  z0 <- (`xor` k0) <$> peek64 (s_len `plusPtr` (-56))
  !v0 <- weakHashPtr (s_len `plusPtr` (-64)) len64 y0
  !w0 <- weakHashPtr (s_len `plusPtr` (-32)) (len64 * k1) k0
  let z1 = shiftMix (snd v0) * k1
      x1 = rotateR (z1 + x0) 39 * k1
      y1 = rotateR y0 33 * k1
  loop s ((len0 - 1) .&. mask) x1 y1 z1 v0 w0
 where
  mask = 0x7fffffffffffffc0
  loop !p !len !x0 !y0 !z0 !v0 !w0 = do
    p_16 <- peek64 (p `plusPtr` 16)
    p_48 <- peek64 (p `plusPtr` 48)
    let x1 = rotateR (x0 + y0 + fst v0 + p_16) 37 * k1
        y1 = rotateR (y0 + snd v0 + p_48) 42 * k1
        x2 = x1 `xor` snd w0
        y2 = y1 `xor` fst v0
        z1 = rotateR (z0 `xor` fst w0) 33
    !v1 <- weakHashPtr p (snd v0 * k1) (x2 + fst w0)
    !w1 <- weakHashPtr (p `plusPtr` 32) (z1 + snd w0) y2
    let (# x3, z2 #) = (# z0, x2 #)
    if len == 64
      then return $! hash16 (fst v1 `hash16` fst w1 + shiftMix y2 * k1 + z2)
                            (snd v1 `hash16` snd w1 + x2)
      else loop (p `plusPtr` 64) (len - 64) x3 y2 z2 v1 w1

shiftMix :: Word64 -> Word64
shiftMix v = v `xor` (v `shiftR` 47)

weakHash :: Word64 -> Word64 -> Word64 -> Word64 -> Word64 -> Word64 -> Pair64
weakHash w x y z a0 b0 = (a + z) :* (b + a1)
  where
    a1 = a0 + w
    b1 = rotate (b0 + a1 + z) 21
    a  = a1 + x + y
    b  = rotate a 44
{-# INLINE weakHash #-}

weakHashPtr :: Ptr a -> Word64 -> Word64 -> IO Pair64
weakHashPtr p a b = weakHash <$> peek64 p
                             <*> peek64 (p `plusPtr` 8)
                             <*> peek64 (p `plusPtr` 16)
                             <*> peek64 (p `plusPtr` 24)
                             <*> pure a
                             <*> pure b
{-# INLINE weakHashPtr #-}

hash16 :: Word64 -> Word64 -> Word64
hash16 hi lo = b * k
  where
    a  = a0 `xor` (a0 `shiftR` 47)
    b  = b0 `xor` (b0 `shiftR` 47)
    b0 = hi `xor` a * k
    a0 = (lo `xor` hi) * k
    k  = 0x9ddfea08eb382d69

k0, k1, k2, k3 :: Word64
k0 = 0xc3a5c85c97cb3127
k1 = 0xb492b66fbe98f273
{-# INLINE k1 #-}
k2 = 0x9ae16a3b2f90404f
k3 = 0xc949d7c7509e6557

peek64 :: Ptr a -> IO Word64
peek64 = peek . castPtr
{-# INLINE peek64 #-}