ArraySpec.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

-- Here are tests for all instances for all main classes
module Test.Massiv.ArraySpec (
  spec,
) where

import Data.Massiv.Array
import Test.Massiv.Core

prop_Construct_makeArray_Manifest
  :: forall r ix
   . (Ragged L ix Int, Source r Int, Load r ix Int)
  => Comp
  -> Sz ix
  -> Fun Int Int
  -> Property
prop_Construct_makeArray_Manifest comp sz f =
  makeArrayLinearR D comp sz (apply f)
    === delay (setComp Seq (makeArray comp sz (apply f . toLinearIndex sz) :: Array r ix Int))

prop_Construct_makeArray_Delayed
  :: forall r ix
   . (Ragged L ix Int, Load r ix Int)
  => Comp
  -> Sz ix
  -> Fun Int Int
  -> Property
prop_Construct_makeArray_Delayed comp sz f =
  makeArrayLinearR P comp sz (apply f)
    === compute (setComp Seq (makeArrayLinear comp sz (apply f)) :: Array r ix Int)

prop_Functor
  :: forall r ix
   . (Ragged L ix Int, Load r ix Int, Functor (Array r ix))
  => Comp
  -> Sz ix
  -> Fun Int Int
  -> Fun Int Int
  -> Property
prop_Functor comp sz f g =
  makeArrayLinearR P comp sz (apply g . apply f)
    === compute (fmap (apply g) (makeArrayLinear comp sz (apply f) :: Array r ix Int))

prop_Extract
  :: forall r ix
   . ( Ragged L ix Int
     , Load r ix Int
     , Source r Int
     )
  => Comp
  -> Sz ix
  -> Fun Int Int
  -> ix
  -> Sz ix
  -> Property
prop_Extract comp sz f start newSize =
  (computeAs P <$> toStringException (extractM start newSize arrD))
    === (compute <$> toStringException (extractM start newSize arr))
  where
    arrD = makeArrayLinearR D comp sz (apply f)
    arr = makeArrayLinear comp sz (apply f) :: Array r ix Int

prop_IxUnbox
  :: forall ix
   . ( Ragged L ix ix
     , Source U ix
     , Unbox ix
     )
  => Comp
  -> Sz ix
  -> Fun Int ix
  -> Property
prop_IxUnbox comp sz f =
  makeArrayLinearR D comp sz (apply f)
    === delay (makeArrayLinear comp sz (apply f) :: Array U ix ix)

prop_computeWithStride
  :: forall r ix
   . (Ragged L ix Int, StrideLoad r ix Int)
  => Comp
  -> Sz ix
  -> Fun Int Int
  -> Stride ix
  -> Property
prop_computeWithStride comp sz f stride =
  (arr === computeWithStride stride arrL)
    .&&. (arr === compute (fromStrideLoad stride arrL))
  where
    arrL = makeArrayLinear comp sz (apply f) :: Array r ix Int
    arr = computeWithStrideAs P stride (makeArrayLinearR D comp sz (apply f))

specCommon
  :: forall ix
   . (Arbitrary ix, StrideLoad DW ix Int, Ragged L ix Int, Ragged L ix ix, Unbox ix)
  => Spec
specCommon =
  describe "Construct" $ do
    prop "Construct_makeArray B" $ prop_Construct_makeArray_Manifest @B @ix
    prop "Construct_makeArray BN" $ prop_Construct_makeArray_Manifest @BN @ix
    prop "Construct_makeArray BL" $ prop_Construct_makeArray_Manifest @BL @ix
    prop "Construct_makeArray S" $ prop_Construct_makeArray_Manifest @S @ix
    prop "Construct_makeArray P" $ prop_Construct_makeArray_Manifest @P @ix
    prop "Construct_makeArray U" $ prop_Construct_makeArray_Manifest @U @ix
    prop "Construct_makeArray_Delayed DI" $ prop_Construct_makeArray_Delayed @DI @ix
    prop "Construct_makeArray_Delayed DL" $ prop_Construct_makeArray_Delayed @DL @ix
    prop "Construct_makeArray_Delayed DW" $ prop_Construct_makeArray_Delayed @DW @ix
    prop "Functor D" $ prop_Functor @D @ix
    prop "Functor DI" $ prop_Functor @DI @ix
    prop "Functor DL" $ prop_Functor @DL @ix
    prop "Functor DW" $ prop_Functor @DW @ix
    prop "Extract B" $ prop_Extract @B @ix
    prop "Extract BN" $ prop_Extract @BN @ix
    prop "Extract BL" $ prop_Extract @BL @ix
    prop "Extract S" $ prop_Extract @S @ix
    prop "Extract U" $ prop_Extract @U @ix
    prop "computeWithStride DI" $ prop_computeWithStride @DI @ix
    prop "computeWithStride DW" $ prop_computeWithStride @DW @ix
    prop "computeWithStride B" $ prop_computeWithStride @B @ix
    prop "computeWithStride BN" $ prop_computeWithStride @BN @ix
    prop "computeWithStride BL" $ prop_computeWithStride @BL @ix
    prop "computeWithStride S" $ prop_computeWithStride @S @ix
    prop "computeWithStride U" $ prop_computeWithStride @U @ix
    prop "IxUnbox" $ prop_IxUnbox @ix

spec :: Spec
spec = do
  specCommon @Ix1
  specCommon @Ix2
  specCommon @Ix3

-- FIXME: Uses too much RAM when compiling
-- specCommon @Ix4
-- specCommon @Ix5