SetOps.hs

module Tests.SetOps where

import Support ()

import Data.Bit
import Data.Bits
import Data.List.NonEmpty (NonEmpty(..))
import qualified Data.Vector.Unboxed as U
import Test.Tasty
import Test.Tasty.QuickCheck hiding ((.&.))

setOpTests :: TestTree
setOpTests = testGroup "Set operations"
    [ testProperty "union"          prop_union_def
    , testProperty "intersection"   prop_intersection_def
    , testProperty "difference"     prop_difference_def
    , testProperty "symDiff"        prop_symDiff_def

    , testProperty "unions"         prop_unions_def
    , testProperty "intersections"  prop_unions_def

    , testProperty "invert"         prop_invert_def

    , testProperty "select"         prop_select_def
    , testProperty "exclude"        prop_exclude_def

    , testProperty "selectBits"     prop_selectBits_def
    , testProperty "excludeBits"    prop_excludeBits_def

    , testProperty "countBits"      prop_countBits_def
    ]

union :: U.Vector Bit -> U.Vector Bit -> U.Vector Bit
union = zipBits (.|.)

prop_union_def :: U.Vector Bit -> U.Vector Bit -> Property
prop_union_def xs ys
    =  U.toList (union xs ys)
    === zipWith (.|.) (U.toList xs) (U.toList ys)

intersection :: U.Vector Bit -> U.Vector Bit -> U.Vector Bit
intersection = zipBits (.&.)

prop_intersection_def :: U.Vector Bit -> U.Vector Bit -> Property
prop_intersection_def xs ys
    =  U.toList (intersection xs ys)
    === zipWith (.&.) (U.toList xs) (U.toList ys)

difference :: U.Vector Bit -> U.Vector Bit -> U.Vector Bit
difference = zipBits (\a b -> a .&. complement b)

prop_difference_def :: U.Vector Bit -> U.Vector Bit -> Property
prop_difference_def xs ys
    =  U.toList (difference xs ys)
    === zipWith diff (U.toList xs) (U.toList ys)
    where
        diff x y = x .&. complement y

symDiff :: U.Vector Bit -> U.Vector Bit -> U.Vector Bit
symDiff = zipBits xor

prop_symDiff_def :: U.Vector Bit -> U.Vector Bit -> Property
prop_symDiff_def xs ys
    =  U.toList (symDiff xs ys)
    === zipWith xor (U.toList xs) (U.toList ys)

unions :: NonEmpty (U.Vector Bit) -> U.Vector Bit
unions (x :| xs) = U.slice 0 l $ U.modify (go xs) x
    where
        l = minimum $ fmap U.length (x :| xs)
        go [] _ = pure ()
        go (y : ys) acc = do
            zipInPlace (.|.) y acc
            go ys acc

prop_unions_def :: U.Vector Bit -> [U.Vector Bit] -> Property
prop_unions_def xs xss
    =   unions (xs :| xss)
    === foldr union xs xss

intersections :: NonEmpty (U.Vector Bit) -> U.Vector Bit
intersections (x :| xs) = U.slice 0 l $ U.modify (go xs) x
    where
        l = minimum $ fmap U.length (x :| xs)
        go [] _ = pure ()
        go (y : ys) acc = do
            zipInPlace (.&.) y acc
            go ys acc

prop_intersections_def :: U.Vector Bit -> [U.Vector Bit] -> Property
prop_intersections_def xs xss
    =   intersections (xs :| xss)
    === foldr intersection xs xss

prop_invert_def :: U.Vector Bit -> Bool
prop_invert_def xs
    =  U.toList (U.modify invertInPlace xs)
    == map complement (U.toList xs)

select :: U.Unbox a => U.Vector Bit -> U.Vector a -> [a]
select mask ws = U.toList (U.map snd (U.filter (unBit . fst) (U.zip mask ws)))

prop_select_def :: U.Vector Bit -> U.Vector Word -> Bool
prop_select_def xs ys
    =  select xs ys
    == [ x | (Bit True, x) <- zip (U.toList xs) (U.toList ys)]

exclude :: U.Unbox a => U.Vector Bit -> U.Vector a -> [a]
exclude mask ws = U.toList (U.map snd (U.filter (not . unBit . fst) (U.zip mask ws)))

prop_exclude_def :: U.Vector Bit -> U.Vector Word -> Bool
prop_exclude_def xs ys
    =  exclude xs ys
    == [ x | (Bit False, x) <- zip (U.toList xs) (U.toList ys)]

prop_selectBits_def :: U.Vector Bit -> U.Vector Bit -> Bool
prop_selectBits_def xs ys
    =  selectBits xs ys
    == U.fromList (select xs ys)

prop_excludeBits_def :: U.Vector Bit -> U.Vector Bit -> Bool
prop_excludeBits_def xs ys
    =  excludeBits xs ys
    == U.fromList (exclude xs ys)

prop_countBits_def :: U.Vector Bit -> Bool
prop_countBits_def xs
    =  countBits xs
    == U.length (selectBits xs xs)