{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeOperators #-} {-# OPTIONS_GHC -fno-warn-orphans #-} module QuickCheck.Arbitrary.Array where import QuickCheck.Arbitrary.Shape import Test.QuickCheck import System.Random ( Random ) import Data.Array.Accelerate.Array.Sugar ( Array(..), Segments, Shape, Elt, Z(..), (:.)(..), (!), DIM1, DIM2 ) import qualified Data.Array.Accelerate.Array.Sugar as Sugar import qualified Data.Set as Set instance (Elt e, Arbitrary e) => Arbitrary (Array Z e) where arbitrary = arbitraryArray Z shrink arr = [ Sugar.fromList Z [x] | x <- shrink (arr ! Z) ] instance (Elt e, Arbitrary e) => Arbitrary (Array DIM1 e) where arbitrary = arbitraryArray =<< arbitrary shrink arr = [ Sugar.fromList (Z :. length sl) sl | sl <- shrink (Sugar.toList arr) ] instance (Elt e, Arbitrary e) => Arbitrary (Array DIM2 e) where arbitrary = arbitraryArray =<< arbitrary shrink arr = let xs = Sugar.toList arr sh@(Z:.h:.w) = Sugar.shape arr in [ q | Z:.m:.n <- shrink sh, m > 0, n > 0 , q <- [ Sugar.fromFunction (Z :. m :. n) (\ix -> arr Sugar.! ix) -- top-left , Sugar.fromFunction (Z :. m :. w-n) (\(Z:.u:.v) -> arr Sugar.! (Z:.u :.(n+v))) -- top-right , Sugar.fromFunction (Z :. h-m :. n) (\(Z:.u:.v) -> arr Sugar.! (Z:.(m+u):.v)) -- bottom-left , Sugar.fromFunction (Z :. h-m :. w-n) (\(Z:.u:.v) -> arr Sugar.! (Z:.(m+u):.(n+v))) -- bottom-right ]] ++ [ Sugar.fromList sh sl | sl <- shrinkOne xs ] instance (Shape sh, Elt e, Arbitrary sh, Arbitrary e) => Arbitrary (Array (sh:.Int:.Int:.Int) e) where arbitrary = arbitraryArray =<< arbitrary shrink arr@(Array _ adata) = let sh = Sugar.shape arr xs = Sugar.toList arr in [ Array (Sugar.fromElt sh') adata | sh' <- shrink sh ] ++ [ Sugar.fromList sh sl | sl <- shrinkOne xs ] -- Generate an arbitrary array of the given shape using the default element -- generator -- arbitraryArray :: (Shape sh, Elt e, Arbitrary e) => sh -> Gen (Array sh e) arbitraryArray sh = arbitraryArrayOf sh arbitrary -- Generate an array of the given shape using the supplied element generator -- function. -- arbitraryArrayOf :: (Shape sh, Elt e) => sh -> Gen e -> Gen (Array sh e) arbitraryArrayOf sh gen = Sugar.fromList sh `fmap` vectorOf (Sugar.size sh) gen {-- -- A version that does not use fromList. It does not gain us anything while -- being much more complex in implementation. -- arbitraryArrayOf :: (Shape sh, Elt e, Arbitrary e) => sh -> Gen e -> Gen (Array sh e) arbitraryArrayOf sh (MkGen gen) = MkGen $ \g k -> let !n = Sugar.size sh (adata, _) = runArrayData $ do arr <- newArrayData n let go _ !i | i >= n = return () go !r !i = let (r1,r2) = split r v = gen r1 k in unsafeWriteArrayData arr i (Sugar.fromElt v) >> go r2 (i+1) -- go g 0 return (arr, undefined) in adata `seq` Sugar.Array (Sugar.fromElt sh) adata --} -- Generate an array where the outermost dimension satisfies the given segmented -- array descriptor. -- arbitrarySegmentedArray :: (Integral i, Shape sh, Elt e, Arbitrary sh, Arbitrary e) => Segments i -> Gen (Array (sh :. Int) e) arbitrarySegmentedArray segs = do let sz = fromIntegral . sum $ Sugar.toList segs sh <- sized $ \n -> arbitraryShape (n `div` 2) arbitraryArray (sh :. sz) shrinkSegmentedArray :: (Shape sh, Elt e, Arbitrary sh, Arbitrary e) => Array (sh :. Int) e -> [Array (sh :. Int) e] shrinkSegmentedArray arr@(Array _ adata) = let sh@(sz :. n) = Sugar.shape arr in [ Array (Sugar.fromElt sz',n) adata | sz' <- shrink sz ] ++ [ Sugar.fromList sh sl | sl <- shrinkOne (Sugar.toList arr) ] -- Generate a segment descriptor. Both the array and individual segments might -- be empty. -- arbitrarySegments :: (Elt i, Integral i, Random i) => Gen (Segments i) arbitrarySegments = sized $ \n -> do k <- choose (0,n) arbitraryArrayOf (Z:.k) (choose (0, fromIntegral n)) shrinkSegments :: (Elt i, Integral i, Arbitrary i) => Segments i -> [Segments i] shrinkSegments arr = let shrinkSeg (x:y:zs) = [ xy:zs | xy <- (x+y) : shrink (x+y), xy >= 0 ] ++ [ x:yzs' | yzs' <- shrinkSeg (y:zs) ] shrinkSeg _ = [] in [ Sugar.fromList (Z :. length sl) sl | sl <- shrinkSeg (Sugar.toList arr) ] ++ [ Sugar.fromList (Z :. length sl) sl | sl <- shrinkList (\x -> [ s | s <- shrink x, s >= 0]) (Sugar.toList arr) ] -- Generate a possibly empty segment descriptor, where each segment is non-empty -- arbitrarySegments1 :: (Elt i, Integral i, Random i) => Gen (Segments i) arbitrarySegments1 = sized $ \n -> do k <- choose (0,n) arbitraryArrayOf (Z:.k) (choose (1, 1 `max` fromIntegral n)) shrinkSegments1 :: (Elt i, Integral i, Arbitrary i) => Segments i -> [Segments i] shrinkSegments1 arr = let shrinkSeg (x:y:zs) = [ xy:zs | xy <- (x+y) : shrink (x+y), xy >= 1 ] ++ [ x:yzs' | yzs' <- shrinkSeg (y:zs) ] shrinkSeg _ = [] in [ Sugar.fromList (Z :. length sl) sl | sl <- shrinkSeg (Sugar.toList arr) ] ++ [ Sugar.fromList (Z :. length sl) sl | sl <- shrinkList (\x -> [ s | s <- shrink x, s >= 1]) (Sugar.toList arr) ] -- Generate an vector where every element in the array is unique. The maximum -- size is based on the current 'sized' parameter. -- arbitraryUniqueVectorOf :: (Elt e, Ord e) => Gen e -> Gen (Array DIM1 e) arbitraryUniqueVectorOf gen = sized $ \n -> do set <- fmap Set.fromList (vectorOf n gen) k <- choose (0, Set.size set) return $! Sugar.fromList (Z :. k) (Set.toList set) -- Generate an arbitrary CSR matrix. The first parameter is the segment -- descriptor, the second a sparse vector of (index,value) pairs, and the third -- the matrix width (number of columns). -- -- The matrix size is based on the current `sized` parameter. -- arbitraryCSRMatrix :: (Elt i, Integral i, Random i, Elt e, Arbitrary e) => Gen ( Array DIM1 i, Array DIM1 (i,e), Int ) arbitraryCSRMatrix = sized $ \cols -> do segd <- arbitrarySegments let nnz = fromIntegral . sum $ Sugar.toList segd smat <- arbitraryArrayOf (Z :. nnz) $ do val <- arbitrary ind <- choose (0, fromIntegral cols - 1) return (ind, val) return (segd, smat, cols) -- Helpers -- ------- shrinkOne :: Arbitrary a => [a] -> [[a]] shrinkOne [] = [] shrinkOne (x:xs) = [ x':xs | x' <- shrink x ] ++ [ x:xs' | xs' <- shrinkOne xs ]