{-# LANGUAGE ScopedTypeVariables #-}
module Test.Data.WideWord.Word128
  ( testWord128
  ) where

import Control.Exception (evaluate)

import Data.Bits ((.&.), (.|.), bit, complement, countLeadingZeros, countTrailingZeros, popCount, rotateL, rotateR, shiftL, shiftR, testBit, xor)
import Data.Int (Int16)
import Data.Word (Word32, Word64)
import Data.WideWord

import Foreign (allocaBytes)
import Foreign.Storable (Storable (..))

import Test.Hspec (Spec, describe, errorCall, it, shouldBe, shouldThrow)
import Test.Hspec.QuickCheck (prop)
import Test.QuickCheck.Modifiers (NonZero (..))


testWord128 :: Spec
testWord128 = describe "Word128:" $ do
  prop "constructor and accessors" $ \ (h, l) ->
    let w128 = Word128 h l in
    (word128Hi64 w128, word128Lo64 w128) `shouldBe` (h, l)

  prop "byte swap" $ \ (h, l) ->
    let w128 = byteSwapWord128 $ byteSwapWord128 (Word128 h l) in
    (word128Hi64 w128, word128Lo64 w128) `shouldBe` (h, l)

  prop "derivied Eq instance" $ \ (a1, a0, b1, b0) ->
    (Word128 a1 a0 == Word128 b1 b0) `shouldBe` (a1 == b1 && a0 == b0)

  prop "Ord instance" $ \ (a1, a0, b1, b0) ->
    compare (Word128 a1 a0) (Word128 b1 b0) `shouldBe` compare (mkInteger a1 a0) (mkInteger b1 b0)

  prop "show" $ \ (a1, a0) ->
    show (Word128 a1 a0) `shouldBe` show (mkInteger a1 a0)

  prop "read" $ \ (a1, a0) ->
    read (show $ Word128 a1 a0) `shouldBe` Word128 a1 a0

  prop "succ" $ \ (a1, a0) ->
    if a1 == maxBound && a0 == maxBound
      then evaluate (succ $ Word128 a1 a0) `shouldThrow` errorCall "Enum.succ{Word128}: tried to take `succ' of maxBound"
      else toInteger128 (succ $ Word128 a1 a0) `shouldBe` succ (mkInteger a1 a0)

  prop "pred" $ \ (a1, a0) ->
    if a1 == 0 && a0 == 0
      then evaluate (pred $ Word128 a1 a0) `shouldThrow` errorCall "Enum.pred{Word128}: tried to take `pred' of minBound"
      else toInteger128 (pred $ Word128 a1 a0) `shouldBe` pred (mkInteger a1 a0)

  it "succ maxBound throws error" $
    evaluate (succ $ Word128 maxBound maxBound) `shouldThrow` errorCall "Enum.succ{Word128}: tried to take `succ' of maxBound"

  it "pred minBount throws error" $
    evaluate (pred $ Word128 0 0) `shouldThrow` errorCall "Enum.pred{Word128}: tried to take `pred' of minBound"

  prop "toEnum / fromEnum" $ \ (a0 :: Word32) -> do
    let w128 = Word128 0 (fromIntegral a0)
        e128 = fromEnum w128
    toInteger e128 `shouldBe` toInteger a0
    toInteger128 (toEnum e128 :: Word128) `shouldBe` toInteger a0

  prop "addition" $ \ (a1, a0, b1, b0) ->
    toInteger128 (Word128 a1 a0 + Word128 b1 b0) `shouldBe` correctWord128 (mkInteger a1 a0 + mkInteger b1 b0)

  prop "subtraction" $ \ (a1, a0, b1, b0) -> do
    let ai = mkInteger a1 a0
        bi = mkInteger b1 b0
        expected = ai + (1 `shiftL` 128) - bi
    toInteger128 (Word128 a1 a0 - Word128 b1 b0) `shouldBe` correctWord128 expected

  prop "multiplication" $ \ (a1, a0, b1, b0) ->
    toInteger128 (Word128 a1 a0 * Word128 b1 b0) `shouldBe` correctWord128 (mkInteger a1 a0 * mkInteger b1 b0)

  prop "negate" $ \ (a1, a0) ->
    toInteger128 (negate (Word128 a1 a0)) `shouldBe` correctWord128 (negate $ mkInteger a1 a0)

  prop "abs" $ \ (a1, a0) ->
    toInteger128 (abs (Word128 a1 a0)) `shouldBe` correctWord128 (abs $ mkInteger a1 a0)

  prop "signum" $ \ (a1, a0) ->
    toInteger128 (signum $ Word128 a1 a0) `shouldBe` signum (mkInteger a1 a0)

  prop "fromInteger" $ \ (a1, a0) -> do
    let w128 = fromInteger $ mkInteger a1 a0
    (word128Hi64 w128, word128Lo64 w128) `shouldBe` (a1, a0)

  prop "logical and/or/xor" $ \ (a1, a0, b1, b0) -> do
    toInteger128 (Word128 a1 a0 .&. Word128 b1 b0) `shouldBe` (mkInteger a1 a0 .&. mkInteger b1 b0)
    toInteger128 (Word128 a1 a0 .|. Word128 b1 b0) `shouldBe` (mkInteger a1 a0 .|. mkInteger b1 b0)
    toInteger128 (xor (Word128 a1 a0) (Word128 b1 b0)) `shouldBe` xor (mkInteger a1 a0) (mkInteger b1 b0)

  prop "complement" $ \ (a1, a0) ->
    toInteger128 (complement $ Word128 a1 a0) `shouldBe` mkInteger (complement a1) (complement a0)

  prop "logical shiftL" $ \ (a1, a0) shift ->
    let safeShift = if shift < 0 then 128 - (abs shift `mod` 128) else shift in
    toInteger128 (shiftL (Word128 a1 a0) shift) `shouldBe` correctWord128 (shiftL (mkInteger a1 a0) safeShift)

  prop "logical shiftR" $ \ (a1, a0) shift ->
    let expected = if shift < 0 then 0 else correctWord128 (shiftR (mkInteger a1 a0) shift) in
    toInteger128 (shiftR (Word128 a1 a0) shift) `shouldBe` expected

  -- Use `Int16` here to force a uniform distribution across the `Int16` range
  -- (standard QuickCkeck generator for `Int` doesn't give an even distribution).
  prop "logical rotateL" $ \ (a1, a0) (r :: Int16) -> do
    let rot = fromIntegral r
        i128 = mkInteger a1 a0
        expected
          | rot < 0 = 0
          | otherwise =
              correctWord128 (i128 `shiftL` erot + i128 `shiftR` (128 - (erot `mod` 128)))
              where
                erot
                  | rot < 0 = 128 - (abs rot `mod` 128)
                  | otherwise = rot `mod` 128
    toInteger128 (rotateL (Word128 a1 a0) rot) `shouldBe` expected

  prop "logical rotateR" $ \ (a1, a0) (r :: Int16) -> do
    let rot = fromIntegral r
        i128 = mkInteger a1 a0
        expected =
          correctWord128 $ i128 `shiftR` erot + i128 `shiftL` (128 - erot)
          where
            erot
              | rot < 0 = 128 - (abs rot `mod` 128)
              | otherwise = rot `mod` 128
    toInteger128 (rotateR (Word128 a1 a0) rot) `shouldBe` expected

  prop "testBit" $ \ (a1, a0) (b :: Int16) -> do
    let idx = fromIntegral b
        expected
          | idx < 0 = False
          | idx >= 128 = False
          | otherwise = testBit (mkInteger a1 a0) idx
    testBit (Word128 a1 a0) idx `shouldBe` expected

  prop "bit" $ \ (b :: Int16) -> do
    let idx = fromIntegral b
        expected
          | idx < 0 = 0
          | idx >= 128 = 0
          | otherwise = bit idx
    toInteger128 (bit idx :: Word128) `shouldBe` expected

  prop "popCount" $ \ (a1, a0) ->
    popCount (Word128 a1 a0) `shouldBe` popCount (mkInteger a1 a0)

  prop "countLeadingZeros" $ \ (a1, a0) -> do
    let expected = if a1 == 0
                    then 64 + countLeadingZeros a0
                    else countLeadingZeros a1
    countLeadingZeros (Word128 a1 a0) `shouldBe` expected

  prop "countTrailingZeros" $ \ (a1, a0) -> do
    let expected = if a0 == 0
                    then 64 + countTrailingZeros a1
                    else countTrailingZeros a0
    countTrailingZeros (Word128 a1 a0) `shouldBe` expected

  prop "quotRem (both upper words zero)" $ \ (a0, NonZero b0) -> do
    let (aq128, ar128) = quotRem (Word128 0 a0) (Word128 0 b0)
    (toInteger128 aq128, toInteger128 ar128) `shouldBe` quotRem (mkInteger 0 a0) (mkInteger 0 b0)

  prop "quotRem (denominator upper word zero)" $ \ (NonZero a1, a0, NonZero b0) -> do
    let (aq128, ar128) = quotRem (Word128 a1 a0) (Word128 0 b0)
    (toInteger128 aq128, toInteger128 ar128) `shouldBe` quotRem (mkInteger a1 a0) (mkInteger 0 b0)

  -- Don't need to test `quot` or `rem` because they are implemented by applying
  -- `fst` or `snd` to the output of `quotRem`.
  prop "quotRem (full)" $ \ (a1, a0, NonZero b1, b0) -> do
    let (aq128, ar128) = quotRem (Word128 a1 a0) (Word128 b1 b0)
    (toInteger128 aq128, toInteger128 ar128) `shouldBe` quotRem (mkInteger a1 a0) (mkInteger b1 b0)

  -- For unsigned values `quotRem` and `divMod` should give the same results.
  prop "divMod (full)" $ \ (a1, a0, NonZero b1, b0) -> do
    let (aq128, ar128) = divMod (Word128 a1 a0) (Word128 b1 b0)
    (toInteger128 aq128, toInteger128 ar128) `shouldBe` divMod (mkInteger a1 a0) (mkInteger b1 b0)

  prop "peek / poke" $ \ (a1, a0) -> do
    ar <- allocaBytes (sizeOf zeroWord128) $ \ ptr -> do
                    poke ptr $ Word128 a1 a0
                    peek ptr
    toInteger128 ar `shouldBe` mkInteger a1 a0

  prop "peekElemOff / pokeElemOff" $ \ (a1, a0, b1, b0) -> do
    (ar, br) <- allocaBytes (2 * sizeOf zeroWord128) $ \ ptr -> do
                    pokeElemOff ptr 0 $ Word128 a1 a0
                    pokeElemOff ptr 1 $ Word128 b1 b0
                    (,) <$> peekElemOff ptr 0 <*>  peekElemOff ptr 1
    (toInteger128 ar, toInteger128 br) `shouldBe` (mkInteger a1 a0, mkInteger b1 b0)

-- -----------------------------------------------------------------------------

mkInteger :: Word64 -> Word64 -> Integer
mkInteger a1 a0 = fromIntegral a1 `shiftL` 64 + fromIntegral a0

correctWord128 :: Integer -> Integer
correctWord128 i
  | i >= 0 && i <= maxWord128 = i
  | otherwise = i .&. maxWord128
  where
    maxWord128 = (1 `shiftL` 128) - 1

toInteger128 :: Word128 -> Integer
toInteger128 = toInteger