content: {-# LANGUAGE CPP        #-}
  {-# LANGUAGE RankNTypes #-}
  
  #ifndef BITVEC_THREADSAFE
  module Tests.SetOps (setOpTests) where
  #else
  module Tests.SetOpsTS (setOpTests) where
  #endif
  
  import Support (twoTimesMore)
  
  import Control.Monad
  import Control.Monad.ST
  import Data.Bit
  import Data.Bits
  import qualified Data.Vector.Unboxed as U
  import qualified Data.Vector.Unboxed.Mutable as MU
  import Test.Tasty
  import Test.Tasty.QuickCheck hiding ((.&.))
  
  setOpTests :: TestTree
  setOpTests = testGroup "Set operations"
    [ testProperty "generalize1"              prop_generalize1
    , testProperty "generalize2"              prop_generalize2
    , twoTimesMore
    $ testProperty "zipBits"                  prop_zipBits
    , testProperty "zipInPlace"               prop_zipInPlace
    , testProperty "zipInPlace around"        prop_zipInPlace_around
  
    , testProperty "mapBits"                  prop_mapBits
    , testProperty "mapInPlace"               prop_mapInPlace
  
    , testProperty "union"                    prop_union_def
    , testProperty "intersection"             prop_intersection_def
    , testProperty "difference"               prop_difference_def
    , testProperty "symDiff"                  prop_symDiff_def
  
    , mkGroup "invertBits" prop_invertBits
  
    , testProperty "invertInPlace"            prop_invertInPlace
    , testProperty "invertInPlace around"     prop_invertInPlace_around
    , testProperty "invertInPlaceWords"       prop_invertInPlaceWords
    , testProperty "invertInPlace middle"     prop_invertInPlace_middle
    , testProperty "invertInPlaceLong middle" prop_invertInPlaceLong_middle
  
    , mkGroup "reverseBits" prop_reverseBits
  
    , testProperty "reverseInPlace"            prop_reverseInPlace
    , testProperty "reverseInPlaceWords"       prop_reverseInPlaceWords
    , testProperty "reverseInPlace middle"     prop_reverseInPlace_middle
    , testProperty "reverseInPlaceLong middle" prop_reverseInPlaceLong_middle
  
    , mkGroup2 "selectBits"  prop_selectBits_def
    , mkGroup2 "excludeBits" prop_excludeBits_def
  
    , mkGroup "countBits" prop_countBits_def
    ]
  
  mkGroup :: String -> (U.Vector Bit -> Property) -> TestTree
  mkGroup name prop = testGroup name
    [ testProperty "simple" prop
    , testProperty "simple_long" (prop . getLarge)
    , testProperty "middle" propMiddle
    , testProperty "middle_long" propMiddleLong
    ]
    where
      f m = let n = fromIntegral m :: Double in
        odd (truncate (exp (abs (sin n) * 10)) :: Integer)
      propMiddle (NonNegative from) (NonNegative len) (NonNegative excess) =
        prop (U.slice from len (U.generate (from + len + excess) (Bit . f)))
      propMiddleLong (NonNegative x) (NonNegative y) (NonNegative z) =
        propMiddle (NonNegative $ x * 31) (NonNegative $ y * 37) (NonNegative $ z * 29)
  
  mkGroup2 :: String -> (U.Vector Bit -> U.Vector Bit -> Property) -> TestTree
  mkGroup2 name prop = testGroup name
    [ testProperty "simple" prop
    , testProperty "simple_long" (\(Large xs) (Large ys) -> prop xs ys)
    , testProperty "middle" propMiddle
    , testProperty "middle_long" propMiddleLong
    ]
    where
      f m = let n = fromIntegral m :: Double in
        odd (truncate (exp (abs (sin n) * 10)) :: Integer)
      propMiddle (NonNegative from1) (NonNegative len1) (NonNegative excess1) (NonNegative from2) (NonNegative len2) (NonNegative excess2) =
        prop (U.slice from1 len1 (U.generate (from1 + len1 + excess1) (Bit . f))) (U.slice from2 len2 (U.generate (from2 + len2 + excess2) (Bit . f)))
      propMiddleLong (NonNegative x1) (NonNegative y1) (NonNegative z1) (NonNegative x2) (NonNegative y2) (NonNegative z2) =
        propMiddle (NonNegative $ x1 * 31) (NonNegative $ y1 * 37) (NonNegative $ z1 * 29) (NonNegative $ x2 * 31) (NonNegative $ y2 * 37) (NonNegative $ z2 * 29)
  
  prop_generalize1 :: Fun Bit Bit -> Bit -> Property
  prop_generalize1 fun x =
    applyFun fun x === generalize1 (applyFun fun) x
  
  prop_generalize2 :: Fun (Bit, Bit) Bit -> Bit -> Bit -> Property
  prop_generalize2 fun x y =
    curry (applyFun fun) x y === generalize2 (curry (applyFun fun)) x y
  
  prop_union_def :: U.Vector Bit -> U.Vector Bit -> Property
  prop_union_def xs ys =
    xs .|. ys === U.zipWith (.|.) xs ys
  
  prop_intersection_def :: U.Vector Bit -> U.Vector Bit -> Property
  prop_intersection_def xs ys =
    xs .&. ys === U.zipWith (.&.) xs ys
  
  prop_difference_def :: U.Vector Bit -> U.Vector Bit -> Property
  prop_difference_def xs ys =
    zipBits diff xs ys === U.zipWith diff xs ys
    where
      diff x y = x .&. complement y
  
  prop_symDiff_def :: U.Vector Bit -> U.Vector Bit -> Property
  prop_symDiff_def xs ys =
    xs `xor` ys === U.zipWith xor xs ys
  
  prop_zipBits :: Fun (Bit, Bit) Bit -> U.Vector Bit -> U.Vector Bit -> Property
  prop_zipBits fun xs ys =
    U.zipWith f xs ys === zipBits (generalize2 f) xs ys
    where
      f = curry $ applyFun fun
  
  prop_zipInPlace :: Fun (Bit, Bit) Bit -> U.Vector Bit -> U.Vector Bit -> Property
  prop_zipInPlace fun xs ys =
    U.zipWith f xs ys === U.take (min (U.length xs) (U.length ys)) (U.modify (zipInPlace (generalize2 f) xs) ys)
    where
      f = curry $ applyFun fun
  
  prop_zipInPlace_around :: Fun (Bit, Bit) Bit -> U.Vector Bit -> U.Vector Bit -> Property
  prop_zipInPlace_around fun ts xs = ioProperty $ do
    let l = U.length xs
        f = curry $ applyFun fun
    when (l < 2) discard
    ys <- U.thaw xs
    let zs = MU.slice 1 (l - 2) ys
    zipInPlace (generalize2 f) ts zs
    hd <- MU.read ys 0
    lst <- MU.read ys (MU.length ys - 1)
    pure $ hd === U.head xs .&&. lst === U.last xs
  
  prop_mapBits :: Fun Bit Bit -> U.Vector Bit -> Property
  prop_mapBits fun xs =
    U.map (applyFun fun) xs === mapBits (generalize1 (applyFun fun)) xs
  
  prop_mapInPlace :: Fun Bit Bit -> U.Vector Bit -> Property
  prop_mapInPlace fun xs =
    U.map (applyFun fun) xs === U.modify (mapInPlace (generalize1 (applyFun fun))) xs
  
  prop_invertBits :: U.Vector Bit -> Property
  prop_invertBits xs =
    U.map complement xs === complement xs
  
  prop_invertInPlace :: U.Vector Bit -> Property
  prop_invertInPlace xs =
    U.map complement xs === U.modify invertInPlace xs
  
  prop_invertInPlace_around :: U.Vector Bit -> Property
  prop_invertInPlace_around xs = ioProperty $ do
    let l = U.length xs
    when (l < 2) discard
    ys <- U.thaw xs
    let zs = MU.slice 1 (l - 2) ys
    invertInPlace zs
    hd <- MU.read ys 0
    lst <- MU.read ys (MU.length ys - 1)
    pure $ hd === U.head xs .&&. lst === U.last xs
  
  prop_invertInPlaceWords :: Large (U.Vector Bit) -> Property
  prop_invertInPlaceWords = prop_invertInPlace . getLarge
  
  prop_invertInPlace_middle :: NonNegative Int -> NonNegative Int -> NonNegative Int -> Property
  prop_invertInPlace_middle (NonNegative from) (NonNegative len) (NonNegative excess) = runST $ do
    let totalLen = from + len + excess
    vec <- MU.new totalLen
    forM_ [0 .. totalLen - 1] $ \i ->
      MU.write vec i (Bit (odd i))
    ref <- U.freeze vec
  
    let middle = MU.slice from len vec
    invertInPlace middle
    wec <- U.unsafeFreeze vec
  
    let refLeft   = U.take from ref
        wecLeft   = U.take from wec
        refRight  = U.drop (from + len) ref
        wecRight  = U.drop (from + len) wec
        refMiddle = U.map complement (U.take len (U.drop from ref))
        wecMiddle = U.take len (U.drop from wec)
    pure $ refLeft === wecLeft .&&. refRight === wecRight .&&. refMiddle === wecMiddle
  
  prop_invertInPlaceLong_middle :: NonNegative Int -> NonNegative Int -> NonNegative Int -> Property
  prop_invertInPlaceLong_middle (NonNegative x) (NonNegative y) (NonNegative z) =
    prop_invertInPlace_middle (NonNegative $ x * 31) (NonNegative $ y * 37) (NonNegative $ z * 29)
  
  prop_reverseBits :: U.Vector Bit -> Property
  prop_reverseBits xs =
    U.reverse xs === reverseBits xs
  
  prop_reverseInPlace :: U.Vector Bit -> Property
  prop_reverseInPlace xs =
    U.reverse xs === U.modify reverseInPlace xs
  
  prop_reverseInPlaceWords :: Large (U.Vector Bit) -> Property
  prop_reverseInPlaceWords = prop_reverseInPlace . getLarge
  
  prop_reverseInPlace_middle :: NonNegative Int -> NonNegative Int -> NonNegative Int -> Property
  prop_reverseInPlace_middle (NonNegative from) (NonNegative len) (NonNegative excess) = runST $ do
    let totalLen = from + len + excess
    vec <- MU.new totalLen
    forM_ [0 .. totalLen - 1] $ \i ->
      MU.write vec i (Bit (odd i))
    ref <- U.freeze vec
  
    let middle = MU.slice from len vec
    reverseInPlace middle
    wec <- U.unsafeFreeze vec
  
    let refLeft   = U.take from ref
        wecLeft   = U.take from wec
        refRight  = U.drop (from + len) ref
        wecRight  = U.drop (from + len) wec
        refMiddle = U.reverse (U.take len (U.drop from ref))
        wecMiddle = U.take len (U.drop from wec)
    pure $ refLeft === wecLeft .&&. refRight === wecRight .&&. refMiddle === wecMiddle
  
  prop_reverseInPlaceLong_middle :: NonNegative Int -> NonNegative Int -> NonNegative Int -> Property
  prop_reverseInPlaceLong_middle (NonNegative x) (NonNegative y) (NonNegative z) =
    prop_reverseInPlace_middle (NonNegative $ x * 31) (NonNegative $ y * 37) (NonNegative $ z * 29)
  
  select :: U.Unbox a => U.Vector Bit -> U.Vector a -> U.Vector a
  select mask ws = U.map snd (U.filter (unBit . fst) (U.zip mask ws))
  
  exclude :: U.Unbox a => U.Vector Bit -> U.Vector a -> U.Vector a
  exclude mask ws = U.map snd (U.filter (not . unBit . fst) (U.zip mask ws))
  
  prop_selectBits_def :: U.Vector Bit -> U.Vector Bit -> Property
  prop_selectBits_def xs ys = selectBits xs ys === select xs ys
  
  prop_excludeBits_def :: U.Vector Bit -> U.Vector Bit -> Property
  prop_excludeBits_def xs ys = excludeBits xs ys === exclude xs ys
  
  prop_countBits_def :: U.Vector Bit -> Property
  prop_countBits_def xs = countBits xs === U.length (selectBits xs xs)
  
  -------------------------------------------------------------------------------
  
  generalize1 :: (Bit -> Bit) -> (forall a. Bits a => a -> a)
  generalize1 f = case (f (Bit False), f (Bit True)) of
    (Bit False, Bit False) -> const zeroBits
    (Bit False, Bit True)  -> id
    (Bit True,  Bit False) -> complement
    (Bit True,  Bit True)  -> const $ complement zeroBits
  
  generalize2 :: (Bit -> Bit -> Bit) -> (forall a. Bits a => a -> a -> a)
  generalize2 f = case (f (Bit False) (Bit False), f (Bit False) (Bit True), f (Bit True) (Bit False), f (Bit True) (Bit True)) of
    (Bit False, Bit False, Bit False, Bit False) -> \_ _ -> zeroBits
    (Bit False, Bit False, Bit False, Bit True)  -> \x y -> x .&. y
    (Bit False, Bit False, Bit True,  Bit False) -> \x y -> x .&. complement y
    (Bit False, Bit False, Bit True,  Bit True)  -> \x _ -> x
  
    (Bit False, Bit True,  Bit False, Bit False) -> \x y -> complement x .&. y
    (Bit False, Bit True,  Bit False, Bit True)  -> \_ y -> y
    (Bit False, Bit True,  Bit True,  Bit False) -> \x y -> x `xor` y
    (Bit False, Bit True,  Bit True,  Bit True)  -> \x y -> x .|. y
  
    (Bit True,  Bit False, Bit False, Bit False) -> \x y -> complement (x .|. y)
    (Bit True,  Bit False, Bit False, Bit True)  -> \x y -> complement (x `xor` y)
    (Bit True,  Bit False, Bit True,  Bit False) -> \_ y -> complement y
    (Bit True,  Bit False, Bit True,  Bit True)  -> \x y -> x .|. complement y
  
    (Bit True,  Bit True,  Bit False, Bit False) -> \x _ -> complement x
    (Bit True,  Bit True,  Bit False, Bit True)  -> \x y -> complement x .|. y
    (Bit True,  Bit True,  Bit True,  Bit False) -> \x y -> complement (x .&. y)
    (Bit True,  Bit True,  Bit True,  Bit True)  -> \_ _ -> complement zeroBits
  