{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE UndecidableInstances #-}
module Test.Massiv.Core.Common
  ( ArrNE(..)
  , ArrTiny(..)
  , ArrTinyNE(..)
  , ArrIx(..)
  , ArrDW(..)
  , module X
  ) where

import Data.Massiv.Array
import Test.Massiv.Core.Index as X
import Test.Massiv.Utils




-- | Arbitrary non-empty array. Computation strategy can be either `Seq` or `Par`.
--
-- @since 0.1.0
newtype ArrNE r ix e = ArrNE
  { ArrNE r ix e -> Array r ix e
unArr :: Array r ix e
  }

-- | Arbitrary small and possibly empty array. Computation strategy can be either `Seq` or `Par`.
--
-- @since 0.1.0
newtype ArrTiny r ix e = ArrTiny
  { ArrTiny r ix e -> Array r ix e
unArrTiny :: Array r ix e
  }

-- | Tiny but non-empty
--
-- @since 0.1.0
newtype ArrTinyNE r ix e = ArrTinyNE
  { ArrTinyNE r ix e -> Array r ix e
unArrTinyNE :: Array r ix e
  }

-- | Arbitrary non-empty array with a valid index. Can be either `Seq` or `Par`
--
-- @since 0.1.0
data ArrIx r ix e = ArrIx (Array r ix e) ix

deriving instance (Show (Array r ix e)) => Show (ArrNE r ix e)
deriving instance (Show (Array r ix e)) => Show (ArrTiny r ix e)
deriving instance (Show (Array r ix e)) => Show (ArrTinyNE r ix e)
deriving instance (Show (Array r ix e), Show ix) => Show (ArrIx r ix e)



instance Arbitrary Comp where
  arbitrary :: Gen Comp
arbitrary =
    [(Int, Gen Comp)] -> Gen Comp
forall a. [(Int, Gen a)] -> Gen a
frequency
      [ (Int
20, Comp -> Gen Comp
forall (f :: * -> *) a. Applicative f => a -> f a
pure Comp
Seq)
      , (Int
10, Comp -> Gen Comp
forall (f :: * -> *) a. Applicative f => a -> f a
pure Comp
Par)
      , (Int
15, [Int] -> Comp
ParOn ([Int] -> Comp) -> Gen [Int] -> Gen Comp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen [Int]
forall a. Arbitrary a => Gen a
arbitrary)
      , (Int
15, Word16 -> Comp
ParN (Word16 -> Comp)
-> (Small Word16 -> Word16) -> Small Word16 -> Comp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Small Word16 -> Word16
forall a. Small a -> a
getSmall (Small Word16 -> Comp) -> Gen (Small Word16) -> Gen Comp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (Small Word16)
forall a. Arbitrary a => Gen a
arbitrary)
      ]


arbitraryArray :: (Load r ix e, Arbitrary e) => Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray :: Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray Gen (Sz ix)
szGen = Comp -> Sz ix -> (Int -> e) -> Array r ix e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear (Comp -> Sz ix -> (Int -> e) -> Array r ix e)
-> Gen Comp -> Gen (Sz ix -> (Int -> e) -> Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Comp
forall a. Arbitrary a => Gen a
arbitrary Gen (Sz ix -> (Int -> e) -> Array r ix e)
-> Gen (Sz ix) -> Gen ((Int -> e) -> Array r ix e)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen (Sz ix)
szGen Gen ((Int -> e) -> Array r ix e)
-> Gen (Int -> e) -> Gen (Array r ix e)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen (Int -> e)
forall a. Arbitrary a => Gen a
arbitrary

-- | Arbitrary array
instance (Arbitrary ix, Load r ix e, Arbitrary e) =>
         Arbitrary (Array r ix e) where
  arbitrary :: Gen (Array r ix e)
arbitrary = Comp -> Sz ix -> (Int -> e) -> Array r ix e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear (Comp -> Sz ix -> (Int -> e) -> Array r ix e)
-> Gen Comp -> Gen (Sz ix -> (Int -> e) -> Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Comp
forall a. Arbitrary a => Gen a
arbitrary Gen (Sz ix -> (Int -> e) -> Array r ix e)
-> Gen (Sz ix) -> Gen ((Int -> e) -> Array r ix e)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen (Sz ix)
forall a. Arbitrary a => Gen a
arbitrary Gen ((Int -> e) -> Array r ix e)
-> Gen (Int -> e) -> Gen (Array r ix e)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen (Int -> e)
forall a. Arbitrary a => Gen a
arbitrary


instance (Arbitrary ix, Load r ix e, Arbitrary e) => Arbitrary (ArrTiny r ix e) where
  arbitrary :: Gen (ArrTiny r ix e)
arbitrary = Array r ix e -> ArrTiny r ix e
forall r ix e. Array r ix e -> ArrTiny r ix e
ArrTiny (Array r ix e -> ArrTiny r ix e)
-> Gen (Array r ix e) -> Gen (ArrTiny r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (Sz ix) -> Gen (Array r ix e)
forall r ix e.
(Load r ix e, Arbitrary e) =>
Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray ((Int -> Int) -> Sz ix -> Sz ix
forall ix. Index ix => (Int -> Int) -> Sz ix -> Sz ix
liftSz (Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
10) (Sz ix -> Sz ix) -> Gen (Sz ix) -> Gen (Sz ix)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (Sz ix)
forall a. Arbitrary a => Gen a
arbitrary)

-- | Arbitrary small and possibly empty array. Computation strategy can be either `Seq` or `Par`.
instance (Arbitrary ix, Load r ix e, Arbitrary e) =>
         Arbitrary (ArrTinyNE r ix e) where
  arbitrary :: Gen (ArrTinyNE r ix e)
arbitrary = Array r ix e -> ArrTinyNE r ix e
forall r ix e. Array r ix e -> ArrTinyNE r ix e
ArrTinyNE (Array r ix e -> ArrTinyNE r ix e)
-> Gen (Array r ix e) -> Gen (ArrTinyNE r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (Sz ix) -> Gen (Array r ix e)
forall r ix e.
(Load r ix e, Arbitrary e) =>
Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray ((Int -> Int) -> Sz ix -> Sz ix
forall ix. Index ix => (Int -> Int) -> Sz ix -> Sz ix
liftSz (Int -> Int
forall a. Enum a => a -> a
succ (Int -> Int) -> (Int -> Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
10)) (Sz ix -> Sz ix) -> Gen (Sz ix) -> Gen (Sz ix)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (Sz ix)
forall a. Arbitrary a => Gen a
arbitrary)

instance (Arbitrary ix, Load r ix e, Arbitrary e) =>
         Arbitrary (ArrNE r ix e) where
  arbitrary :: Gen (ArrNE r ix e)
arbitrary = Array r ix e -> ArrNE r ix e
forall r ix e. Array r ix e -> ArrNE r ix e
ArrNE (Array r ix e -> ArrNE r ix e)
-> Gen (Array r ix e) -> Gen (ArrNE r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (Sz ix) -> Gen (Array r ix e)
forall r ix e.
(Load r ix e, Arbitrary e) =>
Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray (SzNE ix -> Sz ix
forall ix. SzNE ix -> Sz ix
unSzNE (SzNE ix -> Sz ix) -> Gen (SzNE ix) -> Gen (Sz ix)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (SzNE ix)
forall a. Arbitrary a => Gen a
arbitrary)


instance (Arbitrary ix, Load r ix e, Arbitrary e) =>
         Arbitrary (ArrIx r ix e) where
  arbitrary :: Gen (ArrIx r ix e)
arbitrary = do
    SzIx Sz ix
sz ix
ix <- Gen (SzIx ix)
forall a. Arbitrary a => Gen a
arbitrary
    Int -> e
func <- Gen (Int -> e)
forall a. Arbitrary a => Gen a
arbitrary
    Comp
comp <- Gen Comp
forall a. Arbitrary a => Gen a
arbitrary
    ArrIx r ix e -> Gen (ArrIx r ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (ArrIx r ix e -> Gen (ArrIx r ix e))
-> ArrIx r ix e -> Gen (ArrIx r ix e)
forall a b. (a -> b) -> a -> b
$ Array r ix e -> ix -> ArrIx r ix e
forall r ix e. Array r ix e -> ix -> ArrIx r ix e
ArrIx (Comp -> Sz ix -> (Int -> e) -> Array r ix e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear Comp
comp Sz ix
sz Int -> e
func) ix
ix


data ArrDW ix e = ArrDW (Array D ix e) (Array DW ix e)

instance (Show ix, Index ix, Ragged L ix e, Load DW ix e, Show e) =>
         Show (ArrDW ix e) where
  show :: ArrDW ix e -> String
show (ArrDW Array D ix e
d Array DW ix e
dw) =
    String
"Delayed:\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++
    Array D ix e -> String
forall a. Show a => a -> String
show Array D ix e
d String -> ShowS
forall a. [a] -> [a] -> [a]
++
    String
"\nCorresponding Windowed:\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++
    --show dw ++
    String
windowInfo
    where
      windowInfo :: String
windowInfo =
        String -> (Window ix e -> String) -> Maybe (Window ix e) -> String
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
          String
"\n No Window"
          (\Window {ix
windowStart :: forall ix e. Window ix e -> ix
windowStart :: ix
windowStart, Sz ix
windowSize :: forall ix e. Window ix e -> Sz ix
windowSize :: Sz ix
windowSize} ->
             String
"\n With Window starting index (" String -> ShowS
forall a. [a] -> [a] -> [a]
++
             ix -> String
forall a. Show a => a -> String
show ix
windowStart String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") and size (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Sz ix -> String
forall a. Show a => a -> String
show Sz ix
windowSize String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")") (Maybe (Window ix e) -> String) -> Maybe (Window ix e) -> String
forall a b. (a -> b) -> a -> b
$
        Array DW ix e -> Maybe (Window ix e)
forall ix e. Array DW ix e -> Maybe (Window ix e)
getWindow Array DW ix e
dw

instance (Arbitrary ix, CoArbitrary ix, Load DW ix e, Arbitrary e, Typeable e) =>
         Arbitrary (ArrDW ix e) where
  arbitrary :: Gen (ArrDW ix e)
arbitrary = do
    ArrTiny (Array D ix e
arr :: Array D ix e) <- Gen (ArrTiny D ix e)
forall a. Arbitrary a => Gen a
arbitrary
    let sz :: Sz ix
sz = Array D ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array D ix e
arr
    Array D ix e -> Array DW ix e -> ArrDW ix e
forall ix e. Array D ix e -> Array DW ix e -> ArrDW ix e
ArrDW Array D ix e
arr (Array DW ix e -> ArrDW ix e)
-> Gen (Array DW ix e) -> Gen (ArrDW ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
      if Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then Array DW ix e -> Gen (Array DW ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Comp -> Sz ix -> (ix -> e) -> Array DW ix e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array D ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array D ix e
arr) Sz ix
sz (Array D ix e -> ix -> e
forall ix r e.
(HasCallStack, Index ix, Source r e) =>
Array r ix e -> ix -> e
evaluate' Array D ix e
arr))
        else do
          ix
wix <- (ix -> ix -> ix) -> ix -> ix -> ix
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Integral a => a -> a -> a
mod) (Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
sz) (ix -> ix) -> Gen ix -> Gen ix
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen ix
forall a. Arbitrary a => Gen a
arbitrary
          ix
wsz <- (Int -> Int) -> ix -> ix
forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (ix -> ix) -> (ix -> ix) -> ix -> ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ix -> ix -> ix) -> ix -> ix -> ix
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Integral a => a -> a -> a
mod) ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (-) (Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
sz) ix
wix) (ix -> ix) -> Gen ix -> Gen ix
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen ix
forall a. Arbitrary a => Gen a
arbitrary
          Array DW ix e -> Gen (Array DW ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array DW ix e -> Gen (Array DW ix e))
-> Array DW ix e -> Gen (Array DW ix e)
forall a b. (a -> b) -> a -> b
$ Array D ix e -> ix -> Sz ix -> (ix -> e) -> Array DW ix e
forall ix r e.
(Index ix, Source r e) =>
Array r ix e -> ix -> Sz ix -> (ix -> e) -> Array DW ix e
makeWindowedArray Array D ix e
arr ix
wix (ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ix
wsz) (Array D ix e -> ix -> e
forall ix r e.
(HasCallStack, Index ix, Source r e) =>
Array r ix e -> ix -> e
evaluate' Array D ix e
arr)