{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE CPP #-}
module Foundation.Check.Arbitrary
    ( Arbitrary(..)
    , frequency
    , oneof
    , elements
    , between
    ) where

import           Basement.Imports
import           Foundation.Primitive
import           Basement.Nat
import           Basement.Cast (cast)
import           Basement.IntegralConv
import           Basement.Bounded
import           Basement.Types.OffsetSize
import qualified Basement.Types.Char7 as Char7
import           Basement.Types.Word128 (Word128(..))
import           Basement.Types.Word256 (Word256(..))
#if __GLASGOW_HASKELL__ >= 710
import qualified Basement.Sized.List as ListN
#endif
import           Foundation.Check.Gen
import           Foundation.Random
import           Foundation.Bits
import           Foundation.Collection
import           Foundation.Numerical
import           Control.Monad (replicateM)

-- | How to generate an arbitrary value for 'a'
class Arbitrary a where
    arbitrary :: Gen a

instance Arbitrary Integer where
    arbitrary :: Gen Integer
arbitrary = Gen Integer
arbitraryInteger
instance Arbitrary Natural where
    arbitrary :: Gen Natural
arbitrary = Gen Natural
arbitraryNatural

instance (NatWithinBound Word64 n, KnownNat n) => Arbitrary (Zn64 n) where
    arbitrary :: Gen (Zn64 n)
arbitrary = forall (n :: Natural).
(KnownNat n, NatWithinBound Word64 n) =>
Word64 -> Zn64 n
zn64 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
arbitrary
instance KnownNat n => Arbitrary (Zn n) where
    arbitrary :: Gen (Zn n)
arbitrary = forall (n :: Natural). KnownNat n => Natural -> Zn n
zn forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Natural
arbitraryNatural

-- prim types
instance Arbitrary Int where
    arbitrary :: Gen Int
arbitrary = Int64 -> Int
int64ToInt forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Int64
arbitraryInt64
instance Arbitrary Word where
    arbitrary :: Gen Word
arbitrary = Word64 -> Word
word64ToWord forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word64
arbitraryWord64
instance Arbitrary Word256 where
    arbitrary :: Gen Word256
arbitrary = Word64 -> Word64 -> Word64 -> Word64 -> Word256
Word256 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word64
arbitraryWord64 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen Word64
arbitraryWord64 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen Word64
arbitraryWord64 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen Word64
arbitraryWord64
instance Arbitrary Word128 where
    arbitrary :: Gen Word128
arbitrary = Word64 -> Word64 -> Word128
Word128 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word64
arbitraryWord64 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen Word64
arbitraryWord64
instance Arbitrary Word64 where
    arbitrary :: Gen Word64
arbitrary = Gen Word64
arbitraryWord64
instance Arbitrary Word32 where
    arbitrary :: Gen Word32
arbitrary = forall a b. IntegralDownsize a b => a -> b
integralDownsize forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word64
arbitraryWord64
instance Arbitrary Word16 where
    arbitrary :: Gen Word16
arbitrary = forall a b. IntegralDownsize a b => a -> b
integralDownsize forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word64
arbitraryWord64
instance Arbitrary Word8 where
    arbitrary :: Gen Word8
arbitrary = forall a b. IntegralDownsize a b => a -> b
integralDownsize forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word64
arbitraryWord64
instance Arbitrary Int64 where
    arbitrary :: Gen Int64
arbitrary = Gen Int64
arbitraryInt64
instance Arbitrary Int32 where
    arbitrary :: Gen Int32
arbitrary = forall a b. IntegralDownsize a b => a -> b
integralDownsize forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Int64
arbitraryInt64
instance Arbitrary Int16 where
    arbitrary :: Gen Int16
arbitrary = forall a b. IntegralDownsize a b => a -> b
integralDownsize forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Int64
arbitraryInt64
instance Arbitrary Int8 where
    arbitrary :: Gen Int8
arbitrary = forall a b. IntegralDownsize a b => a -> b
integralDownsize forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Int64
arbitraryInt64
instance Arbitrary Char where
    arbitrary :: Gen Char
arbitrary = Gen Char
arbitraryChar
instance Arbitrary Char7 where
    arbitrary :: Gen Char7
arbitrary = Word8 -> Char7
Char7.fromByteMask forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a b. IntegralDownsize a b => a -> b
integralDownsize forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word64
arbitraryWord64
instance Arbitrary (CountOf ty) where
    arbitrary :: Gen (CountOf ty)
arbitrary = forall ty. Int -> CountOf ty
CountOf forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
arbitrary

instance Arbitrary Bool where
    arbitrary :: Gen Bool
arbitrary = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Bits a => a -> Int -> Bool
testBit Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word64
arbitraryWord64

instance Arbitrary String where
    arbitrary :: Gen String
arbitrary = forall a. (GenParams -> Gen a) -> Gen a
genWithParams forall a b. (a -> b) -> a -> b
$ \GenParams
params ->
        forall l. IsList l => [Item l] -> l
fromList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Word -> Gen Word
genMax (GenParams -> Word
genMaxSizeString GenParams
params) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Word
i -> forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (forall source destination.
Cast source destination =>
source -> destination
cast Word
i) forall a. Arbitrary a => Gen a
arbitrary)

instance Arbitrary AsciiString where
    arbitrary :: Gen AsciiString
arbitrary = forall a. (GenParams -> Gen a) -> Gen a
genWithParams forall a b. (a -> b) -> a -> b
$ \GenParams
params ->
        forall l. IsList l => [Item l] -> l
fromList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Word -> Gen Word
genMax (GenParams -> Word
genMaxSizeString GenParams
params) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Word
i -> forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (forall source destination.
Cast source destination =>
source -> destination
cast Word
i) forall a. Arbitrary a => Gen a
arbitrary)

instance Arbitrary Float where
    arbitrary :: Gen Float
arbitrary = Gen Float
arbitraryF32
instance Arbitrary Double where
    arbitrary :: Gen Double
arbitrary = Gen Double
arbitraryF64

instance Arbitrary a => Arbitrary (Maybe a) where
    arbitrary :: Gen (Maybe a)
arbitrary = forall a. NonEmpty [(Word, Gen a)] -> Gen a
frequency forall a b. (a -> b) -> a -> b
$ forall c. Collection c => c -> NonEmpty c
nonEmpty_ [ (Word
1, forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing), (Word
4, forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
arbitrary) ]

instance (Arbitrary l, Arbitrary r) => Arbitrary (Either l r) where
    arbitrary :: Gen (Either l r)
arbitrary = forall a. NonEmpty [Gen a] -> Gen a
oneof forall a b. (a -> b) -> a -> b
$ forall c. Collection c => c -> NonEmpty c
nonEmpty_ [ forall a b. a -> Either a b
Left forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
arbitrary, forall a b. b -> Either a b
Right forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
arbitrary ]

instance (Arbitrary a, Arbitrary b)
    => Arbitrary (a,b) where
    arbitrary :: Gen (a, b)
arbitrary = (,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Arbitrary a => Gen a
arbitrary
instance (Arbitrary a, Arbitrary b, Arbitrary c)
    => Arbitrary (a,b,c) where
    arbitrary :: Gen (a, b, c)
arbitrary = (,,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Arbitrary a => Gen a
arbitrary
instance (Arbitrary a, Arbitrary b, Arbitrary c, Arbitrary d)
    => Arbitrary (a,b,c,d) where
    arbitrary :: Gen (a, b, c, d)
arbitrary = (,,,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Arbitrary a => Gen a
arbitrary
instance (Arbitrary a, Arbitrary b, Arbitrary c, Arbitrary d, Arbitrary e)
    => Arbitrary (a,b,c,d,e) where
    arbitrary :: Gen (a, b, c, d, e)
arbitrary = (,,,,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Arbitrary a => Gen a
arbitrary
instance (Arbitrary a, Arbitrary b, Arbitrary c, Arbitrary d, Arbitrary e, Arbitrary f)
    => Arbitrary (a,b,c,d,e,f) where
    arbitrary :: Gen (a, b, c, d, e, f)
arbitrary = (,,,,,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Arbitrary a => Gen a
arbitrary

instance Arbitrary a => Arbitrary [a] where
    arbitrary :: Gen [a]
arbitrary = forall a. (GenParams -> Gen a) -> Gen a
genWithParams forall a b. (a -> b) -> a -> b
$ \GenParams
params ->
        forall l. IsList l => [Item l] -> l
fromList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Word -> Gen Word
genMax (GenParams -> Word
genMaxSizeArray GenParams
params) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Word
i -> forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (forall source destination.
Cast source destination =>
source -> destination
cast Word
i) forall a. Arbitrary a => Gen a
arbitrary)
#if __GLASGOW_HASKELL__ >= 710
instance (Arbitrary a, KnownNat n, NatWithinBound Int n) => Arbitrary (ListN.ListN n a) where
    arbitrary :: Gen (ListN n a)
arbitrary = forall (n :: Natural) (m :: * -> *) a.
(NatWithinBound Int n, Monad m, KnownNat n) =>
m a -> m (ListN n a)
ListN.replicateM forall a. Arbitrary a => Gen a
arbitrary
#endif

arbitraryInteger :: Gen Integer
arbitraryInteger :: Gen Integer
arbitraryInteger =
    -- TODO use the sized parameter
    forall a. NonEmpty [(Word, Gen a)] -> Gen a
frequency forall a b. (a -> b) -> a -> b
$ forall c. Collection c => c -> NonEmpty c
nonEmpty_
        [ (Word
4, Bool -> Word -> Gen Integer
integerOfSize Bool
True Word
2)
        , (Word
4, Bool -> Word -> Gen Integer
integerOfSize Bool
False Word
2)
        , (Word
4, Bool -> Word -> Gen Integer
integerOfSize Bool
True Word
4)
        , (Word
4, Bool -> Word -> Gen Integer
integerOfSize Bool
False Word
4)
        , (Word
2, Bool -> Word -> Gen Integer
integerOfSize Bool
True Word
8)
        , (Word
2, Bool -> Word -> Gen Integer
integerOfSize Bool
False Word
8)
        , (Word
1, Bool -> Word -> Gen Integer
integerOfSize Bool
True Word
16)
        , (Word
1, Bool -> Word -> Gen Integer
integerOfSize Bool
False Word
16)
        ]
  where
    integerOfSize :: Bool -> Word -> Gen Integer
    integerOfSize :: Bool -> Word -> Gen Integer
integerOfSize Bool
toSign Word
n = ((if Bool
toSign then forall a. HasNegation a => a -> a
negate else forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id) forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall collection a.
Foldable collection =>
(a -> Element collection -> a) -> a -> collection -> a
foldl' (\Integer
x Element [Item (UArray Word8)]
y -> Integer
x forall a. Additive a => a -> a -> a
+ forall a b. IntegralUpsize a b => a -> b
integralUpsize Element [Item (UArray Word8)]
y) Integer
0 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall l. IsList l => l -> [Item l]
toList)
                         forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall ty. (PrimType ty, Arbitrary ty) => Word -> Gen (UArray ty)
arbitraryUArrayOf Word
n :: Gen (UArray Word8))

arbitraryNatural :: Gen Natural
arbitraryNatural :: Gen Natural
arbitraryNatural = forall a b. IntegralDownsize a b => a -> b
integralDownsize forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. Signed a => a -> a
abs forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Integer
arbitraryInteger

arbitraryChar :: Gen Char
arbitraryChar :: Gen Char
arbitraryChar = forall a. NonEmpty [(Word, Gen a)] -> Gen a
frequency forall a b. (a -> b) -> a -> b
$ forall c. Collection c => c -> NonEmpty c
nonEmpty_
    [ (Word
6, Word -> Char
wordToChar forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word -> Gen Word
genMax Word
128)
    , (Word
1, Word -> Char
wordToChar forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word -> Gen Word
genMax Word
0x10ffff)
    ]

arbitraryWord64 :: Gen Word64
arbitraryWord64 :: Gen Word64
arbitraryWord64 = forall a.
(forall (randomly :: * -> *). MonadRandom randomly => randomly a)
-> Gen a
genWithRng forall (m :: * -> *). MonadRandom m => m Word64
getRandomWord64

arbitraryInt64 :: Gen Int64
arbitraryInt64 :: Gen Int64
arbitraryInt64 = forall source destination.
Cast source destination =>
source -> destination
cast forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Word64
arbitraryWord64

arbitraryF64 :: Gen Double
arbitraryF64 :: Gen Double
arbitraryF64 = forall a.
(forall (randomly :: * -> *). MonadRandom randomly => randomly a)
-> Gen a
genWithRng forall (m :: * -> *). MonadRandom m => m Double
getRandomF64

arbitraryF32 :: Gen Float
arbitraryF32 :: Gen Float
arbitraryF32 = forall a.
(forall (randomly :: * -> *). MonadRandom randomly => randomly a)
-> Gen a
genWithRng forall (m :: * -> *). MonadRandom m => m Float
getRandomF32

arbitraryUArrayOf :: (PrimType ty, Arbitrary ty) => Word -> Gen (UArray ty)
arbitraryUArrayOf :: forall ty. (PrimType ty, Arbitrary ty) => Word -> Gen (UArray ty)
arbitraryUArrayOf Word
size = (Word, Word) -> Gen Word
between (Word
0, Word
size) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
    \Word
sz -> forall l. IsList l => [Item l] -> l
fromList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (forall source destination.
Cast source destination =>
source -> destination
cast Word
sz) forall a. Arbitrary a => Gen a
arbitrary

-- | Call one of the generator weighted
frequency :: NonEmpty [(Word, Gen a)] -> Gen a
frequency :: forall a. NonEmpty [(Word, Gen a)] -> Gen a
frequency (forall a. NonEmpty a -> a
getNonEmpty -> [(Word, Gen a)]
l) = (Word, Word) -> Gen Word
between (Word
0, Word
sum) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall {t} {b}.
(Difference t ~ t, Ord t, Subtractive t) =>
[(t, b)] -> t -> b
pickOne [(Word, Gen a)]
l
  where
    sum :: Word
    !sum :: Word
sum = forall collection a.
Foldable collection =>
(a -> Element collection -> a) -> a -> collection -> a
foldl' forall a. Additive a => a -> a -> a
(+) Word
0 forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst [(Word, Gen a)]
l

    pickOne :: [(t, b)] -> t -> b
pickOne ((t
k,b
x):[(t, b)]
xs) t
n
        | t
n forall a. Ord a => a -> a -> Bool
<= t
k    = b
x
        | Bool
otherwise = [(t, b)] -> t -> b
pickOne [(t, b)]
xs (t
nforall a. Subtractive a => a -> a -> Difference a
-t
k)
    pickOne [(t, b)]
_ t
_ = forall a. HasCallStack => String -> a
error String
"frequency"

oneof :: NonEmpty [Gen a] -> Gen a
oneof :: forall a. NonEmpty [Gen a] -> Gen a
oneof NonEmpty [Gen a]
ne = forall a. NonEmpty [(Word, Gen a)] -> Gen a
frequency (forall (f :: * -> *) a b.
Functor f =>
(a -> b) -> NonEmpty (f a) -> NonEmpty (f b)
nonEmptyFmap (\Gen a
x -> (Word
1, Gen a
x)) NonEmpty [Gen a]
ne)

elements :: NonEmpty [a] -> Gen a
elements :: forall a. NonEmpty [a] -> Gen a
elements NonEmpty [a]
l = forall a. NonEmpty [(Word, Gen a)] -> Gen a
frequency (forall (f :: * -> *) a b.
Functor f =>
(a -> b) -> NonEmpty (f a) -> NonEmpty (f b)
nonEmptyFmap (\a
x -> (Word
1, forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x)) NonEmpty [a]
l)

between :: (Word, Word) -> Gen Word
between :: (Word, Word) -> Gen Word
between (Word
x,Word
y)
    | Difference Word
range forall a. Eq a => a -> a -> Bool
== Word
0 = forall (f :: * -> *) a. Applicative f => a -> f a
pure Word
x
    | Bool
otherwise = forall a. Additive a => a -> a -> a
(+) Word
x forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word -> Gen Word
genMax Difference Word
range
  where range :: Difference Word
range = Word
y forall a. Subtractive a => a -> a -> Difference a
- Word
x

genMax :: Word -> Gen Word
genMax :: Word -> Gen Word
genMax Word
m = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. IDivisible a => a -> a -> a
mod Word
m forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
arbitrary