content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE DeriveTraversable #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeApplications #-}
  module Main (main) where
  
  import Control.Arrow (first)
  import Control.Monad ((<=<))
  import Control.Monad.Trans.State (State, runState, state)
  import Data.Hashable (Hashable)
  import Data.Coerce (coerce)
  import Data.Function (on)
  import Data.Functor.Compose (Compose (..))
  import Data.List (nub, nubBy)
  import Data.Maybe (fromMaybe)
  import Data.Proxy (Proxy (..))
  import Data.Typeable (Typeable, typeRep)
  import Test.QuickCheck (Arbitrary (..), Fun, Property, applyFun, Function (..), functionMap, CoArbitrary, (===))
  import Test.QuickCheck.Instances ()
  import Test.Tasty (defaultMain, testGroup, TestTree)
  import Test.Tasty.QuickCheck (testProperty)
  
  import qualified Data.HashMap.Lazy as HashMap
  import qualified Data.IntMap as IntMap
  import qualified Data.Map.Lazy as Map
  import qualified Data.Vector as V
  import qualified Data.Sequence as Seq
  
  import Witherable
  import Prelude hiding (filter)
  
  main :: IO ()
  main = defaultMain $ testGroup "witherable"
    [ testGroup "Filterable"
      [ filterableLaws (Proxy @[])
      , filterableLaws (Proxy @Maybe)
      , filterableLaws (Proxy @(Either String))
      , filterableLaws (Proxy @V.Vector)
      , filterableLaws (Proxy @Seq.Seq)
      , filterableLaws (Proxy @IntMap.IntMap)
      , filterableLaws (Proxy @(Map.Map K))
      , filterableLaws (Proxy @(HashMap.HashMap K))
      , filterableLaws (Proxy @Wicked)
      ]
  
    , testGroup "Witherable"
      [ witherableLaws (Proxy @[])
      , witherableLaws (Proxy @Maybe)
      , witherableLaws (Proxy @(Either String))
      , witherableLaws (Proxy @V.Vector)
      , witherableLaws (Proxy @Seq.Seq)
  #if MIN_VERSION_containers(0,6,3)
      -- traverse @IntMap is broken
      , witherableLaws (Proxy @IntMap.IntMap)
  #endif
      , witherableLaws (Proxy @(Map.Map K))
      , witherableLaws (Proxy @(HashMap.HashMap K))
      -- Wicked is not Witherable, see https://github.com/fumieval/witherable/issues/63#issuecomment-834631975
      -- , witherableLaws (Proxy @Wicked)
      ]
  
    , nubProperties
    ]
  
  -------------------------------------------------------------------------------
  -- Filterable laws
  -------------------------------------------------------------------------------
  
  filterableLaws
    :: forall f.
       ( Filterable f, Typeable f
       , Arbitrary (f A), Show (f A), Eq (f A)
       , Arbitrary (f (Maybe A)), Show (f (Maybe A))
       , Show (f B), Eq (f B), Show (f C), Eq (f C)
       )
    => Proxy f
    -> TestTree
  filterableLaws p = testGroup (show (typeRep p))
    [ testProperty "conservation" prop_conservation
    , testProperty "composition" prop_composition
    , testProperty "default filter" prop_default_filter
    , testProperty "default mapMaybe" prop_default_mapMaybe
    , testProperty "default catMaybes" prop_default_catMaybes
    ]
    where
      prop_conservation :: Fun A B -> f A -> Property
      prop_conservation f' xs =
          mapMaybe (Just . f) xs === fmap f xs
        where
          f = applyFun f'
  
      prop_composition :: Fun B (Maybe C) -> Fun A (Maybe B) -> f A -> Property
      prop_composition f' g' xs =
          mapMaybe f (mapMaybe g xs) === mapMaybe (f <=< g) xs
        where
          f = applyFun f'
          g = applyFun g'
  
      prop_default_filter :: Fun A Bool -> f A -> Property
      prop_default_filter f' xs =
          filter f xs === mapMaybe (\a -> if f a then Just a else Nothing) xs
        where
          f = applyFun f'
  
      prop_default_mapMaybe :: Fun A (Maybe B) -> f A -> Property
      prop_default_mapMaybe f' xs =
          mapMaybe f xs === catMaybes (fmap f xs)
        where
          f = applyFun f'
  
      prop_default_catMaybes :: f (Maybe A) -> Property
      prop_default_catMaybes xs = catMaybes xs === mapMaybe id xs
  
  -------------------------------------------------------------------------------
  -- Witherable laws
  -------------------------------------------------------------------------------
  
  witherableLaws
    :: forall f.
       ( Witherable f, Typeable f
       , Arbitrary (f A), Show (f A), Eq (f A)
       , Arbitrary (f (Maybe A)), Show (f (Maybe A))
       , Show (f B), Eq (f B), Show (f C), Eq (f C)
       )
    => Proxy f
    -> TestTree
  witherableLaws p = testGroup (show (typeRep p))
    [ testProperty "default wither" prop_default_wither
    , testProperty "default witherM" prop_default_witherM
    , testProperty "default filterA" prop_default_filterA
    , testProperty "identity" prop_identity
    , testProperty "composition" prop_composition
    ]
    where
      prop_default_wither :: S -> Fun (A, S) (Maybe B, S) -> f A -> Property
      prop_default_wither s0 f' xs = equalState s0 xs
          (wither f)
          (fmap catMaybes . traverse f)
        where
          f :: A -> State S (Maybe B)
          f a = state $ \s -> applyFun f' (a, s)
  
      prop_default_witherM :: S -> Fun (A, S) (Maybe B, S) -> f A -> Property
      prop_default_witherM s0 f' xs = equalState s0 xs
          (witherM f)
          (wither f)
        where
          f a = state $ \s -> applyFun f' (a, s)
  
      prop_default_filterA :: S -> Fun (A, S) (Bool, S) -> f A -> Property
      prop_default_filterA s0 f' xs = equalState s0 xs
          (filterA f)
          (wither (\a -> (\b -> if b then Just a else Nothing) <$> f a))
        where
          f a = state $ \s -> applyFun f' (a, s)
  
      prop_identity :: S -> Fun (A, S) (B, S) -> f A -> Property
      prop_identity s0 f' xs = equalState s0 xs
          (wither (fmap Just . f))
          (traverse f)
        where
          f a = state $ \s -> applyFun f' (a, s)
  
      prop_composition :: S -> S -> Fun (B, S) (Maybe C, S) -> Fun (A, S) (Maybe B, S) -> f A -> Property
      prop_composition s0 s1 f' g' xs = equalStateC s0 s1 xs
           (Compose . fmap (wither f) . wither g)
           (wither (Compose . fmap (wither f) . g))
        where
          f a = state $ \s -> applyFun f' (a, s)
          g b = state $ \s -> applyFun g' (b, s)
  
      equalState
          :: (Eq b, Show b)
          => S -> a -> (a -> State S b) -> (a -> State S b) -> Property
      equalState s0 xs f g = runState (f xs) s0 === runState (g xs) s0
  
      equalStateC
          :: forall a b. (Eq b, Show b)
          => S -> S -> a -> (a -> Compose (State S) (State S) b) -> (a -> Compose (State S) (State S) b) -> Property
      equalStateC s0 s1 xs f g = run (f xs) === run (g xs)
        where
          run :: Compose (State S) (State S) b -> ((b, S), S)
          run m = first (\x -> runState x s1) (runState (getCompose m) s0)
  
  -------------------------------------------------------------------------------
  -- Nub "laws"
  -------------------------------------------------------------------------------
  
  nubProperties :: TestTree
  nubProperties = testGroup "nub"
    [ testProperty "ordNub" prop_ordNub
    , testProperty "ordNubOn" prop_ordNubOn
    , testProperty "hashNub" prop_hashNub
    , testProperty "hashNubOn" prop_hashNubOn
    , testProperty "ordNub is lazy" prop_lazy_ordNub
    , testProperty "hashNub is lazy" prop_lazy_hashNub
    ]
    where
      prop_ordNub :: [A] -> Property
      prop_ordNub xs = nub xs === ordNub xs
  
      prop_hashNub :: [A] -> Property
      prop_hashNub xs = nub xs === hashNub xs
  
      prop_ordNubOn :: Fun A B -> [A] -> Property
      prop_ordNubOn f' xs = nubBy ((==) `on` f) xs === ordNubOn f xs
        where
          f = applyFun f'
  
      prop_hashNubOn :: Fun A B -> [A] -> Property
      prop_hashNubOn f' xs = nubBy ((==) `on` f) xs === hashNubOn f xs
        where
          f = applyFun f'
  
      prop_lazy_ordNub :: Property
      prop_lazy_ordNub = take 3 (ordNub ('x' : 'y' : 'z' : 'z' : error "bottom")) === "xyz"
  
      prop_lazy_hashNub :: Property
      prop_lazy_hashNub = take 3 (hashNub ('x' : 'y' : 'z' : 'z' : error "bottom")) === "xyz"
  
  -------------------------------------------------------------------------------
  -- "Poly"
  -------------------------------------------------------------------------------
  
  newtype A = A Int
    deriving (Eq, Ord, Show, Hashable, Arbitrary, CoArbitrary)
  
  instance Function A where
    function = functionMap coerce A
  
  newtype B = B Int
    deriving (Eq, Ord, Show, Hashable, Arbitrary, CoArbitrary)
  
  instance Function B where
    function = functionMap coerce B
  
  newtype C = C Int
    deriving (Eq, Ord, Show, Hashable, Arbitrary, CoArbitrary)
  
  instance Function C where
    function = functionMap coerce C
  
  newtype K = K Int
    deriving (Eq, Ord, Show, Hashable, Arbitrary, CoArbitrary)
  
  instance Function K where
    function = functionMap coerce K
  
  newtype S = S Int
    deriving (Eq, Ord, Show, Hashable, Arbitrary, CoArbitrary)
  
  instance Function S where
    function = functionMap coerce S
  
  -------------------------------------------------------------------------------
  -- Wicked
  -------------------------------------------------------------------------------
  
  newtype Wicked a = W [a]
    deriving (Eq, Show, Functor, Foldable, Traversable)
  
  instance Filterable Wicked where
    -- mapMaybe f (W [a1,a2,...]) = W [b1, b2, ...]
    -- if all of [f a1, f a2, ...] are Just. Otherwise, it returns (W []).
    mapMaybe f = fromMaybe (W []) . traverse f
  
  -- default implementation in terms of Filterable
  instance Witherable Wicked
  
  instance Arbitrary a => Arbitrary (Wicked a) where
      arbitrary = W <$> arbitrary
      shrink (W xs) = map W (shrink xs)
  