SIP hash in pure Haskell

SIP.hs

{-# LANGUAGE UnicodeSyntax #-}

-- SIP hash in pure Haskell
-- Original C reference implementation taken from (github.com/veorq/SipHash)
-- Translated to Haskell by Adam R. Nelson (github.com/ar-nelson)
--------------------------------------------------------------------------------

module Data.Digest.SIP(sipHash) where

import Control.Monad
import Control.Monad.ST
import Data.Bits
import Data.ByteString
import Data.Function
import Data.STRef
import Data.Word
import Prelude hiding (drop, length, null, splitAt, take)

cRound ∷ Int
cRound = 2

dRound ∷ Int
dRound = 4

-- Computes the SIP hash of a ByteString. The second argument is a 128-bit
-- secret key, in the form of two Word64s.

sipHash ∷ ByteString → (Word64, Word64) → Word64
sipHash bytes (k0, k1) = runST $
  do v0 ← newSTRef (0x736f6d6570736575 `xor` k0)
     v1 ← newSTRef (0x646f72616e646f6d `xor` k1)
     v2 ← newSTRef (0x6c7967656e657261 `xor` k0)
     v3 ← newSTRef (0x7465646279746573 `xor` k1)

     let sipRound = do { v0 += v1; rotl v1 13; v1 ^= v0; rotl v0 32
                       ; v2 += v3; rotl v3 16; v3 ^= v2
                       ; v0 += v3; rotl v3 21; v3 ^= v0
                       ; v2 += v1; rotl v1 17; v1 ^= v2; rotl v2 32
                       }
           where a += b = readSTRef b >>= \b' → modifySTRef' a (+ b')
                 a ^= b = readSTRef b >>= \b' → modifySTRef' a (`xor` b')
                 rotl x b = modifySTRef' x $
                   \n → shiftL n b .|. shiftR n (64 - b)

     -- Redefine ^= as non-monadic in second arg for the rest of the function.
     let a ^= b = modifySTRef' a (`xor` b)

     end ← flip fix bytes $ \loop nextBytes →
       let (chunk, remaining) = splitAt 8 nextBytes
           m = take64Bits chunk
       in if null remaining then return chunk
                            else do v3 ^= m
                                    replicateM_ cRound sipRound
                                    v0 ^= m
                                    loop remaining

     let b = shiftL (fromIntegral (length bytes)) 56 + take64Bits end

     v3 ^= b
     replicateM_ cRound sipRound
     v0 ^= b
     v2 ^= 0xff
     replicateM_ dRound sipRound

     v0' ← readSTRef v0
     v1' ← readSTRef v1
     v2' ← readSTRef v2
     v3' ← readSTRef v3

     return (v0' `xor` v1' `xor` v2' `xor` v3')

  where take64Bits = fst . foldl' accum (0, 0) . take 8
          where accum (word, idx) byte =
                  (word .|. shift (fromIntegral byte) (idx * 8), idx + 1)

In order to pass the test vectors of reference implementation I had to change line 51 as follows:

in if length chunk < 8 then return chunk