{-# LANGUAGE FlexibleContexts, UndecidableInstances, TypeOperators, DataKinds, TypeFamilies, ScopedTypeVariables#-}

{- |

Generic implementation of the 'arbitrary' method. Example usage:

@
data Foo = Foo
  { _fooX :: X
  , _fooY :: Y
  } deriving (Generic)

instance Arbitrary Foo where
  arbitrary = genericArbitrary
  shrink = genericShrink
@

This instance can also be derived using DerivingVia language extension

@
data Foo = Foo
  { _fooX :: X
  , _fooY :: Y
  } deriving (Generic)
    deriving (Arbitrary) via GenericArbitrary Foo
@

The generated 'arbitrary' method is equivalent to

@Foo <$> arbitrary <*> arbitrary@.

-}

module Test.QuickCheck.Arbitrary.Generic
  ( GenericArbitrary(..)
  , Arbitrary(..)
  , genericArbitrary
  , genericShrink
  ) where

import Control.Applicative
import Data.Coerce (coerce)
import Data.Proxy
import GHC.Generics as G
import GHC.TypeLits
import Test.QuickCheck as QC
import Test.QuickCheck.Arbitrary (GSubterms, RecursivelyShrink)

newtype GenericArbitrary a = GenericArbitrary { GenericArbitrary a -> a
unGenericArbitrary :: a }
  deriving (Int -> GenericArbitrary a -> ShowS
[GenericArbitrary a] -> ShowS
GenericArbitrary a -> String
(Int -> GenericArbitrary a -> ShowS)
-> (GenericArbitrary a -> String)
-> ([GenericArbitrary a] -> ShowS)
-> Show (GenericArbitrary a)
forall a. Show a => Int -> GenericArbitrary a -> ShowS
forall a. Show a => [GenericArbitrary a] -> ShowS
forall a. Show a => GenericArbitrary a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [GenericArbitrary a] -> ShowS
$cshowList :: forall a. Show a => [GenericArbitrary a] -> ShowS
show :: GenericArbitrary a -> String
$cshow :: forall a. Show a => GenericArbitrary a -> String
showsPrec :: Int -> GenericArbitrary a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> GenericArbitrary a -> ShowS
Show, GenericArbitrary a -> GenericArbitrary a -> Bool
(GenericArbitrary a -> GenericArbitrary a -> Bool)
-> (GenericArbitrary a -> GenericArbitrary a -> Bool)
-> Eq (GenericArbitrary a)
forall a. Eq a => GenericArbitrary a -> GenericArbitrary a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: GenericArbitrary a -> GenericArbitrary a -> Bool
$c/= :: forall a. Eq a => GenericArbitrary a -> GenericArbitrary a -> Bool
== :: GenericArbitrary a -> GenericArbitrary a -> Bool
$c== :: forall a. Eq a => GenericArbitrary a -> GenericArbitrary a -> Bool
Eq)

instance
  ( Generic a,
    GArbitrary (Rep a),
    RecursivelyShrink (Rep a),
    GSubterms (Rep a) a
  ) => Arbitrary (GenericArbitrary a) where
  arbitrary :: Gen (GenericArbitrary a)
arbitrary = Gen a -> Gen (GenericArbitrary a)
coerce (Gen a
forall a (ga :: * -> *).
(Generic a, GArbitrary ga, ga ~ Rep a) =>
Gen a
genericArbitrary :: Gen a)
  shrink :: GenericArbitrary a -> [GenericArbitrary a]
shrink = (a -> [a]) -> GenericArbitrary a -> [GenericArbitrary a]
coerce (a -> [a]
forall a.
(Generic a, RecursivelyShrink (Rep a), GSubterms (Rep a) a) =>
a -> [a]
genericShrink :: a -> [a])

class GArbitrary a where
  gArbitrary :: QC.Gen (a x)

instance GArbitrary G.U1 where
  gArbitrary :: Gen (U1 x)
gArbitrary = U1 x -> Gen (U1 x)
forall (f :: * -> *) a. Applicative f => a -> f a
pure U1 x
forall k (p :: k). U1 p
G.U1

instance Arbitrary c => GArbitrary (G.K1 i c) where
  gArbitrary :: Gen (K1 i c x)
gArbitrary = c -> K1 i c x
forall k i c (p :: k). c -> K1 i c p
G.K1 (c -> K1 i c x) -> Gen c -> Gen (K1 i c x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> Int) -> Gen c -> Gen c
forall a. (Int -> Int) -> Gen a -> Gen a
scale Int -> Int
forall p. (Eq p, Num p, Enum p) => p -> p
predNat Gen c
forall a. Arbitrary a => Gen a
arbitrary
    where
      predNat :: p -> p
predNat p
0 = p
0
      predNat p
n = p -> p
forall a. Enum a => a -> a
pred p
n

instance GArbitrary f => GArbitrary (G.M1 i c f) where
  gArbitrary :: Gen (M1 i c f x)
gArbitrary = f x -> M1 i c f x
forall k i (c :: Meta) (f :: k -> *) (p :: k). f p -> M1 i c f p
G.M1 (f x -> M1 i c f x) -> Gen (f x) -> Gen (M1 i c f x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (f x)
forall (a :: * -> *) x. GArbitrary a => Gen (a x)
gArbitrary

instance (GArbitrary a, GArbitrary b) => GArbitrary (a G.:*: b) where
  gArbitrary :: Gen ((:*:) a b x)
gArbitrary = (a x -> b x -> (:*:) a b x)
-> Gen (a x) -> Gen (b x) -> Gen ((:*:) a b x)
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a x -> b x -> (:*:) a b x
forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
(G.:*:) Gen (a x)
forall (a :: * -> *) x. GArbitrary a => Gen (a x)
gArbitrary Gen (b x)
forall (a :: * -> *) x. GArbitrary a => Gen (a x)
gArbitrary

-- | Calculates count of constructors encoded by particular ':+:'.
-- Internal use only.
type family SumLen a :: Nat where
  SumLen (a G.:+: b) = (SumLen a) + (SumLen b)
  SumLen a           = 1

instance (GArbitrary a, GArbitrary b, KnownNat (SumLen a), KnownNat (SumLen b)
         ) => GArbitrary (a G.:+: b) where
  gArbitrary :: Gen ((:+:) a b x)
gArbitrary = [(Int, Gen ((:+:) a b x))] -> Gen ((:+:) a b x)
forall a. [(Int, Gen a)] -> Gen a
frequency
    [ (Int
lfreq, a x -> (:+:) a b x
forall k (f :: k -> *) (g :: k -> *) (p :: k). f p -> (:+:) f g p
G.L1 (a x -> (:+:) a b x) -> Gen (a x) -> Gen ((:+:) a b x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (a x)
forall (a :: * -> *) x. GArbitrary a => Gen (a x)
gArbitrary)
    , (Int
rfreq, b x -> (:+:) a b x
forall k (f :: k -> *) (g :: k -> *) (p :: k). g p -> (:+:) f g p
G.R1 (b x -> (:+:) a b x) -> Gen (b x) -> Gen ((:+:) a b x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (b x)
forall (a :: * -> *) x. GArbitrary a => Gen (a x)
gArbitrary) ]
    where
      lfreq :: Int
lfreq = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Proxy (SumLen a) -> Integer
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Integer
natVal (Proxy (SumLen a)
forall k (t :: k). Proxy t
Proxy :: Proxy (SumLen a))
      rfreq :: Int
rfreq = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Proxy (SumLen b) -> Integer
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Integer
natVal (Proxy (SumLen b)
forall k (t :: k). Proxy t
Proxy :: Proxy (SumLen b))

genericArbitrary :: (Generic a, GArbitrary ga, ga ~ G.Rep a) => Gen a
genericArbitrary :: Gen a
genericArbitrary = ga Any -> a
forall a x. Generic a => Rep a x -> a
G.to (ga Any -> a) -> Gen (ga Any) -> Gen a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (ga Any)
forall (a :: * -> *) x. GArbitrary a => Gen (a x)
gArbitrary