{-# LANGUAGE TypeOperators, FlexibleInstances, MultiParamTypeClasses, ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Data.Array.Repa.Arbitrary
        ( -- * Arbitrary Unboxed Arrays
          arbitraryUShaped
        , forAllUShaped
        , forAll2UShaped
        , forAll3UShaped
        , forAll4UShaped
        , forAll5UShaped

          -- * Arbitrary Boxed Arrays
        , arbitraryVShaped
        , forAllVShaped
        , forAll2VShaped
        , forAll3VShaped
        , forAll4VShaped
        , forAll5VShaped)
where
import Data.Array.Repa.Base
import Data.Array.Repa.Repr.Unboxed
import Data.Array.Repa.Shape
import Data.Array.Repa.Index
import Test.QuickCheck.Arbitrary
import Test.QuickCheck.Gen
import Test.QuickCheck.Property                 (forAll)
import Control.Monad
import qualified Data.Array.Repa.Repr.Vector    as V
import qualified Data.Vector.Unboxed            as U


-- Aribrary -------------------------------------------------------------------
-- | This module exports instances of @Arbitrary@ and @CoArbitrary@ for
--   unboxed Repa arrays.
instance Arbitrary Z where
  arbitrary :: Gen Z
arbitrary = Z -> Gen Z
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return Z
Z


-- Note: this is a shape that is "sized", and then random array for a given
-- shape is generated.
instance Arbitrary a 
      => Arbitrary (a :. Int) where
 arbitrary :: Gen (a :. Int)
arbitrary 
        = (Int -> Gen (a :. Int)) -> Gen (a :. Int)
forall a. (Int -> Gen a) -> Gen a
sized (\Int
n -> do 
                Int
b <- if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                         then Int -> Gen Int
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
1
                         else (Int, Int) -> Gen Int
forall a. Random a => (a, a) -> Gen a
choose (Int
1, Int
n)
                a
a <- Int -> Gen a -> Gen a
forall a. HasCallStack => Int -> Gen a -> Gen a
resize ((Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
b Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
b) Gen a
forall a. Arbitrary a => Gen a
arbitrary
                -- each dimension should be at least 1-wide
                (a :. Int) -> Gen (a :. Int)
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return ((a :. Int) -> Gen (a :. Int)) -> (a :. Int) -> Gen (a :. Int)
forall a b. (a -> b) -> a -> b
$ a
a a -> Int -> a :. Int
forall tail head. tail -> head -> tail :. head
:. Int
b)


-- | Generates a random unboxed array of a given shape
arbitraryUShaped :: sh -> Gen (Array U sh a)
arbitraryUShaped sh
sh =   sh -> [a] -> Array U sh a
forall sh a. (Shape sh, Unbox a) => sh -> [a] -> Array U sh a
fromListUnboxed sh
sh ([a] -> Array U sh a) -> Gen [a] -> Gen (Array U sh a)
forall a b. (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` Int -> Gen [a]
forall a. Arbitrary a => Int -> Gen [a]
vector (sh -> Int
forall sh. Shape sh => sh -> Int
size sh
sh)


-- | Generates a random boxed array of a given shape
arbitraryVShaped :: sh -> Gen (Array V sh a)
arbitraryVShaped sh
sh = sh -> [a] -> Array V sh a
forall sh a. Shape sh => sh -> [a] -> Array V sh a
V.fromListVector  sh
sh ([a] -> Array V sh a) -> Gen [a] -> Gen (Array V sh a)
forall a b. (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` Int -> Gen [a]
forall a. Arbitrary a => Int -> Gen [a]
vector (sh -> Int
forall sh. Shape sh => sh -> Int
size sh
sh)


instance (Arbitrary sh, Arbitrary a, U.Unbox a, Shape sh) 
       => Arbitrary (Array U sh a) where
  arbitrary :: Gen (Array U sh a)
arbitrary = Gen sh
forall a. Arbitrary a => Gen a
arbitrary Gen sh -> (sh -> Gen (Array U sh a)) -> Gen (Array U sh a)
forall a b. Gen a -> (a -> Gen b) -> Gen b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= sh -> Gen (Array U sh a)
forall {sh} {a}.
(Shape sh, Unbox a, Arbitrary a) =>
sh -> Gen (Array U sh a)
arbitraryUShaped


instance (Arbitrary sh, Arbitrary a, Shape sh) 
       => Arbitrary (Array V.V sh a) where
  arbitrary :: Gen (Array V sh a)
arbitrary = Gen sh
forall a. Arbitrary a => Gen a
arbitrary Gen sh -> (sh -> Gen (Array V sh a)) -> Gen (Array V sh a)
forall a b. Gen a -> (a -> Gen b) -> Gen b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= sh -> Gen (Array V sh a)
forall {sh} {a}.
(Shape sh, Arbitrary a) =>
sh -> Gen (Array V sh a)
arbitraryVShaped


-- CoArbitrary ----------------------------------------------------------------
instance CoArbitrary Z where
  coarbitrary :: forall b. Z -> Gen b -> Gen b
coarbitrary Z
_ = Gen b -> Gen b
forall a. a -> a
id 

instance (CoArbitrary a) 
       => CoArbitrary (a :. Int) where
  coarbitrary :: forall b. (a :. Int) -> Gen b -> Gen b
coarbitrary (a
a :. Int
b) = a -> Gen b -> Gen b
forall b. a -> Gen b -> Gen b
forall a b. CoArbitrary a => a -> Gen b -> Gen b
coarbitrary a
a (Gen b -> Gen b) -> (Gen b -> Gen b) -> Gen b -> Gen b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Gen b -> Gen b
forall b. Int -> Gen b -> Gen b
forall a b. CoArbitrary a => a -> Gen b -> Gen b
coarbitrary Int
b

instance (CoArbitrary sh, CoArbitrary a, Source r a, Shape sh) 
       => CoArbitrary (Array r sh a) where
  coarbitrary :: forall b. Array r sh a -> Gen b -> Gen b
coarbitrary Array r sh a
arr 
        = (sh -> Gen b -> Gen b
forall b. sh -> Gen b -> Gen b
forall a b. CoArbitrary a => a -> Gen b -> Gen b
coarbitrary (sh -> Gen b -> Gen b)
-> (Array r sh a -> sh) -> Array r sh a -> Gen b -> Gen b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r sh a -> sh
forall sh. Shape sh => Array r sh a -> sh
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent (Array r sh a -> Gen b -> Gen b) -> Array r sh a -> Gen b -> Gen b
forall a b. (a -> b) -> a -> b
$ Array r sh a
arr) (Gen b -> Gen b) -> (Gen b -> Gen b) -> Gen b -> Gen b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a] -> Gen b -> Gen b
forall b. [a] -> Gen b -> Gen b
forall a b. CoArbitrary a => a -> Gen b -> Gen b
coarbitrary ([a] -> Gen b -> Gen b)
-> (Array r sh a -> [a]) -> Array r sh a -> Gen b -> Gen b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r sh a -> [a]
forall sh r e. (Shape sh, Source r e) => Array r sh e -> [e]
toList (Array r sh a -> Gen b -> Gen b) -> Array r sh a -> Gen b -> Gen b
forall a b. (a -> b) -> a -> b
$ Array r sh a
arr)


-- Wrappers -------------------------------------------------------------------
-- | Convenience functions for writing tests on 2-,3-,4-tuples of arrays
--   of the same size (or just of a fixed size.)

-- | These are helper functions:
forAll2 :: Gen a2 -> ((a2, a2) -> prop) -> Property
forAll2 Gen a2
arbf = Gen (a2, a2) -> ((a2, a2) -> prop) -> Property
forall a prop.
(Show a, Testable prop) =>
Gen a -> (a -> prop) -> Property
forAll (Gen (a2, a2) -> ((a2, a2) -> prop) -> Property)
-> Gen (a2, a2) -> ((a2, a2) -> prop) -> Property
forall a b. (a -> b) -> a -> b
$ (a2 -> a2 -> (a2, a2)) -> Gen a2 -> Gen a2 -> Gen (a2, a2)
forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (,)    Gen a2
arbf Gen a2
arbf
forAll3 :: Gen a3 -> ((a3, a3, a3) -> prop) -> Property
forAll3 Gen a3
arbf = Gen (a3, a3, a3) -> ((a3, a3, a3) -> prop) -> Property
forall a prop.
(Show a, Testable prop) =>
Gen a -> (a -> prop) -> Property
forAll (Gen (a3, a3, a3) -> ((a3, a3, a3) -> prop) -> Property)
-> Gen (a3, a3, a3) -> ((a3, a3, a3) -> prop) -> Property
forall a b. (a -> b) -> a -> b
$ (a3 -> a3 -> a3 -> (a3, a3, a3))
-> Gen a3 -> Gen a3 -> Gen a3 -> Gen (a3, a3, a3)
forall (m :: * -> *) a1 a2 a3 r.
Monad m =>
(a1 -> a2 -> a3 -> r) -> m a1 -> m a2 -> m a3 -> m r
liftM3 (,,)   Gen a3
arbf Gen a3
arbf Gen a3
arbf
forAll4 :: Gen a4 -> ((a4, a4, a4, a4) -> prop) -> Property
forAll4 Gen a4
arbf = Gen (a4, a4, a4, a4) -> ((a4, a4, a4, a4) -> prop) -> Property
forall a prop.
(Show a, Testable prop) =>
Gen a -> (a -> prop) -> Property
forAll (Gen (a4, a4, a4, a4) -> ((a4, a4, a4, a4) -> prop) -> Property)
-> Gen (a4, a4, a4, a4) -> ((a4, a4, a4, a4) -> prop) -> Property
forall a b. (a -> b) -> a -> b
$ (a4 -> a4 -> a4 -> a4 -> (a4, a4, a4, a4))
-> Gen a4 -> Gen a4 -> Gen a4 -> Gen a4 -> Gen (a4, a4, a4, a4)
forall (m :: * -> *) a1 a2 a3 a4 r.
Monad m =>
(a1 -> a2 -> a3 -> a4 -> r) -> m a1 -> m a2 -> m a3 -> m a4 -> m r
liftM4 (,,,)  Gen a4
arbf Gen a4
arbf Gen a4
arbf Gen a4
arbf
forAll5 :: Gen a5 -> ((a5, a5, a5, a5, a5) -> prop) -> Property
forAll5 Gen a5
arbf = Gen (a5, a5, a5, a5, a5)
-> ((a5, a5, a5, a5, a5) -> prop) -> Property
forall a prop.
(Show a, Testable prop) =>
Gen a -> (a -> prop) -> Property
forAll (Gen (a5, a5, a5, a5, a5)
 -> ((a5, a5, a5, a5, a5) -> prop) -> Property)
-> Gen (a5, a5, a5, a5, a5)
-> ((a5, a5, a5, a5, a5) -> prop)
-> Property
forall a b. (a -> b) -> a -> b
$ (a5 -> a5 -> a5 -> a5 -> a5 -> (a5, a5, a5, a5, a5))
-> Gen a5
-> Gen a5
-> Gen a5
-> Gen a5
-> Gen a5
-> Gen (a5, a5, a5, a5, a5)
forall (m :: * -> *) a1 a2 a3 a4 a5 r.
Monad m =>
(a1 -> a2 -> a3 -> a4 -> a5 -> r)
-> m a1 -> m a2 -> m a3 -> m a4 -> m a5 -> m r
liftM5 (,,,,) Gen a5
arbf Gen a5
arbf Gen a5
arbf Gen a5
arbf Gen a5
arbf


-- | Property tested for unboxed random arrays with a given shape.
forAllUShaped :: sh -> (Array U sh a -> prop) -> Property
forAllUShaped sh
sh  = Gen (Array U sh a) -> (Array U sh a -> prop) -> Property
forall a prop.
(Show a, Testable prop) =>
Gen a -> (a -> prop) -> Property
forAll  (Gen (Array U sh a) -> (Array U sh a -> prop) -> Property)
-> Gen (Array U sh a) -> (Array U sh a -> prop) -> Property
forall a b. (a -> b) -> a -> b
$ sh -> Gen (Array U sh a)
forall {sh} {a}.
(Shape sh, Unbox a, Arbitrary a) =>
sh -> Gen (Array U sh a)
arbitraryUShaped sh
sh

-- | Property tested for pair of unboxed random arrays with a given shape.
forAll2UShaped :: sh -> ((Array U sh a, Array U sh a) -> prop) -> Property
forAll2UShaped sh
sh = Gen (Array U sh a)
-> ((Array U sh a, Array U sh a) -> prop) -> Property
forall {a2} {prop}.
(Show a2, Testable prop) =>
Gen a2 -> ((a2, a2) -> prop) -> Property
forAll2 (Gen (Array U sh a)
 -> ((Array U sh a, Array U sh a) -> prop) -> Property)
-> Gen (Array U sh a)
-> ((Array U sh a, Array U sh a) -> prop)
-> Property
forall a b. (a -> b) -> a -> b
$ sh -> Gen (Array U sh a)
forall {sh} {a}.
(Shape sh, Unbox a, Arbitrary a) =>
sh -> Gen (Array U sh a)
arbitraryUShaped sh
sh

-- | Property tested for triple of unboxed random arrays with a given shape.
forAll3UShaped :: sh
-> ((Array U sh a, Array U sh a, Array U sh a) -> prop) -> Property
forAll3UShaped sh
sh = Gen (Array U sh a)
-> ((Array U sh a, Array U sh a, Array U sh a) -> prop) -> Property
forall {a3} {prop}.
(Show a3, Testable prop) =>
Gen a3 -> ((a3, a3, a3) -> prop) -> Property
forAll3 (Gen (Array U sh a)
 -> ((Array U sh a, Array U sh a, Array U sh a) -> prop)
 -> Property)
-> Gen (Array U sh a)
-> ((Array U sh a, Array U sh a, Array U sh a) -> prop)
-> Property
forall a b. (a -> b) -> a -> b
$ sh -> Gen (Array U sh a)
forall {sh} {a}.
(Shape sh, Unbox a, Arbitrary a) =>
sh -> Gen (Array U sh a)
arbitraryUShaped sh
sh

-- | Property tested for quadruple of unboxed random arrays with a given shape.
forAll4UShaped :: sh
-> ((Array U sh a, Array U sh a, Array U sh a, Array U sh a)
    -> prop)
-> Property
forAll4UShaped sh
sh = Gen (Array U sh a)
-> ((Array U sh a, Array U sh a, Array U sh a, Array U sh a)
    -> prop)
-> Property
forall {a4} {prop}.
(Show a4, Testable prop) =>
Gen a4 -> ((a4, a4, a4, a4) -> prop) -> Property
forAll4 (Gen (Array U sh a)
 -> ((Array U sh a, Array U sh a, Array U sh a, Array U sh a)
     -> prop)
 -> Property)
-> Gen (Array U sh a)
-> ((Array U sh a, Array U sh a, Array U sh a, Array U sh a)
    -> prop)
-> Property
forall a b. (a -> b) -> a -> b
$ sh -> Gen (Array U sh a)
forall {sh} {a}.
(Shape sh, Unbox a, Arbitrary a) =>
sh -> Gen (Array U sh a)
arbitraryUShaped sh
sh

-- | Property tested for 5-tuple of unboxed random arrays with a given shape.
forAll5UShaped :: sh
-> ((Array U sh a, Array U sh a, Array U sh a, Array U sh a,
     Array U sh a)
    -> prop)
-> Property
forAll5UShaped sh
sh = Gen (Array U sh a)
-> ((Array U sh a, Array U sh a, Array U sh a, Array U sh a,
     Array U sh a)
    -> prop)
-> Property
forall {a5} {prop}.
(Show a5, Testable prop) =>
Gen a5 -> ((a5, a5, a5, a5, a5) -> prop) -> Property
forAll5 (Gen (Array U sh a)
 -> ((Array U sh a, Array U sh a, Array U sh a, Array U sh a,
      Array U sh a)
     -> prop)
 -> Property)
-> Gen (Array U sh a)
-> ((Array U sh a, Array U sh a, Array U sh a, Array U sh a,
     Array U sh a)
    -> prop)
-> Property
forall a b. (a -> b) -> a -> b
$ sh -> Gen (Array U sh a)
forall {sh} {a}.
(Shape sh, Unbox a, Arbitrary a) =>
sh -> Gen (Array U sh a)
arbitraryUShaped sh
sh


-- | Property tested for unboxed random arrays with a given shape.
forAllVShaped :: sh -> (Array V sh a -> prop) -> Property
forAllVShaped sh
sh  = Gen (Array V sh a) -> (Array V sh a -> prop) -> Property
forall a prop.
(Show a, Testable prop) =>
Gen a -> (a -> prop) -> Property
forAll  (Gen (Array V sh a) -> (Array V sh a -> prop) -> Property)
-> Gen (Array V sh a) -> (Array V sh a -> prop) -> Property
forall a b. (a -> b) -> a -> b
$ sh -> Gen (Array V sh a)
forall {sh} {a}.
(Shape sh, Arbitrary a) =>
sh -> Gen (Array V sh a)
arbitraryVShaped sh
sh

-- | Property tested for pair of unboxed random arrays with a given shape.
forAll2VShaped :: sh -> ((Array V sh a, Array V sh a) -> prop) -> Property
forAll2VShaped sh
sh = Gen (Array V sh a)
-> ((Array V sh a, Array V sh a) -> prop) -> Property
forall {a2} {prop}.
(Show a2, Testable prop) =>
Gen a2 -> ((a2, a2) -> prop) -> Property
forAll2 (Gen (Array V sh a)
 -> ((Array V sh a, Array V sh a) -> prop) -> Property)
-> Gen (Array V sh a)
-> ((Array V sh a, Array V sh a) -> prop)
-> Property
forall a b. (a -> b) -> a -> b
$ sh -> Gen (Array V sh a)
forall {sh} {a}.
(Shape sh, Arbitrary a) =>
sh -> Gen (Array V sh a)
arbitraryVShaped sh
sh

-- | Property tested for triple of unboxed random arrays with a given shape.
forAll3VShaped :: sh
-> ((Array V sh a, Array V sh a, Array V sh a) -> prop) -> Property
forAll3VShaped sh
sh = Gen (Array V sh a)
-> ((Array V sh a, Array V sh a, Array V sh a) -> prop) -> Property
forall {a3} {prop}.
(Show a3, Testable prop) =>
Gen a3 -> ((a3, a3, a3) -> prop) -> Property
forAll3 (Gen (Array V sh a)
 -> ((Array V sh a, Array V sh a, Array V sh a) -> prop)
 -> Property)
-> Gen (Array V sh a)
-> ((Array V sh a, Array V sh a, Array V sh a) -> prop)
-> Property
forall a b. (a -> b) -> a -> b
$ sh -> Gen (Array V sh a)
forall {sh} {a}.
(Shape sh, Arbitrary a) =>
sh -> Gen (Array V sh a)
arbitraryVShaped sh
sh

-- | Property tested for quadruple of unboxed random arrays with a given shape.
forAll4VShaped :: sh
-> ((Array V sh a, Array V sh a, Array V sh a, Array V sh a)
    -> prop)
-> Property
forAll4VShaped sh
sh = Gen (Array V sh a)
-> ((Array V sh a, Array V sh a, Array V sh a, Array V sh a)
    -> prop)
-> Property
forall {a4} {prop}.
(Show a4, Testable prop) =>
Gen a4 -> ((a4, a4, a4, a4) -> prop) -> Property
forAll4 (Gen (Array V sh a)
 -> ((Array V sh a, Array V sh a, Array V sh a, Array V sh a)
     -> prop)
 -> Property)
-> Gen (Array V sh a)
-> ((Array V sh a, Array V sh a, Array V sh a, Array V sh a)
    -> prop)
-> Property
forall a b. (a -> b) -> a -> b
$ sh -> Gen (Array V sh a)
forall {sh} {a}.
(Shape sh, Arbitrary a) =>
sh -> Gen (Array V sh a)
arbitraryVShaped sh
sh

-- | Property tested for 5-tuple of unboxed random arrays with a given shape.
forAll5VShaped :: sh
-> ((Array V sh a, Array V sh a, Array V sh a, Array V sh a,
     Array V sh a)
    -> prop)
-> Property
forAll5VShaped sh
sh = Gen (Array V sh a)
-> ((Array V sh a, Array V sh a, Array V sh a, Array V sh a,
     Array V sh a)
    -> prop)
-> Property
forall {a5} {prop}.
(Show a5, Testable prop) =>
Gen a5 -> ((a5, a5, a5, a5, a5) -> prop) -> Property
forAll5 (Gen (Array V sh a)
 -> ((Array V sh a, Array V sh a, Array V sh a, Array V sh a,
      Array V sh a)
     -> prop)
 -> Property)
-> Gen (Array V sh a)
-> ((Array V sh a, Array V sh a, Array V sh a, Array V sh a,
     Array V sh a)
    -> prop)
-> Property
forall a b. (a -> b) -> a -> b
$ sh -> Gen (Array V sh a)
forall {sh} {a}.
(Shape sh, Arbitrary a) =>
sh -> Gen (Array V sh a)
arbitraryVShaped sh
sh