{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Data.Massiv.Array.Mutable
(
read
, read'
, write
, write'
, modify
, modify'
, swap
, swap'
, Mutable
, MArray
, msize
, new
, thaw
, freeze
, createArray_
, createArray
, createArrayST_
, createArrayST
, generateArray
, generateArrayIO
, unfoldlPrim_
, unfoldlPrim
, withMArray
, withMArrayST
, RealWorld
, computeInto
) where
import Prelude hiding (mapM, read)
import Control.Monad (unless)
import Control.Monad.Primitive (PrimMonad (..))
import Control.Monad.ST
import Data.Massiv.Array.Manifest.Internal
import Data.Massiv.Array.Unsafe
import Data.Massiv.Core.Common
import Data.Massiv.Core.Scheduler
new :: (Mutable r ix e, PrimMonad m) => ix -> m (MArray (PrimState m) r ix e)
new sz = unsafeNewZero (liftIndex (max 0) sz)
{-# INLINE new #-}
thaw :: (Mutable r ix e, PrimMonad m) => Array r ix e -> m (MArray (PrimState m) r ix e)
thaw = unsafeThaw . clone
{-# INLINE thaw #-}
freeze :: (Mutable r ix e, PrimMonad m) => Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
freeze comp marr = clone <$> unsafeFreeze comp marr
{-# INLINE freeze #-}
createArray_ ::
(Mutable r ix e, PrimMonad m)
=> Comp
-> ix
-> (MArray (PrimState m) r ix e -> m a)
-> m (Array r ix e)
createArray_ comp sz action = fmap snd $ createArray comp sz action
{-# INLINE createArray_ #-}
createArray ::
(Mutable r ix e, PrimMonad m)
=> Comp
-> ix
-> (MArray (PrimState m) r ix e -> m a)
-> m (a, Array r ix e)
createArray comp sz action = do
marr <- new sz
a <- action marr
arr <- unsafeFreeze comp marr
return (a, arr)
{-# INLINE createArray #-}
createArrayST_ ::
Mutable r ix e => Comp -> ix -> (forall s. MArray s r ix e -> ST s a) -> Array r ix e
createArrayST_ comp sz action = runST $ createArray_ comp sz action
{-# INLINE createArrayST_ #-}
createArrayST ::
Mutable r ix e => Comp -> ix -> (forall s. MArray s r ix e -> ST s a) -> (a, Array r ix e)
createArrayST comp sz action = runST $ createArray comp sz action
{-# INLINE createArrayST #-}
generateArray ::
(Mutable r ix e, PrimMonad m)
=> Comp
-> ix
-> (ix -> m e)
-> m (Array r ix e)
generateArray comp sz' gen = do
let sz = liftIndex (max 0) sz'
marr <- unsafeNew sz
iterM_ zeroIndex (msize marr) (pureIndex 1) (<) $ \ix -> gen ix >>= write marr ix
unsafeFreeze comp marr
{-# INLINE generateArray #-}
generateArrayIO ::
(Mutable r ix e)
=> Comp
-> ix
-> (ix -> IO e)
-> IO (Array r ix e)
generateArrayIO comp sz' gen = do
case comp of
Seq -> generateArray comp sz' gen
ParOn wids -> do
let sz = liftIndex (max 0) sz'
marr <- unsafeNew sz
withScheduler_ wids $ \scheduler ->
splitLinearlyWithM_
(numWorkers scheduler)
(scheduleWork scheduler)
(totalElem sz)
(gen . fromLinearIndex sz)
(unsafeLinearWrite marr)
unsafeFreeze comp marr
{-# INLINE generateArrayIO #-}
unfoldlPrim_ ::
(Mutable r ix e, PrimMonad m)
=> Comp
-> ix
-> (a -> ix -> m (a, e))
-> a
-> m (Array r ix e)
unfoldlPrim_ comp sz gen acc0 = fmap snd $ unfoldlPrim comp sz gen acc0
{-# INLINE unfoldlPrim_ #-}
unfoldlPrim ::
(Mutable r ix e, PrimMonad m)
=> Comp
-> ix
-> (a -> ix -> m (a, e))
-> a
-> m (a, Array r ix e)
unfoldlPrim comp sz gen acc0 =
createArray comp sz $ \marr ->
let sz' = msize marr
in iterLinearM sz' 0 (totalElem sz') 1 (<) acc0 $ \i ix acc -> do
(acc', e) <- gen acc ix
unsafeLinearWrite marr i e
return acc'
{-# INLINE unfoldlPrim #-}
withMArray ::
(Mutable r ix e, PrimMonad m)
=> Array r ix e
-> (MArray (PrimState m) r ix e -> m a)
-> m (Array r ix e)
withMArray arr action = do
marr <- thaw arr
_ <- action marr
unsafeFreeze (getComp arr) marr
{-# INLINE withMArray #-}
withMArrayST ::
Mutable r ix e
=> Array r ix e
-> (forall s . MArray s r ix e -> ST s a)
-> Array r ix e
withMArrayST arr f = runST $ withMArray arr f
{-# INLINE withMArrayST #-}
read :: (Mutable r ix e, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> m (Maybe e)
read marr ix =
if isSafeIndex (msize marr) ix
then Just <$> unsafeRead marr ix
else return Nothing
{-# INLINE read #-}
read' :: (Mutable r ix e, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> m e
read' marr ix = do
mval <- read marr ix
case mval of
Just e -> return e
Nothing -> errorIx "Data.Massiv.Array.Mutable.read'" (msize marr) ix
{-# INLINE read' #-}
write :: (Mutable r ix e, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> e -> m Bool
write marr ix e =
if isSafeIndex (msize marr) ix
then unsafeWrite marr ix e >> return True
else return False
{-# INLINE write #-}
write' :: (Mutable r ix e, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> e -> m ()
write' marr ix e =
write marr ix e >>= (`unless` errorIx "Data.Massiv.Array.Mutable.write'" (msize marr) ix)
{-# INLINE write' #-}
modify :: (Mutable r ix e, PrimMonad m) =>
MArray (PrimState m) r ix e -> (e -> e) -> ix -> m Bool
modify marr f ix =
if isSafeIndex (msize marr) ix
then do
val <- unsafeRead marr ix
unsafeWrite marr ix $ f val
return True
else return False
{-# INLINE modify #-}
modify' :: (Mutable r ix e, PrimMonad m) =>
MArray (PrimState m) r ix e -> (e -> e) -> ix -> m ()
modify' marr f ix =
modify marr f ix >>= (`unless` errorIx "Data.Massiv.Array.Mutable.modify'" (msize marr) ix)
{-# INLINE modify' #-}
swap :: (Mutable r ix e, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> ix -> m Bool
swap marr ix1 ix2 = do
let sz = msize marr
if isSafeIndex sz ix1 && isSafeIndex sz ix2
then do
val1 <- unsafeRead marr ix1
val2 <- unsafeRead marr ix2
unsafeWrite marr ix1 val2
unsafeWrite marr ix2 val1
return True
else return False
{-# INLINE swap #-}
swap' :: (Mutable r ix e, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> ix -> m ()
swap' marr ix1 ix2 = do
success <- swap marr ix1 ix2
unless success $
errorIx "Data.Massiv.Array.Mutable.swap'" (msize marr) $
if isSafeIndex (msize marr) ix1
then ix2
else ix1
{-# INLINE swap' #-}