{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# 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
  { forall r ix e. 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
  { forall r ix e. 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
  { forall r ix e. 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. HasCallStack => [(Int, Gen a)] -> Gen a
frequency
      [ (Int
20, Comp -> Gen Comp
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Comp
Seq)
      , (Int
10, Comp -> Gen Comp
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Comp
Par)
      , (Int
10, Comp -> Gen Comp
forall a. a -> Gen a
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 :: forall r ix e.
(Load r ix e, Arbitrary e) =>
Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray Gen (Sz ix)
szGen =
  [Gen (Array r ix e)] -> Gen (Array r ix e)
forall a. HasCallStack => [Gen a] -> Gen a
oneof
    [ do
        Comp
comp <- Gen Comp
forall a. Arbitrary a => Gen a
arbitrary
        Sz ix
sz <- Gen (Sz ix)
szGen
        Fun Int e
f <- Gen (Fun Int e)
forall a. Arbitrary a => Gen a
arbitrary
        Array r ix e -> Gen (Array r ix e)
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array r ix e -> Gen (Array r ix e))
-> Array r ix e -> Gen (Array r ix e)
forall a b. (a -> b) -> a -> b
$ Comp -> Sz ix -> (ix -> e) -> Array r ix e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
comp Sz ix
sz (Fun Int e -> Int -> e
forall a b. Fun a b -> a -> b
applyFun Fun Int e
f (Int -> e) -> (ix -> Int) -> ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
sz)
    , 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 a b. Gen (a -> b) -> Gen a -> Gen b
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 a b. Gen (a -> b) -> Gen a -> Gen b
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, Index ix, Arbitrary e) => Arbitrary (Array D ix e) where
  arbitrary :: Gen (Array D ix e)
arbitrary = Gen (Sz ix) -> Gen (Array D ix e)
forall r ix e.
(Load r ix e, Arbitrary e) =>
Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray Gen (Sz ix)
forall a. Arbitrary a => Gen a
arbitrary

instance (Arbitrary ix, Index ix, Arbitrary e) => Arbitrary (Array DL ix e) where
  arbitrary :: Gen (Array DL ix e)
arbitrary = Gen (Sz ix) -> Gen (Array DL ix e)
forall r ix e.
(Load r ix e, Arbitrary e) =>
Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray Gen (Sz ix)
forall a. Arbitrary a => Gen a
arbitrary
instance (Arbitrary ix, Index ix, Arbitrary e) => Arbitrary (Array DI ix e) where
  arbitrary :: Gen (Array DI ix e)
arbitrary = Gen (Sz ix) -> Gen (Array DI ix e)
forall r ix e.
(Load r ix e, Arbitrary e) =>
Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray Gen (Sz ix)
forall a. Arbitrary a => Gen a
arbitrary
instance (Arbitrary ix, Load DW ix e, Arbitrary e) => Arbitrary (Array DW ix e) where
  arbitrary :: Gen (Array DW ix e)
arbitrary = Gen (Sz ix) -> Gen (Array DW ix e)
forall r ix e.
(Load r ix e, Arbitrary e) =>
Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray Gen (Sz ix)
forall a. Arbitrary a => Gen a
arbitrary
instance (ix ~ Ix1, Arbitrary e) => Arbitrary (Array DS ix e) where
  arbitrary :: Gen (Array DS ix e)
arbitrary = Gen (Sz ix) -> Gen (Array DS ix e)
forall r ix e.
(Load r ix e, Arbitrary e) =>
Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray Gen (Sz ix)
forall a. Arbitrary a => Gen a
arbitrary
instance (Arbitrary ix, Index ix, Arbitrary e) => Arbitrary (Array B ix e) where
  arbitrary :: Gen (Array B ix e)
arbitrary = Gen (Sz ix) -> Gen (Array B ix e)
forall r ix e.
(Load r ix e, Arbitrary e) =>
Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray Gen (Sz ix)
forall a. Arbitrary a => Gen a
arbitrary
instance (Arbitrary ix, Index ix, Arbitrary e) => Arbitrary (Array BL ix e) where
  arbitrary :: Gen (Array BL ix e)
arbitrary = Gen (Sz ix) -> Gen (Array BL ix e)
forall r ix e.
(Load r ix e, Arbitrary e) =>
Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray Gen (Sz ix)
forall a. Arbitrary a => Gen a
arbitrary
instance (Arbitrary ix, Index ix, Arbitrary e, NFData e) => Arbitrary (Array BN ix e) where
  arbitrary :: Gen (Array BN ix e)
arbitrary = Gen (Sz ix) -> Gen (Array BN ix e)
forall r ix e.
(Load r ix e, Arbitrary e) =>
Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray Gen (Sz ix)
forall a. Arbitrary a => Gen a
arbitrary
instance (Arbitrary ix, Index ix, Arbitrary e, Prim e) => Arbitrary (Array P ix e) where
  arbitrary :: Gen (Array P ix e)
arbitrary = Gen (Sz ix) -> Gen (Array P ix e)
forall r ix e.
(Load r ix e, Arbitrary e) =>
Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray Gen (Sz ix)
forall a. Arbitrary a => Gen a
arbitrary
instance (Arbitrary ix, Index ix, Arbitrary e, Storable e) => Arbitrary (Array S ix e) where
  arbitrary :: Gen (Array S ix e)
arbitrary = Gen (Sz ix) -> Gen (Array S ix e)
forall r ix e.
(Load r ix e, Arbitrary e) =>
Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray Gen (Sz ix)
forall a. Arbitrary a => Gen a
arbitrary
instance (Arbitrary ix, Index ix, Arbitrary e, Unbox e) => Arbitrary (Array U ix e) where
  arbitrary :: Gen (Array U ix e)
arbitrary = Gen (Sz ix) -> Gen (Array U ix e)
forall r ix e.
(Load r ix e, Arbitrary e) =>
Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray Gen (Sz ix)
forall a. Arbitrary a => Gen a
arbitrary
instance (Arbitrary ix, Load L ix e, Arbitrary e) => Arbitrary (Array L ix e) where
  arbitrary :: Gen (Array L ix e)
arbitrary = Gen (Sz ix) -> Gen (Array L ix e)
forall r ix e.
(Load r ix e, Arbitrary e) =>
Gen (Sz ix) -> Gen (Array r ix e)
arbitraryArray Gen (Sz ix)
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.
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 a. a -> Gen a
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 :: ix
windowStart :: forall ix e. Window ix e -> ix
windowStart, Sz ix
windowSize :: Sz ix
windowSize :: forall ix e. Window ix e -> 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
forall ix e. Array D 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 a. a -> Gen a
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
forall ix e. Array D 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 a. a -> Gen a
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)