{-# 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 (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 (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. 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 (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 (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 (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 (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 (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 :: Z -> Gen b -> Gen b
coarbitrary Z
_ = Gen b -> Gen b
forall a. a -> a
id 

instance (CoArbitrary a) 
       => CoArbitrary (a :. Int) where
  coarbitrary :: (a :. Int) -> Gen b -> Gen b
coarbitrary (a
a :. Int
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 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 :: Array r sh a -> Gen b -> Gen b
coarbitrary Array r sh a
arr 
        = (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 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 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 b -> ((b, b) -> prop) -> Property
forAll2 Gen b
arbf = Gen (b, b) -> ((b, b) -> prop) -> Property
forall a prop.
(Show a, Testable prop) =>
Gen a -> (a -> prop) -> Property
forAll (Gen (b, b) -> ((b, b) -> prop) -> Property)
-> Gen (b, b) -> ((b, b) -> prop) -> Property
forall a b. (a -> b) -> a -> b
$ (b -> b -> (b, b)) -> Gen b -> Gen b -> Gen (b, b)
forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (,)    Gen b
arbf Gen b
arbf
forAll3 :: Gen c -> ((c, c, c) -> prop) -> Property
forAll3 Gen c
arbf = Gen (c, c, c) -> ((c, c, c) -> prop) -> Property
forall a prop.
(Show a, Testable prop) =>
Gen a -> (a -> prop) -> Property
forAll (Gen (c, c, c) -> ((c, c, c) -> prop) -> Property)
-> Gen (c, c, c) -> ((c, c, c) -> prop) -> Property
forall a b. (a -> b) -> a -> b
$ (c -> c -> c -> (c, c, c))
-> Gen c -> Gen c -> Gen c -> Gen (c, c, c)
forall (m :: * -> *) a1 a2 a3 r.
Monad m =>
(a1 -> a2 -> a3 -> r) -> m a1 -> m a2 -> m a3 -> m r
liftM3 (,,)   Gen c
arbf Gen c
arbf Gen c
arbf
forAll4 :: Gen d -> ((d, d, d, d) -> prop) -> Property
forAll4 Gen d
arbf = Gen (d, d, d, d) -> ((d, d, d, d) -> prop) -> Property
forall a prop.
(Show a, Testable prop) =>
Gen a -> (a -> prop) -> Property
forAll (Gen (d, d, d, d) -> ((d, d, d, d) -> prop) -> Property)
-> Gen (d, d, d, d) -> ((d, d, d, d) -> prop) -> Property
forall a b. (a -> b) -> a -> b
$ (d -> d -> d -> d -> (d, d, d, d))
-> Gen d -> Gen d -> Gen d -> Gen d -> Gen (d, d, d, d)
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 d
arbf Gen d
arbf Gen d
arbf Gen d
arbf
forAll5 :: Gen e -> ((e, e, e, e, e) -> prop) -> Property
forAll5 Gen e
arbf = Gen (e, e, e, e, e) -> ((e, e, e, e, e) -> prop) -> Property
forall a prop.
(Show a, Testable prop) =>
Gen a -> (a -> prop) -> Property
forAll (Gen (e, e, e, e, e) -> ((e, e, e, e, e) -> prop) -> Property)
-> Gen (e, e, e, e, e) -> ((e, e, e, e, e) -> prop) -> Property
forall a b. (a -> b) -> a -> b
$ (e -> e -> e -> e -> e -> (e, e, e, e, e))
-> Gen e -> Gen e -> Gen e -> Gen e -> Gen e -> Gen (e, e, e, e, e)
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 e
arbf Gen e
arbf Gen e
arbf Gen e
arbf Gen e
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 b prop.
(Show b, Testable prop) =>
Gen b -> ((b, b) -> 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 c prop.
(Show c, Testable prop) =>
Gen c -> ((c, c, c) -> 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 d prop.
(Show d, Testable prop) =>
Gen d -> ((d, d, d, d) -> 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 e prop.
(Show e, Testable prop) =>
Gen e -> ((e, e, e, e, e) -> 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 b prop.
(Show b, Testable prop) =>
Gen b -> ((b, b) -> 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 c prop.
(Show c, Testable prop) =>
Gen c -> ((c, c, c) -> 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 d prop.
(Show d, Testable prop) =>
Gen d -> ((d, d, d, d) -> 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 e prop.
(Show e, Testable prop) =>
Gen e -> ((e, e, e, e, e) -> 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