content: {-# LANGUAGE TupleSections #-}
  module Main where
  
  import Data.Binary.IEEE754 (doubleToWord, floatToWord, wordToDouble, wordToFloat)
  import Data.Bits ((.|.), (.&.), Bits, FiniteBits, bit)
  import Data.Discrimination (Sort, runSort)
  import Data.Discrimination.IEEE754 (sortingFloat, sortingDouble)
  import Data.Word (Word32, Word64)
  
  import Test.Framework (defaultMain, testGroup)
  import Test.Framework.Providers.API (Test)
  import Test.Framework.Providers.QuickCheck2 (testProperty)
  import Test.QuickCheck (Arbitrary)
  
  sortCompare :: Sort a -> a -> a -> Ordering
  sortCompare s l r = case runSort s [(l, LT), (r, GT)] of
      [x]:_ -> x
      _     -> EQ
  
  propAgreesWithOrd :: RealFloat a => Sort a -> (a, a) -> Bool
  propAgreesWithOrd s (l, r)
      | l < r = sortCompare s l r == LT
      | l > r = sortCompare s l r == GT
      | l == r = sortCompare s l r == EQ
      -- one or the other is NaN; Ord doesn't give coherent comparisons for these:
      -- compare NaN 0 = GT, but compare 0 NaN == GT as well.  So don't trust it.
      | otherwise = True
  
  propNegLTPos :: Num a => Sort a -> (a, a) -> Bool
  propNegLTPos s (l, r) = sortCompare s (negate (abs l)) (abs r) == LT
  
  propNegZeroLTZero :: Fractional a => Sort a -> Bool
  propNegZeroLTZero s = sortCompare s (-0.0) 0.0 == LT
  
  inf :: Fractional a => a
  inf = 1.0 / 0
  
  negInf :: Fractional a => a
  negInf = -1.0 / 0
  
  propNegInfLTNumbers :: RealFloat a => Sort a -> a -> Bool
  propNegInfLTNumbers s x = isNaN x || x == negInf || sortCompare s negInf x == LT
  
  propInfGTNumbers :: RealFloat a => Sort a -> a -> Bool
  propInfGTNumbers s x = isNaN x || x == inf || sortCompare s inf x == GT
  
  -- Trim a word to be a valid NaN payload of n bits.
  makeNaNPayload :: (Num w, Bits w) => Int -> w -> w
  makeNaNPayload n x = case x .&. (bit n - 1) of
    -- Zero significand is infinity, so choose 1 arbitrarily instead.
    0 -> 1
    payload -> payload
  
  makeFloatNaNPayload :: Word32 -> Word32
  makeFloatNaNPayload = makeNaNPayload 23
  
  -- Build a 'Float' NaN with the given sign and payload.
  makeFloatNaN :: Bool -> Word32 -> Float
  makeFloatNaN sign payload =
      wordToFloat $ exponentAndSign .|. makeFloatNaNPayload payload
    where
      exponentAndSign = if sign then 0xff800000 else 0x7f800000
  
  makeDoubleNaNPayload :: Word64 -> Word64
  makeDoubleNaNPayload = makeNaNPayload 52
  
  -- Build a 'Double' NaN with the given sign and payload.
  makeDoubleNaN :: Bool -> Word64 -> Double
  makeDoubleNaN sign payload =
      wordToDouble $ exponentAndSign .|. makeDoubleNaNPayload payload
    where
      exponentAndSign = if sign then 0xfff0000000000000 else 0x7ff0000000000000
  
  -- Let's obey the IEEE754 total order.  See
  -- https://github.com/rust-lang/rust/issues/5585 for a description of it.
  -- It's consistent with Ord for all non-NaN arguments, so we only have to write
  -- properties for the cases involving NaNs.
  nansTestGroup
      :: ( Arbitrary a, RealFloat a, Show a
         , Arbitrary b, FiniteBits b, Ord b, Show b
         )
      => Sort a -> (a -> b) -> (Bool -> b -> a) -> (b -> b) -> Test
  nansTestGroup s toWord makeNaN makePayload = testGroup "NaNs"
      [ testProperty "-NaN < everything" $ \x payload ->
          sortCompare s (makeNaN True payload) x == LT
      , testProperty "-NaN < -Infinity" $ \payload ->
          sortCompare s (makeNaN True payload) negInf == LT
      , testProperty "+NaN > everything" $ \x payload ->
          sortCompare s (makeNaN False payload) x == GT
      , testProperty "+NaN > +Infinity" $ \payload ->
          sortCompare s (makeNaN False payload) inf == GT
      , testProperty "-NaN < +NaN" $ \l r ->
          sortCompare s (makeNaN True l) (makeNaN False r) == LT &&
          sortCompare s (makeNaN False l) (makeNaN True r) == GT
      , testProperty "+NaNs compare by payload" $ \l r ->
          sortCompare s (makeNaN False l) (makeNaN False r) ==
          compare (makePayload l) (makePayload r)
      , testProperty "-NaNs compare by reverse payload" $ \l r ->
          sortCompare s (makeNaN True l) (makeNaN True r) ==
          compare (makePayload r) (makePayload l)
  
      -- Sanity checks...
      , testProperty "NaNs are NaN" $ \sign payload ->
          isNaN (makeNaN sign payload)
      , testProperty "makePayload is correct" $ \sign payload ->
          toWord (makeNaN sign payload) ==
          toWord (makeNaN sign (makePayload payload))
      ]
  
  comparisonsTestGroup
      :: ( RealFloat a, Arbitrary a, Show a
         , FiniteBits b, Ord b, Show b, Arbitrary b
         )
      => String -> Sort a -> (a -> b) -> (Bool -> b -> a) -> (b -> b) -> Test
  comparisonsTestGroup nm s toWord makeNaN makePayload = testGroup nm
      [ testProperty "negatives LT positives" $
          propNegLTPos s
      , testProperty "-0 < 0" $
          propNegZeroLTZero s
      , testProperty "-INF < x" $
          propNegInfLTNumbers s
      , testProperty "INF > x" $
          propInfGTNumbers s
  
      , testGroup "agrees with Ord" $ ($ propAgreesWithOrd s) <$>
          [ testProperty "l = 0" . (. (0,))
          , testProperty "r = 0" . (. (,0))
          , testProperty "l = -INF" . (. (negInf,))
          , testProperty "r = -INF" . (. (,negInf))
          , testProperty "l = INF" . (. (inf,))
          , testProperty "r = INF" . (. (,inf))
          , testProperty "l = r" . (. (\x -> (x, x)))
          , testProperty "arbitrary"
          ]
  
      , nansTestGroup s toWord makeNaN makePayload
      ]
  
  main :: IO ()
  main = defaultMain
      [ comparisonsTestGroup "Float" sortingFloat floatToWord
                                     makeFloatNaN makeFloatNaNPayload
      , comparisonsTestGroup "Double" sortingDouble doubleToWord
                                      makeDoubleNaN makeDoubleNaNPayload
      ]
  