MapSpec.hs

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

module Test.Massiv.Array.Ops.MapSpec (spec) where

import Control.Monad.ST
import Control.Scheduler
import Data.Foldable as F
import Data.Massiv.Array as A
import Data.Massiv.Array.Unsafe
import Test.Massiv.Core
import Prelude as P

prop_zipUnzip
  :: (Index ix, Show (Array D ix Int))
  => Array D ix Int
  -> Array D ix Int
  -> Property
prop_zipUnzip arr1 arr2 =
  (extract' zeroIndex sz arr1, extract' zeroIndex sz arr2) === A.unzip (A.zip arr1 arr2)
  where
    sz = Sz (liftIndex2 min (unSz (size arr1)) (unSz (size arr2)))

prop_zipFlip
  :: (Index ix, Show (Array D ix (Int, Int)))
  => Array D ix Int
  -> Array D ix Int
  -> Property
prop_zipFlip arr1 arr2 =
  A.zip arr1 arr2
    === A.map (\(e2, e1) -> (e1, e2)) (A.zip arr2 arr1)

prop_zipUnzip3
  :: (Index ix, Show (Array D ix Int))
  => Array D ix Int
  -> Array D ix Int
  -> Array D ix Int
  -> Property
prop_zipUnzip3 arr1 arr2 arr3 =
  (extract' zeroIndex sz arr1, extract' zeroIndex sz arr2, extract' zeroIndex sz arr3)
    === A.unzip3 (A.zip3 arr1 arr2 arr3)
  where
    sz =
      Sz (liftIndex2 min (liftIndex2 min (unSz (size arr1)) (unSz (size arr2))) (unSz (size arr3)))

prop_zipFlip3
  :: (Index ix, Show (Array D ix (Int, Int, Int)))
  => Array D ix Int
  -> Array D ix Int
  -> Array D ix Int
  -> Property
prop_zipFlip3 arr1 arr2 arr3 =
  A.zip3 arr1 arr2 arr3 === A.map (\(e3, e2, e1) -> (e1, e2, e3)) (A.zip3 arr3 arr2 arr1)

prop_zipUnzip4
  :: (Index ix, Show (Array D ix Int))
  => Array D ix Int
  -> Array D ix Int
  -> Array D ix Int
  -> Array D ix Int
  -> Property
prop_zipUnzip4 arr1 arr2 arr3 arr4 =
  ( extract' zeroIndex sz arr1
  , extract' zeroIndex sz arr2
  , extract' zeroIndex sz arr3
  , extract' zeroIndex sz arr4
  )
    === A.unzip4 (A.zip4 arr1 arr2 arr3 arr4)
  where
    sz = sfoldl (liftSz2 min) (size arr1) $ smap size $ sfromList [arr2, arr3, arr4]

prop_zipFlip4
  :: (Index ix, Show (Array D ix (Int, Int, Int, Int)))
  => Array D ix Int
  -> Array D ix Int
  -> Array D ix Int
  -> Array D ix Int
  -> Property
prop_zipFlip4 arr1 arr2 arr3 arr4 =
  A.zip4 arr1 arr2 arr3 arr4
    === A.map (\(e4, e3, e2, e1) -> (e1, e2, e3, e4)) (A.zip4 arr4 arr3 arr2 arr1)

prop_zip4
  :: (Index ix, Show (Array D ix (Int, Int, Int, Int)))
  => Array D ix Int
  -> Array D ix Int
  -> Array D ix Int
  -> Array D ix Int
  -> Property
prop_zip4 arr1 arr2 arr3 arr4 =
  let f = (,,,)
   in A.zip4 arr1 arr2 arr3 arr4
        === A.zipWith (\(e1, e2) (e3, e4) -> f e1 e2 e3 e4) (A.zip arr1 arr2) (A.zip arr3 arr4)

prop_itraverseA
  :: (Index ix, Show (Array U ix Int)) => Array D ix Int -> Fun (ix, Int) Int -> Property
prop_itraverseA arr fun =
  alt_imapM (\ix -> Just . applyFun2Compat fun ix) arr
    === itraverseA @U (\ix -> Just . applyFun2Compat fun ix) arr

mapSpec
  :: forall ix
   . ( Arbitrary ix
     , CoArbitrary ix
     , Index ix
     , Function ix
     , Show (Array U ix Int)
     , Show (Array D ix Int)
     , Show (Array D ix (Int, Int))
     , Show (Array D ix (Int, Int, Int))
     , Show (Array D ix (Int, Int, Int, Int))
     )
  => Spec
mapSpec = do
  describe "Zipping" $ do
    it "zipUnzip" $ property $ prop_zipUnzip @ix
    it "zipFlip" $ property $ prop_zipFlip @ix
    it "zipUnzip3" $ property $ prop_zipUnzip3 @ix
    it "zipFlip3" $ property $ prop_zipFlip3 @ix
    it "zipUnzip4" $ property $ prop_zipUnzip4 @ix
    it "zipFlip4" $ property $ prop_zipFlip4 @ix
    it "zip" $ property $ prop_zip4 @ix
  describe "Traversing" $ do
    it "itraverseA" $ property $ prop_itraverseA @ix
  describe "StatefulMapping" $ do
    it "mapWS" $ property $ prop_MapWS @ix

spec :: Spec
spec = do
  describe "Ix1" $ mapSpec @Ix1
  describe "Ix2" $ mapSpec @Ix2
  describe "Ix3" $ mapSpec @Ix3
  describe "Ix4" $ mapSpec @Ix4

alt_imapM
  :: (Applicative f, Index ix, Manifest r2 b, Source r1 a)
  => (ix -> a -> f b)
  -> Array r1 ix a
  -> f (Array r2 ix b)
alt_imapM f arr = fmap loadList $ P.traverse (uncurry f) $ foldrS (:) [] (imap (,) arr)
  where
    loadList xs =
      runST $ do
        marr <- unsafeNew (size arr)
        _ <- F.foldlM (\i e -> unsafeLinearWrite marr i e >> return (i + 1)) 0 xs
        unsafeFreeze (getComp arr) marr
    {-# INLINE loadList #-}

prop_MapWS :: (Show (Array U ix Int), Index ix) => Array U ix Int -> Property
prop_MapWS arr =
  monadicIO $
    run $ do
      let comp = getComp arr
      count <- getCompWorkers comp
      arrStates <- newMArray' @P (Sz count)
      states <- initWorkerStates comp (\(WorkerId i) -> pure $ \f -> modifyM_ arrStates f i)
      arr' <-
        forWS states arr $ \e smod -> do
          smod $ \acc -> pure (acc + e)
          pure e
      accsArr <- freeze Seq arrStates
      pure (A.sum arr' === A.sum accsArr .&&. arr === arr')