{-| V2 implementation of ChibiHash 

    This is a 64-bit non-cryptographic hash function optimized for:
    - Fast performance on short strings
    - Good distribution of hash values
    - Simple implementation with no lookup tables

    Version 2 improvements over V1, from the original C implementation:

    - Faster performance on short strings (42 cycles/hash vs 34 cycles/hash)
    - Improved seeding that affects all 256 bits of internal state
    - Better mixing in bulk data processing
    - Passes all 252 tests in smhasher3 (commit 34093a3), v1 failed 3.

-}
module ChibiHash.V2
    ( chibihash64
    ) where

import Data.Word
import Data.Bits
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS

-- | Prime-like constant used for mixing, derived from digits of e
k :: Word64
k :: Word64
k = Word64
0x2B7E151628AED2A7

-- | Convert bytes to Word64 using little-endian ordering
-- Takes 8 bytes and combines them into a single 64-bit word
load64le :: [Word8] -> Word64
load64le :: [Word8] -> Word64
load64le [Word8]
bytes = 
    let lo :: Word64
lo = [Word8] -> Word64
load32le [Word8]
bytes
        hi :: Word64
hi = [Word8] -> Word64
load32le (Int -> [Word8] -> [Word8]
forall a. Int -> [a] -> [a]
drop Int
4 [Word8]
bytes)
    in Word64
lo Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Word64
hi Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int
32)

-- | Convert bytes to Word32 using little-endian ordering
-- Takes 4 bytes and combines them into the lower 32 bits of a Word64
load32le :: [Word8] -> Word64
load32le :: [Word8] -> Word64
load32le [Word8]
bytes = 
    let b0 :: Word64
b0 = Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Word8] -> Word8
forall a. HasCallStack => [a] -> a
head [Word8]
bytes)
        b1 :: Word64
b1 = Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Word8]
bytes [Word8] -> Int -> Word8
forall a. HasCallStack => [a] -> Int -> a
!! Int
1) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int
8
        b2 :: Word64
b2 = Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Word8]
bytes [Word8] -> Int -> Word8
forall a. HasCallStack => [a] -> Int -> a
!! Int
2) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int
16
        b3 :: Word64
b3 = Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Word8]
bytes [Word8] -> Int -> Word8
forall a. HasCallStack => [a] -> Int -> a
!! Int
3) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int
24
    in Word64
b0 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64
b1 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64
b2 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64
b3

-- | Basic arithmetic operations used throughout the hash function
add, subtract, mul :: Word64 -> Word64 -> Word64
add :: Word64 -> Word64 -> Word64
add Word64
a Word64
b = Word64
a Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
b
subtract :: Word64 -> Word64 -> Word64
subtract Word64
a Word64
b = Word64
a Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
b
mul :: Word64 -> Word64 -> Word64
mul Word64
a Word64
b = Word64
a Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
b

-- | Main hash function for V2
-- Takes a ByteString input and 64-bit seed value
-- Returns a 64-bit hash value
chibihash64 :: ByteString -> Word64 -> Word64
chibihash64 :: ByteString -> Word64 -> Word64
chibihash64 ByteString
input Word64
seed = 
    let bytes :: [Word8]
bytes = ByteString -> [Word8]
BS.unpack ByteString
input
        len :: Word64
len = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
BS.length ByteString
input) :: Word64

        -- Initialize state with seed-dependent values
        seed2 :: Word64
seed2 = ((Word64
seed Word64 -> Word64 -> Word64
`ChibiHash.V2.subtract` Word64
k) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`rotateL` Int
15) Word64 -> Word64 -> Word64
`add` 
                ((Word64
seed Word64 -> Word64 -> Word64
`ChibiHash.V2.subtract` Word64
k) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`rotateL` Int
47)

        h0 :: [Word64]
h0 = [ Word64
seed
             , Word64
seed Word64 -> Word64 -> Word64
`add` Word64
k
             , Word64
seed2
             , Word64
seed2 Word64 -> Word64 -> Word64
`add` ((Word64
k Word64 -> Word64 -> Word64
`mul` Word64
k) Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Word64
k)
             ]

        -- Process input in stages
        ([Word64]
h1, [Word8]
remaining) = [Word8] -> [Word64] -> ([Word64], [Word8])
processBlocks [Word8]
bytes [Word64]
h0  -- Process 32-byte blocks
        h2 :: [Word64]
h2 = [Word8] -> Int -> [Word64] -> [Word64]
processRemaining [Word8]
remaining ([Word8] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word8]
remaining) [Word64]
h1  -- Handle remaining bytes

    in case [Word64]
h2 of
        [Word64
ha, Word64
hb, Word64
hc, Word64
hd] -> 
            let -- Final mixing steps
                h_final_0 :: Word64
h_final_0 = Word64
ha Word64 -> Word64 -> Word64
`add` ((Word64
hc Word64 -> Word64 -> Word64
`mul` Word64
k) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`rotateL` Int
31 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64
hc Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
31))
                h_final_1 :: Word64
h_final_1 = Word64
hb Word64 -> Word64 -> Word64
`add` ((Word64
hd Word64 -> Word64 -> Word64
`mul` Word64
k) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`rotateL` Int
31 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64
hd Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
31))
                h_final_0' :: Word64
h_final_0' = Word64
h_final_0 Word64 -> Word64 -> Word64
`mul` Word64
k
                h_final_0'' :: Word64
h_final_0'' = Word64
h_final_0' Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64
h_final_0' Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
31)
                h_final_1' :: Word64
h_final_1' = Word64
h_final_1 Word64 -> Word64 -> Word64
`add` Word64
h_final_0''

                -- Length-dependent mixing
                x :: Word64
x = Word64
len Word64 -> Word64 -> Word64
`mul` Word64
k
                x' :: Word64
x' = Word64
x Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64
x Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`rotateL` Int
29)
                x'' :: Word64
x'' = Word64
x' Word64 -> Word64 -> Word64
`add` Word64
seed
                x''' :: Word64
x''' = Word64
x'' Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Word64
h_final_1'
                x'''' :: Word64
x'''' = Word64
x''' Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64
x''' Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`rotateL` Int
15) Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64
x''' Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`rotateL` Int
42)
                x''''' :: Word64
x''''' = Word64
x'''' Word64 -> Word64 -> Word64
`mul` Word64
k
                final :: Word64
final = Word64
x''''' Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64
x''''' Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`rotateL` Int
13) Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64
x''''' Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`rotateL` Int
31)
            in Word64
final
        [Word64]
_ -> [Char] -> Word64
forall a. HasCallStack => [Char] -> a
error [Char]
"Invalid state: processRemaining must return 4 values"

-- | Process input in 32-byte blocks
-- Each block is split into four 8-byte stripes
processBlocks :: [Word8] -> [Word64] -> ([Word64], [Word8])
processBlocks :: [Word8] -> [Word64] -> ([Word64], [Word8])
processBlocks [Word8]
input [Word64]
h
    | [Word8] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word8]
input Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
32 = ([Word64]
h, [Word8]
input)
    | Bool
otherwise =
        let ([Word8]
block, [Word8]
rest) = Int -> [Word8] -> ([Word8], [Word8])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
32 [Word8]
input
            stripes :: [[Word8]]
stripes = Int -> [Word8] -> [[Word8]]
forall a. Int -> [a] -> [[a]]
chunksOf Int
8 [Word8]
block
            h' :: [Word64]
h' = ([Word64] -> (Int, [Word8]) -> [Word64])
-> [Word64] -> [(Int, [Word8])] -> [Word64]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\[Word64]
acc (Int
i, [Word8]
s) -> ([Word64], Int, [Word8]) -> [Word64]
processStripe ([Word64]
acc, Int
i, [Word8]
s)) 
                      [Word64]
h 
                      ([Int] -> [[Word8]] -> [(Int, [Word8])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..Int
3] [[Word8]]
stripes)
        in [Word8] -> [Word64] -> ([Word64], [Word8])
processBlocks [Word8]
rest [Word64]
h'

-- | Process an 8-byte stripe within a block
-- Updates the 4-element state array based on stripe index
processStripe :: ([Word64], Int, [Word8]) -> [Word64]
processStripe :: ([Word64], Int, [Word8]) -> [Word64]
processStripe ([Word64]
state, Int
i, [Word8]
stripe) | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
4 =
    let v :: Word64
v = [Word8] -> Word64
load64le [Word8]
stripe
        hi' :: Word64
hi' = (Word64
v Word64 -> Word64 -> Word64
`add` ([Word64]
state [Word64] -> Int -> Word64
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)) Word64 -> Word64 -> Word64
`mul` Word64
k
        nextIdx :: Int
nextIdx = (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
3
        next :: Word64
next = ([Word64]
state [Word64] -> Int -> Word64
forall a. HasCallStack => [a] -> Int -> a
!! Int
nextIdx) Word64 -> Word64 -> Word64
`add` (Word64
v Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`rotateL` Int
27)
    in case Int
i of
        Int
0 -> [Word64
hi', Word64
next, [Word64]
state [Word64] -> Int -> Word64
forall a. HasCallStack => [a] -> Int -> a
!! Int
2, [Word64]
state [Word64] -> Int -> Word64
forall a. HasCallStack => [a] -> Int -> a
!! Int
3]
        Int
1 -> [[Word64] -> Word64
forall a. HasCallStack => [a] -> a
head [Word64]
state, Word64
hi', Word64
next, [Word64]
state [Word64] -> Int -> Word64
forall a. HasCallStack => [a] -> Int -> a
!! Int
3]
        Int
2 -> [[Word64] -> Word64
forall a. HasCallStack => [a] -> a
head [Word64]
state, [Word64]
state [Word64] -> Int -> Word64
forall a. HasCallStack => [a] -> Int -> a
!! Int
1, Word64
hi', Word64
next]
        Int
3 -> [Word64
next, [Word64]
state [Word64] -> Int -> Word64
forall a. HasCallStack => [a] -> Int -> a
!! Int
1, [Word64]
state [Word64] -> Int -> Word64
forall a. HasCallStack => [a] -> Int -> a
!! Int
2, Word64
hi']
        Int
_ -> [Char] -> [Word64]
forall a. HasCallStack => [Char] -> a
error [Char]
"Invalid index"
processStripe ([Word64], Int, [Word8])
_ = [Char] -> [Word64]
forall a. HasCallStack => [Char] -> a
error [Char]
"Invalid state: expected 4 hash values"

-- | Process remaining bytes after block processing
-- Handles different cases based on number of remaining bytes:
-- - 8 or more bytes: process in 8-byte chunks
-- - 4-7 bytes: special handling with two 32-bit reads
-- - 1-3 bytes: special handling for very short remainders
processRemaining :: [Word8] -> Int -> [Word64] -> [Word64]
processRemaining :: [Word8] -> Int -> [Word64] -> [Word64]
processRemaining [Word8]
bytes Int
len state :: [Word64]
state@[Word64
ha, Word64
hb, Word64
hc, Word64
hd]
    | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
8 = 
        let ([Word8]
chunk, [Word8]
rest) = Int -> [Word8] -> ([Word8], [Word8])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
8 [Word8]
bytes
            ha' :: Word64
ha' = Word64
ha Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` [Word8] -> Word64
load32le [Word8]
chunk
            ha'' :: Word64
ha'' = Word64
ha' Word64 -> Word64 -> Word64
`mul` Word64
k
            hb' :: Word64
hb' = Word64
hb Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` [Word8] -> Word64
load32le (Int -> [Word8] -> [Word8]
forall a. Int -> [a] -> [a]
drop Int
4 [Word8]
chunk)
            hb'' :: Word64
hb'' = Word64
hb' Word64 -> Word64 -> Word64
`mul` Word64
k
            remaining_len :: Int
remaining_len = Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
8
        in [Word8] -> Int -> [Word64] -> [Word64]
processRemaining [Word8]
rest Int
remaining_len [Word64
ha'', Word64
hb'', Word64
hc, Word64
hd]
    | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
4 = 
        let hc' :: Word64
hc' = Word64
hc Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` [Word8] -> Word64
load32le [Word8]
bytes
            hd' :: Word64
hd' = Word64
hd Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` [Word8] -> Word64
load32le (Int -> [Word8] -> [Word8]
forall a. Int -> [a] -> [a]
drop (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
4) [Word8]
bytes)
        in [Word64
ha, Word64
hb, Word64
hc', Word64
hd']
    | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = 
        let hc' :: Word64
hc' = Word64
hc Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Word8] -> Word8
forall a. HasCallStack => [a] -> a
head [Word8]
bytes)
            mid_byte :: Word64
mid_byte = Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Word8]
bytes [Word8] -> Int -> Word8
forall a. HasCallStack => [a] -> Int -> a
!! (Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2))
            last_byte :: Word64
last_byte = Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Word8] -> Word8
forall a. HasCallStack => [a] -> a
last [Word8]
bytes) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int
8
            hd' :: Word64
hd' = Word64
hd Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64
mid_byte Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64
last_byte)
        in [Word64
ha, Word64
hb, Word64
hc', Word64
hd']
    | Bool
otherwise = [Word64]
state
processRemaining [Word8]
_ Int
_ [Word64]
_ = [Char] -> [Word64]
forall a. HasCallStack => [Char] -> a
error [Char]
"Invalid state"

-- | Split a list into chunks of specified size
chunksOf :: Int -> [a] -> [[a]]
chunksOf :: forall a. Int -> [a] -> [[a]]
chunksOf Int
n = ([a] -> Bool) -> [[a]] -> [[a]]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Bool -> Bool
not (Bool -> Bool) -> ([a] -> Bool) -> [a] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null) ([[a]] -> [[a]]) -> ([a] -> [[a]]) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a] -> Maybe ([a], [a])) -> [a] -> [[a]]
forall {t} {a}. (t -> Maybe (a, t)) -> t -> [a]
unfoldr (([a], [a]) -> Maybe ([a], [a])
forall a. a -> Maybe a
Just (([a], [a]) -> Maybe ([a], [a]))
-> ([a] -> ([a], [a])) -> [a] -> Maybe ([a], [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n)
  where
    unfoldr :: (t -> Maybe (a, t)) -> t -> [a]
unfoldr t -> Maybe (a, t)
f t
x = case t -> Maybe (a, t)
f t
x of
                    Just (a
y, t
ys) -> a
y a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (t -> Maybe (a, t)) -> t -> [a]
unfoldr t -> Maybe (a, t)
f t
ys
                    Maybe (a, t)
Nothing -> []