Mutable.hs

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

module Test.Massiv.Core.Mutable (
  -- * Spec for Mutable instance
  unsafeMutableSpec,
  prop_UnsafeNewMsize,
  prop_UnsafeThawFreeze,
  prop_UnsafeInitializeNew,
  prop_UnsafeArrayLinearCopy,

  -- ** Properties that aren't valid for boxed
  unsafeMutableUnboxedSpec,
  prop_UnsafeInitialize,
) where

import Data.Massiv.Array as A
import Data.Massiv.Array.Unsafe
import Test.Massiv.Core.Common
import Test.Massiv.Utils

prop_UnsafeNewMsize
  :: forall r ix e
   . (Arbitrary ix, Index ix, Manifest r e)
  => Property
prop_UnsafeNewMsize = property $ \sz -> do
  marr :: MArray RealWorld r ix e <- unsafeNew sz
  sz `shouldBe` sizeOfMArray marr

prop_UnsafeNewLinearWriteRead
  :: forall r ix e
   . (Eq e, Show e, Manifest r e, Index ix, Arbitrary ix, Arbitrary e)
  => Property
prop_UnsafeNewLinearWriteRead = property $ \(SzIx sz ix) e1 e2 -> do
  marr :: MArray RealWorld r ix e <- unsafeNew sz
  let i = toLinearIndex sz ix
  unsafeLinearWrite marr i e1
  unsafeLinearRead marr i `shouldReturn` e1
  unsafeLinearModify marr (\ !_ -> pure e2) i `shouldReturn` e1
  unsafeLinearRead marr i `shouldReturn` e2

prop_UnsafeThawFreeze
  :: forall r ix e
   . (Eq (Array r ix e), Show (Array r ix e), Index ix, Manifest r e)
  => Array r ix e
  -> Property
prop_UnsafeThawFreeze arr = arr === runST (unsafeFreeze (getComp arr) =<< unsafeThaw arr)

prop_UnsafeInitializeNew
  :: forall r ix e
   . ( Eq (Array r ix e)
     , Show (Array r ix e)
     , Show e
     , Arbitrary e
     , Arbitrary ix
     , Index ix
     , Manifest r e
     )
  => Property
prop_UnsafeInitializeNew =
  property $ \comp sz e ->
    (compute (A.replicate comp sz e :: Array DL ix e) :: Array r ix e)
      === runST (unsafeFreeze comp =<< initializeNew (Just e) sz)

prop_UnsafeInitialize
  :: forall r ix e
   . ( Eq (Array r ix e)
     , Show (Array r ix e)
     , Arbitrary ix
     , Index ix
     , Manifest r e
     )
  => Property
prop_UnsafeInitialize =
  property $ \comp sz ->
    runST $ do
      marr1 :: MArray s r ix e <- unsafeNew sz
      initialize marr1
      marr2 :: MArray s r ix e <- initializeNew Nothing sz
      (===) <$> unsafeFreeze comp marr1 <*> unsafeFreeze comp marr2

prop_UnsafeLinearCopy
  :: forall r ix e
   . (Eq (Array r ix e), Show (Array r ix e), Index ix, Manifest r e)
  => Array r ix e
  -> Property
prop_UnsafeLinearCopy arr =
  (arr, arr)
    === runST
      ( do
          let sz = size arr
          marrs <- thawS arr
          marrd <- unsafeNew sz
          unsafeLinearCopy marrs 0 marrd 0 (Sz (totalElem sz))
          arrd <- unsafeFreeze (getComp arr) marrd
          arrs <- unsafeFreeze (getComp arr) marrs
          pure (arrs, arrd)
      )

prop_UnsafeLinearCopyPart
  :: forall r ix e
   . ( Eq (Vector r e)
     , Show (Vector r e)
     , Eq (Array r ix e)
     , Show (Array r ix e)
     , Manifest r e
     , Index ix
     )
  => ArrIx r ix e
  -> NonNegative Ix1
  -> Ix1
  -> Property
prop_UnsafeLinearCopyPart (ArrIx arr ix) (NonNegative delta) toOffset =
  arr === arrs .&&. slice' i k (flatten arr) === slice' j k arrd
  where
    sz = size arr
    i = toLinearIndex sz ix
    j = max 0 (i + toOffset)
    k = Sz (totalElem sz - i - delta)
    sz' = Sz (j + unSz k)
    (arrs, arrd) =
      runST $ do
        marrs <- thawS arr -- make sure that the source does not get modified
        marrd <- unsafeNew sz'
        unsafeLinearCopy marrs i marrd j k
        (,) <$> unsafeFreeze (getComp arr) marrs <*> unsafeFreeze (getComp arr) marrd

prop_UnsafeArrayLinearCopy
  :: forall r ix e
   . (Eq (Array r ix e), Show (Array r ix e), Index ix, Manifest r e)
  => Array r ix e
  -> Property
prop_UnsafeArrayLinearCopy arr =
  arr
    === runST
      ( do
          let sz = size arr
          marr <- unsafeNew sz
          unsafeArrayLinearCopy arr 0 marr 0 (Sz (totalElem sz))
          unsafeFreeze (getComp arr) marr
      )

prop_UnsafeArrayLinearCopyPart
  :: forall r ix e
   . (Eq (Vector r e), Show (Vector r e), Index ix, Manifest r e)
  => ArrIx r ix e
  -> NonNegative Ix1
  -> Ix1
  -> Property
prop_UnsafeArrayLinearCopyPart (ArrIx arr ix) (NonNegative delta) toOffset =
  slice' i k (flatten arr) === slice' j k arr'
  where
    sz = size arr
    i = toLinearIndex sz ix
    j = max 0 (i + toOffset)
    k = Sz (totalElem sz - i - delta)
    sz' = Sz (j + unSz k)
    arr' =
      runST $ do
        marr <- unsafeNew sz'
        unsafeArrayLinearCopy arr i marr j k
        unsafeFreeze (getComp arr) marr

prop_UnsafeLinearSet
  :: forall r ix e
   . ( Eq (Vector r e)
     , Show (Vector r e)
     , Index ix
     , Manifest r e
     )
  => Comp
  -> SzIx ix
  -> NonNegative Ix1
  -> e
  -> Property
prop_UnsafeLinearSet comp (SzIx sz ix) (NonNegative delta) e =
  compute (A.replicate Seq k e :: Array DL Ix1 e)
    === slice' i k (flatten (arrd :: Array r ix e))
  where
    i = toLinearIndex sz ix
    k = Sz (totalElem sz - i - delta)
    arrd =
      runST $ do
        marrd <- unsafeNew sz
        unsafeLinearSet marrd i k e
        unsafeFreeze comp marrd

prop_UnsafeLinearShrink
  :: forall r ix e
   . ( Eq (Vector r e)
     , Show (Vector r e)
     , Manifest r e
     , Index ix
     )
  => ArrIx r ix e
  -> Property
prop_UnsafeLinearShrink (ArrIx arr ix) =
  slice' 0 k (flatten arr) === slice' 0 k (flatten arr')
  where
    sz = size arr
    sz' = Sz (liftIndex2 (-) (unSz sz) ix)
    k = Sz (totalElem sz')
    arr' =
      runST $ do
        marr <- thawS arr
        marr' <- unsafeLinearShrink marr sz'
        unsafeFreeze (getComp arr) marr'

prop_UnsafeLinearGrow
  :: forall r ix e
   . ( Eq (Array r ix e)
     , Show (Array r ix e)
     , Eq (Vector r e)
     , Show (Vector r e)
     , Manifest r e
     , Index ix
     )
  => ArrIx r ix e
  -> e
  -> Property
prop_UnsafeLinearGrow (ArrIx arr ix) e =
  slice' 0 k (flatten arr)
    === slice' 0 k (flatten arrGrown)
    .&&. arrCopied
    === arrGrown
  where
    sz = size arr
    sz' = Sz (liftIndex2 (+) (unSz sz) ix)
    k = Sz (totalElem sz)
    (arrCopied, arrGrown) =
      runST $ do
        marrCopied <- unsafeNew sz'
        unsafeArrayLinearCopy arr 0 marrCopied 0 k
        marr <- thawS arr
        marrGrown <- unsafeLinearGrow marr sz'
        when (sz' /= sz) $ do
          unsafeLinearSet marrGrown (totalElem sz) (Sz (totalElem sz' - totalElem sz)) e
          unsafeLinearSet marrCopied (totalElem sz) (Sz (totalElem sz' - totalElem sz)) e
        (,) <$> unsafeFreeze (getComp arr) marrCopied <*> unsafeFreeze (getComp arr) marrGrown

prop_UnsafeLinearSliceMArray
  :: forall r ix e
   . (HasCallStack, Index ix, Manifest r e, Eq (Vector r e), Show (Vector r e))
  => Array r ix e
  -> Property
prop_UnsafeLinearSliceMArray arr =
  forAll genLinearRegion $ \(i, k) ->
    propIO $ do
      marr <- thawS arr
      unsafeFreeze Seq (unsafeLinearSliceMArray i k marr) `shouldReturn` unsafeLinearSlice i k arr
  where
    n = totalElem (size arr)
    genLinearRegion = do
      k <- chooseInt (0, n)
      i <- chooseInt (0, n - k)
      pure (i, Sz k)

unsafeMutableSpec
  :: forall r ix e
   . ( Eq (Vector r e)
     , Show (Vector r e)
     , Eq (Array r ix e)
     , Show (Array r ix e)
     , Manifest r e
     , Show e
     , Eq e
     , Load r ix e
     , Arbitrary (Array r ix e)
     , Arbitrary e
     , Arbitrary ix
     , Typeable e
     )
  => Spec
unsafeMutableSpec =
  describe ("Manifest (" ++ showsArrayType @r @ix @e ") (Unsafe)") $ do
    prop "UnsafeNewMsize" $ prop_UnsafeNewMsize @r @ix @e
    prop "UnsafeNewLinearWriteRead" $ prop_UnsafeNewLinearWriteRead @r @ix @e
    prop "UnsafeThawFreeze" $ prop_UnsafeThawFreeze @r @ix @e
    prop "UnsafeInitializeNew" $ prop_UnsafeInitializeNew @r @ix @e
    prop "UnsafeLinearSet" $ prop_UnsafeLinearSet @r @ix @e
    prop "UnsafeLinearCopy" $ prop_UnsafeLinearCopy @r @ix @e
    prop "UnsafeLinearCopyPart" $ prop_UnsafeLinearCopyPart @r @ix @e
    prop "UnsafeArrayLinearCopy" $ prop_UnsafeArrayLinearCopy @r @ix @e
    prop "UnsafeArrayLinearCopyPart" $ prop_UnsafeArrayLinearCopyPart @r @ix @e
    prop "UnsafeLinearShrink" $ prop_UnsafeLinearShrink @r @ix @e
    prop "UnsafeLinearGrow" $ prop_UnsafeLinearGrow @r @ix @e
    prop "UnsafeLinearSliceMArray" $ prop_UnsafeLinearSliceMArray @r @ix @e

unsafeMutableUnboxedSpec
  :: forall r ix e
   . ( Typeable e
     , Eq (Array r ix e)
     , Show (Array r ix e)
     , Index ix
     , Arbitrary ix
     , Manifest r e
     )
  => Spec
unsafeMutableUnboxedSpec =
  describe ("Manifest Unboxed (" ++ showsArrayType @r @ix @e ") (Unsafe)") $
    it "UnsafeInitialize" $
      prop_UnsafeInitialize @r @ix @e