{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- |
-- Module      : Data.Massiv.Array.Mutable
-- Copyright   : (c) Alexey Kuleshevich 2018-2021
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Data.Massiv.Array.Mutable
  ( -- ** Size
    sizeOfMArray
  , msize
  , resizeMArrayM
  , flattenMArray
  , outerSliceMArrayM
  , outerSlicesMArray
    -- ** Element-wise mutation
  , read
  , readM
  , write
  , write_
  , writeM
  , modify
  , modify_
  , modifyM
  , modifyM_
  , swap
  , swap_
  , swapM
  , swapM_
  , zipSwapM_
  -- ** Operations on @MArray@
  -- *** Immutable conversion
  , thaw
  , thawS
  , freeze
  , freezeS
  -- *** Create mutable
  , newMArray
  , newMArray'
  , makeMArray
  , makeMArrayLinear
  , makeMArrayS
  , makeMArrayLinearS
  -- *** Create pure
  , createArray_
  , createArray
  , createArrayS_
  , createArrayS
  , createArrayST_
  , createArrayST
  -- *** Generate
  , generateArray
  , generateArrayLinear
  , generateArrayS
  , generateArrayLinearS
  -- *** Stateful worker threads
  , generateArrayWS
  , generateArrayLinearWS
  -- *** Unfold
  , unfoldrPrimM_
  , iunfoldrPrimM_
  , unfoldrPrimM
  , iunfoldrPrimM
  , unfoldlPrimM_
  , iunfoldlPrimM_
  , unfoldlPrimM
  , iunfoldlPrimM
  -- *** Mapping
  , forPrimM
  , forPrimM_
  , iforPrimM
  , iforPrimM_
  , iforLinearPrimM
  , iforLinearPrimM_
  , for2PrimM_
  , ifor2PrimM_
  -- *** Modify
  , withMArray
  , withMArray_
  , withLoadMArray_
  , withMArrayS
  , withLoadMArrayS
  , withMArrayS_
  , withLoadMArrayS_
  , withMArrayST
  , withLoadMArrayST
  , withMArrayST_
  , withLoadMArrayST_
  -- *** Initialize
  , initialize
  , initializeNew
  -- ** Computation
  , Manifest
  , MArray
  , RealWorld
  , computeInto
  , loadArray
  , loadArrayS
  ) where

-- TODO: add fromListM, et al.

import Data.Maybe (fromMaybe)
import Control.Monad (void, when, unless, (>=>))
import Control.Monad.ST
import Control.Monad.Primitive
import Control.Scheduler
import Data.Massiv.Core.Common
import Data.Massiv.Array.Mutable.Internal
import Data.Massiv.Array.Delayed.Pull (D)
import Prelude hiding (mapM, read)

-- | /O(1)/ - Change the size of a mutable array. Throws
-- `SizeElementsMismatchException` if total number of elements does not match
-- the supplied array.
--
-- @since 1.0.0
resizeMArrayM ::
     (Manifest r e, Index ix', Index ix, MonadThrow m)
  => Sz ix'
  -> MArray s r ix e
  -> m (MArray s r ix' e)
resizeMArrayM :: Sz ix' -> MArray s r ix e -> m (MArray s r ix' e)
resizeMArrayM Sz ix'
sz MArray s r ix e
marr =
  Sz ix' -> MArray s r ix e -> MArray s r ix' e
forall r e ix' ix s.
(Manifest r e, Index ix', Index ix) =>
Sz ix' -> MArray s r ix e -> MArray s r ix' e
unsafeResizeMArray Sz ix'
sz MArray s r ix e
marr MArray s r ix' e -> m () -> m (MArray s r ix' e)
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Sz ix -> Sz ix' -> m ()
forall (m :: * -> *) ix ix'.
(MonadThrow m, Index ix, Index ix') =>
Sz ix -> Sz ix' -> m ()
guardNumberOfElements (MArray s r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray s r ix e
marr) Sz ix'
sz
{-# INLINE resizeMArrayM #-}


-- | /O(1)/ - Change a mutable array to a mutable vector.
--
-- @since 1.0.0
flattenMArray :: (Manifest r e, Index ix) => MArray s r ix e -> MVector s r e
flattenMArray :: MArray s r ix e -> MVector s r e
flattenMArray MArray s r ix e
marr = Sz Ix1 -> MArray s r ix e -> MVector s r e
forall r e ix' ix s.
(Manifest r e, Index ix', Index ix) =>
Sz ix' -> MArray s r ix e -> MArray s r ix' e
unsafeResizeMArray (Sz ix -> Sz Ix1
forall ix. Index ix => Sz ix -> Sz Ix1
toLinearSz (MArray s r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray s r ix e
marr)) MArray s r ix e
marr
{-# INLINE flattenMArray #-}


-- | /O(1)/ - Slice a mutable array from the outside, while reducing its
-- dimensionality by one. Same as `Data.Massiv.Array.!?>` operator, but for
-- mutable arrays.
--
-- @since 1.0.0
outerSliceMArrayM ::
     forall r ix e m s. (MonadThrow m, Index (Lower ix), Index ix, Manifest r e)
  => MArray s r ix e
  -> Ix1
  -> m (MArray s r (Lower ix) e)
outerSliceMArrayM :: MArray s r ix e -> Ix1 -> m (MArray s r (Lower ix) e)
outerSliceMArrayM !MArray s r ix e
marr !Ix1
i = do
  let (Sz Ix1
k, Sz (Lower ix)
szL) = Sz ix -> (Sz Ix1, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Ix1, Sz (Lower ix))
unconsSz (MArray s r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray s r ix e
marr)
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Sz Ix1 -> Ix1 -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz Ix1
k Ix1
i) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ IndexException -> m ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (IndexException -> m ()) -> IndexException -> m ()
forall a b. (a -> b) -> a -> b
$ Sz Ix1 -> Ix1 -> IndexException
forall ix. Index ix => Sz ix -> ix -> IndexException
IndexOutOfBoundsException Sz Ix1
k Ix1
i
  MArray s r (Lower ix) e -> m (MArray s r (Lower ix) e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MArray s r (Lower ix) e -> m (MArray s r (Lower ix) e))
-> MArray s r (Lower ix) e -> m (MArray s r (Lower ix) e)
forall a b. (a -> b) -> a -> b
$ Sz (Lower ix) -> MArray s r Ix1 e -> MArray s r (Lower ix) e
forall r e ix' ix s.
(Manifest r e, Index ix', Index ix) =>
Sz ix' -> MArray s r ix e -> MArray s r ix' e
unsafeResizeMArray Sz (Lower ix)
szL (MArray s r Ix1 e -> MArray s r (Lower ix) e)
-> MArray s r Ix1 e -> MArray s r (Lower ix) e
forall a b. (a -> b) -> a -> b
$ Ix1 -> Sz Ix1 -> MArray s r ix e -> MArray s r Ix1 e
forall r e ix s.
(Manifest r e, Index ix) =>
Ix1 -> Sz Ix1 -> MArray s r ix e -> MVector s r e
unsafeLinearSliceMArray (Ix1
i Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
* Sz (Lower ix) -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz (Lower ix)
szL) (Sz (Lower ix) -> Sz Ix1
forall ix. Index ix => Sz ix -> Sz Ix1
toLinearSz Sz (Lower ix)
szL) MArray s r ix e
marr
{-# INLINE outerSliceMArrayM #-}

-- | /O(1)/ - Take all outer slices of a mutable array and construct a delayed
-- vector out of them. In other words it applies `outerSliceMArrayM` to each
-- outer index. Same as `Data.Massiv.Array.outerSlices` function, but for
-- mutable arrays.
--
-- ====__Examples__
--
-- >>> import Data.Massiv.Array as A
-- >>> arr <- resizeM (Sz2 4 7) $ makeArrayR P Seq (Sz1 28) (+10)
-- >>> arr
-- Array P Seq (Sz (4 :. 7))
--   [ [ 10, 11, 12, 13, 14, 15, 16 ]
--   , [ 17, 18, 19, 20, 21, 22, 23 ]
--   , [ 24, 25, 26, 27, 28, 29, 30 ]
--   , [ 31, 32, 33, 34, 35, 36, 37 ]
--   ]
--
-- Here we can see we can get individual rows from a mutable matrix
--
-- >>> marr <- thawS arr
-- >>> import Control.Monad ((<=<))
-- >>> mapIO_ (print <=< freezeS)  $ outerSlicesMArray Seq marr
-- Array P Seq (Sz1 7)
--   [ 10, 11, 12, 13, 14, 15, 16 ]
-- Array P Seq (Sz1 7)
--   [ 17, 18, 19, 20, 21, 22, 23 ]
-- Array P Seq (Sz1 7)
--   [ 24, 25, 26, 27, 28, 29, 30 ]
-- Array P Seq (Sz1 7)
--   [ 31, 32, 33, 34, 35, 36, 37 ]
--
-- For the sake of example what if our goal was to mutate array in such a way
-- that rows from the top half were swapped with the bottom half:
--
-- >>> (top, bottom) <- splitAtM 1 2 $ outerSlicesMArray Seq marr
-- >>> mapIO_ (print <=< freezeS) top
-- Array P Seq (Sz1 7)
--   [ 10, 11, 12, 13, 14, 15, 16 ]
-- Array P Seq (Sz1 7)
--   [ 17, 18, 19, 20, 21, 22, 23 ]
-- >>> mapIO_ (print <=< freezeS) bottom
-- Array P Seq (Sz1 7)
--   [ 24, 25, 26, 27, 28, 29, 30 ]
-- Array P Seq (Sz1 7)
--   [ 31, 32, 33, 34, 35, 36, 37 ]
-- >>> szipWithM_ (zipSwapM_ 0) top bottom
-- >>> freezeS marr
-- Array P Seq (Sz (4 :. 7))
--   [ [ 24, 25, 26, 27, 28, 29, 30 ]
--   , [ 31, 32, 33, 34, 35, 36, 37 ]
--   , [ 10, 11, 12, 13, 14, 15, 16 ]
--   , [ 17, 18, 19, 20, 21, 22, 23 ]
--   ]
--
-- @since 1.0.0
outerSlicesMArray ::
     forall r ix e s. (Index (Lower ix), Index ix, Manifest r e)
  => Comp
  -> MArray s r ix e
  -> Vector D (MArray s r (Lower ix) e)
outerSlicesMArray :: Comp -> MArray s r ix e -> Vector D (MArray s r (Lower ix) e)
outerSlicesMArray Comp
comp MArray s r ix e
marr =
  Comp
-> Sz Ix1
-> (Ix1 -> MArray s r (Lower ix) e)
-> Vector D (MArray s r (Lower ix) e)
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
comp Sz Ix1
k (\Ix1
i -> Sz (Lower ix) -> MArray s r Ix1 e -> MArray s r (Lower ix) e
forall r e ix' ix s.
(Manifest r e, Index ix', Index ix) =>
Sz ix' -> MArray s r ix e -> MArray s r ix' e
unsafeResizeMArray Sz (Lower ix)
szL (MArray s r Ix1 e -> MArray s r (Lower ix) e)
-> MArray s r Ix1 e -> MArray s r (Lower ix) e
forall a b. (a -> b) -> a -> b
$ Ix1 -> Sz Ix1 -> MArray s r ix e -> MArray s r Ix1 e
forall r e ix s.
(Manifest r e, Index ix) =>
Ix1 -> Sz Ix1 -> MArray s r ix e -> MVector s r e
unsafeLinearSliceMArray (Ix1
i Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
* Sz Ix1 -> Ix1
forall ix. Sz ix -> ix
unSz Sz Ix1
kL) Sz Ix1
kL MArray s r ix e
marr)
  where
    kL :: Sz Ix1
kL = Sz (Lower ix) -> Sz Ix1
forall ix. Index ix => Sz ix -> Sz Ix1
toLinearSz Sz (Lower ix)
szL
    (Sz Ix1
k, Sz (Lower ix)
szL) = Sz ix -> (Sz Ix1, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Ix1, Sz (Lower ix))
unconsSz (Sz ix -> (Sz Ix1, Sz (Lower ix)))
-> Sz ix -> (Sz Ix1, Sz (Lower ix))
forall a b. (a -> b) -> a -> b
$ MArray s r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray s r ix e
marr
{-# INLINE outerSlicesMArray #-}


-- | /O(n)/ - Initialize a new mutable array. All elements will be set to some default value. For
-- boxed arrays it will be a thunk with `Uninitialized` exception, while for others it will be
-- simply zeros.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> marr <- newMArray' (Sz2 2 6) :: IO (MArray RealWorld P Ix2 Int)
-- >>> freeze Seq marr
-- Array P Seq (Sz (2 :. 6))
--   [ [ 0, 0, 0, 0, 0, 0 ]
--   , [ 0, 0, 0, 0, 0, 0 ]
--   ]
--
-- Or using @TypeApplications@:
--
-- >>> :set -XTypeApplications
-- >>> newMArray' @P @Ix2 @Int (Sz2 2 6) >>= freezeS
-- Array P Seq (Sz (2 :. 6))
--   [ [ 0, 0, 0, 0, 0, 0 ]
--   , [ 0, 0, 0, 0, 0, 0 ]
--   ]
-- >>> newMArray' @B @_ @Int (Sz2 2 6) >>= freezeS
-- *** Exception: Uninitialized
--
-- @since 0.6.0
newMArray' ::
     forall r ix e m. (Manifest r e, Index ix, PrimMonad m)
  => Sz ix
  -> m (MArray (PrimState m) r ix e)
newMArray' :: Sz ix -> m (MArray (PrimState m) r ix e)
newMArray' Sz ix
sz = Sz ix -> m (MArray (PrimState m) r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz m (MArray (PrimState m) r ix e)
-> (MArray (PrimState m) r ix e -> m (MArray (PrimState m) r ix e))
-> m (MArray (PrimState m) r ix e)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \MArray (PrimState m) r ix e
ma -> MArray (PrimState m) r ix e
ma MArray (PrimState m) r ix e
-> m () -> m (MArray (PrimState m) r ix e)
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ MArray (PrimState m) r ix e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> m ()
initialize MArray (PrimState m) r ix e
ma
{-# INLINE newMArray' #-}


-- | /O(n)/ - Make a mutable copy of a pure array. Keep in mind that both `freeze` and `thaw` trigger a
-- copy of the full array.
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array
-- >>> :set -XTypeApplications
-- >>> arr <- fromListsM @U @Ix2 @Double Par [[12,21],[13,31]]
-- >>> marr <- thaw arr
-- >>> modify marr (pure . (+ 10)) (1 :. 0)
-- Just 13.0
-- >>> freeze Par marr
-- Array U Par (Sz (2 :. 2))
--   [ [ 12.0, 21.0 ]
--   , [ 23.0, 31.0 ]
--   ]
--
-- @since 0.1.0
thaw :: forall r ix e m. (Manifest r e, Index ix, MonadIO m) => Array r ix e -> m (MArray RealWorld r ix e)
thaw :: Array r ix e -> m (MArray RealWorld r ix e)
thaw Array r ix e
arr =
  IO (MArray RealWorld r ix e) -> m (MArray RealWorld r ix e)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MArray RealWorld r ix e) -> m (MArray RealWorld r ix e))
-> IO (MArray RealWorld r ix e) -> m (MArray RealWorld r ix e)
forall a b. (a -> b) -> a -> b
$ do
    let sz :: Sz ix
sz = Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
        totalLength :: Ix1
totalLength = Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz
    MArray RealWorld r ix e
marr <- Sz ix -> IO (MArray (PrimState IO) r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz
    Comp -> (Scheduler RealWorld () -> IO ()) -> IO ()
withMassivScheduler_ (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) ((Scheduler RealWorld () -> IO ()) -> IO ())
-> (Scheduler RealWorld () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld ()
scheduler ->
      Ix1 -> Ix1 -> (Ix1 -> Ix1 -> IO ()) -> IO ()
forall a. Ix1 -> Ix1 -> (Ix1 -> Ix1 -> a) -> a
splitLinearly (Scheduler RealWorld () -> Ix1
forall s a. Scheduler s a -> Ix1
numWorkers Scheduler RealWorld ()
scheduler) Ix1
totalLength ((Ix1 -> Ix1 -> IO ()) -> IO ()) -> (Ix1 -> Ix1 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ix1
chunkLength Ix1
slackStart -> do
        Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> IO ()) -> IO ()
forall (m :: * -> *) a.
Monad m =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m a) -> m ()
loopM_ Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Ix1
slackStart) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
chunkLength) ((Ix1 -> IO ()) -> IO ()) -> (Ix1 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ !Ix1
start ->
          Scheduler RealWorld () -> IO () -> IO ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler RealWorld ()
scheduler (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Array r ix e
-> Ix1 -> MArray (PrimState IO) r ix e -> Ix1 -> Sz Ix1 -> IO ()
forall r e ix' ix (m :: * -> *).
(Manifest r e, Index ix', Index ix, PrimMonad m) =>
Array r ix' e
-> Ix1 -> MArray (PrimState m) r ix e -> Ix1 -> Sz Ix1 -> m ()
unsafeArrayLinearCopy Array r ix e
arr Ix1
start MArray RealWorld r ix e
MArray (PrimState IO) r ix e
marr Ix1
start (Ix1 -> Sz Ix1
forall ix. ix -> Sz ix
SafeSz Ix1
chunkLength)
        let slackLength :: Ix1
slackLength = Ix1
totalLength Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
- Ix1
slackStart
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ix1
slackLength Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
> Ix1
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
          Scheduler RealWorld () -> IO () -> IO ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler RealWorld ()
scheduler (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
          Array r ix e
-> Ix1 -> MArray (PrimState IO) r ix e -> Ix1 -> Sz Ix1 -> IO ()
forall r e ix' ix (m :: * -> *).
(Manifest r e, Index ix', Index ix, PrimMonad m) =>
Array r ix' e
-> Ix1 -> MArray (PrimState m) r ix e -> Ix1 -> Sz Ix1 -> m ()
unsafeArrayLinearCopy Array r ix e
arr Ix1
slackStart MArray RealWorld r ix e
MArray (PrimState IO) r ix e
marr Ix1
slackStart (Ix1 -> Sz Ix1
forall ix. ix -> Sz ix
SafeSz Ix1
slackLength)
    MArray RealWorld r ix e -> IO (MArray RealWorld r ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure MArray RealWorld r ix e
marr
{-# INLINE thaw #-}

-- | Same as `thaw`, but restrict computation to sequential only.
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array
-- >>> :set -XOverloadedLists
-- >>> thawS @P @Ix1 @Double [1..10]
-- >>> marr <- thawS @P @Ix1 @Double [1..10]
-- >>> writeM marr 5 100
-- >>> freezeS marr
-- Array P Seq (Sz1 10)
--   [ 1.0, 2.0, 3.0, 4.0, 5.0, 100.0, 7.0, 8.0, 9.0, 10.0 ]
--
-- @since 0.3.0
thawS ::
     forall r ix e m. (Manifest r e, Index ix, PrimMonad m)
  => Array r ix e
  -> m (MArray (PrimState m) r ix e)
thawS :: Array r ix e -> m (MArray (PrimState m) r ix e)
thawS Array r ix e
arr = do
  MArray (PrimState m) r ix e
tmarr <- Sz ix -> m (MArray (PrimState m) r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew (Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr)
  Array r ix e
-> Ix1 -> MArray (PrimState m) r ix e -> Ix1 -> Sz Ix1 -> m ()
forall r e ix' ix (m :: * -> *).
(Manifest r e, Index ix', Index ix, PrimMonad m) =>
Array r ix' e
-> Ix1 -> MArray (PrimState m) r ix e -> Ix1 -> Sz Ix1 -> m ()
unsafeArrayLinearCopy Array r ix e
arr Ix1
0 MArray (PrimState m) r ix e
tmarr Ix1
0 (Ix1 -> Sz Ix1
forall ix. ix -> Sz ix
SafeSz (Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem (Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr)))
  MArray (PrimState m) r ix e -> m (MArray (PrimState m) r ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure MArray (PrimState m) r ix e
tmarr
{-# INLINE thawS #-}


-- | /O(n)/ - Yield an immutable copy of the mutable array. Note that mutable representations
-- have to be the same.
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array
-- >>> marr <- newMArray @P (Sz2 2 6) (0 :: Int)
-- >>> forM_ (range Seq 0 (Ix2 1 4)) $ \ix -> write marr ix 9
-- >>> freeze Seq marr
-- Array P Seq (Sz (2 :. 6))
--   [ [ 9, 9, 9, 9, 0, 0 ]
--   , [ 0, 0, 0, 0, 0, 0 ]
--   ]
--
-- @since 0.1.0
freeze ::
     forall r ix e m. (Manifest r e, Index ix, MonadIO m)
  => Comp
  -> MArray RealWorld r ix e
  -> m (Array r ix e)
freeze :: Comp -> MArray RealWorld r ix e -> m (Array r ix e)
freeze Comp
comp MArray RealWorld r ix e
smarr =
  IO (Array r ix e) -> m (Array r ix e)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Array r ix e) -> m (Array r ix e))
-> IO (Array r ix e) -> m (Array r ix e)
forall a b. (a -> b) -> a -> b
$ do
    let sz :: Sz ix
sz = MArray RealWorld r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray RealWorld r ix e
smarr
        totalLength :: Ix1
totalLength = Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz
    MArray RealWorld r ix e
tmarr <- Sz ix -> IO (MArray (PrimState IO) r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz
    Comp -> (Scheduler RealWorld () -> IO ()) -> IO ()
withMassivScheduler_ Comp
comp ((Scheduler RealWorld () -> IO ()) -> IO ())
-> (Scheduler RealWorld () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld ()
scheduler ->
      Ix1 -> Ix1 -> (Ix1 -> Ix1 -> IO ()) -> IO ()
forall a. Ix1 -> Ix1 -> (Ix1 -> Ix1 -> a) -> a
splitLinearly (Scheduler RealWorld () -> Ix1
forall s a. Scheduler s a -> Ix1
numWorkers Scheduler RealWorld ()
scheduler) Ix1
totalLength ((Ix1 -> Ix1 -> IO ()) -> IO ()) -> (Ix1 -> Ix1 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ix1
chunkLength Ix1
slackStart -> do
        Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> IO ()) -> IO ()
forall (m :: * -> *) a.
Monad m =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m a) -> m ()
loopM_ Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Ix1
slackStart) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
chunkLength) ((Ix1 -> IO ()) -> IO ()) -> (Ix1 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ !Ix1
start ->
          Scheduler RealWorld () -> IO () -> IO ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler RealWorld ()
scheduler (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ MArray (PrimState IO) r ix e
-> Ix1 -> MArray (PrimState IO) r ix e -> Ix1 -> Sz Ix1 -> IO ()
forall r e ix' ix (m :: * -> *).
(Manifest r e, Index ix', Index ix, PrimMonad m) =>
MArray (PrimState m) r ix' e
-> Ix1 -> MArray (PrimState m) r ix e -> Ix1 -> Sz Ix1 -> m ()
unsafeLinearCopy MArray RealWorld r ix e
MArray (PrimState IO) r ix e
smarr Ix1
start MArray RealWorld r ix e
MArray (PrimState IO) r ix e
tmarr Ix1
start (Ix1 -> Sz Ix1
forall ix. ix -> Sz ix
SafeSz Ix1
chunkLength)
        let slackLength :: Ix1
slackLength = Ix1
totalLength Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
- Ix1
slackStart
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ix1
slackLength Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
> Ix1
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
          Scheduler RealWorld () -> IO () -> IO ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler RealWorld ()
scheduler (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
          MArray (PrimState IO) r ix e
-> Ix1 -> MArray (PrimState IO) r ix e -> Ix1 -> Sz Ix1 -> IO ()
forall r e ix' ix (m :: * -> *).
(Manifest r e, Index ix', Index ix, PrimMonad m) =>
MArray (PrimState m) r ix' e
-> Ix1 -> MArray (PrimState m) r ix e -> Ix1 -> Sz Ix1 -> m ()
unsafeLinearCopy MArray RealWorld r ix e
MArray (PrimState IO) r ix e
smarr Ix1
slackStart MArray RealWorld r ix e
MArray (PrimState IO) r ix e
tmarr Ix1
slackStart (Ix1 -> Sz Ix1
forall ix. ix -> Sz ix
SafeSz Ix1
slackLength)
    Comp -> MArray (PrimState IO) r ix e -> IO (Array r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
comp MArray RealWorld r ix e
MArray (PrimState IO) r ix e
tmarr
{-# INLINE freeze #-}


-- | Same as `freeze`, but do the copy of supplied muable array sequentially. Also, unlike `freeze`
-- that has to be done in `IO`, `freezeS` can be used with `ST`.
--
-- @since 0.3.0
freezeS ::
     forall r ix e m. (Manifest r e, Index ix, PrimMonad m)
  => MArray (PrimState m) r ix e
  -> m (Array r ix e)
freezeS :: MArray (PrimState m) r ix e -> m (Array r ix e)
freezeS MArray (PrimState m) r ix e
smarr = do
  let sz :: Sz ix
sz = MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
smarr
  MArray (PrimState m) r ix e
tmarr <- Sz ix -> m (MArray (PrimState m) r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz
  MArray (PrimState m) r ix e
-> Ix1 -> MArray (PrimState m) r ix e -> Ix1 -> Sz Ix1 -> m ()
forall r e ix' ix (m :: * -> *).
(Manifest r e, Index ix', Index ix, PrimMonad m) =>
MArray (PrimState m) r ix' e
-> Ix1 -> MArray (PrimState m) r ix e -> Ix1 -> Sz Ix1 -> m ()
unsafeLinearCopy MArray (PrimState m) r ix e
smarr Ix1
0 MArray (PrimState m) r ix e
tmarr Ix1
0 (Ix1 -> Sz Ix1
forall ix. ix -> Sz ix
SafeSz (Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz))
  Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
Seq MArray (PrimState m) r ix e
tmarr
{-# INLINE freezeS #-}

unsafeNewUpper ::
     (Load r' ix e, Manifest r e, PrimMonad m) => Array r' ix e -> m (MArray (PrimState m) r Ix1 e)
unsafeNewUpper :: Array r' ix e -> m (MArray (PrimState m) r Ix1 e)
unsafeNewUpper !Array r' ix e
arr = Sz Ix1 -> m (MArray (PrimState m) r Ix1 e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew (Sz Ix1 -> Maybe (Sz Ix1) -> Sz Ix1
forall a. a -> Maybe a -> a
fromMaybe Sz Ix1
forall ix. Index ix => Sz ix
zeroSz (Array r' ix e -> Maybe (Sz Ix1)
forall r ix e. Shape r ix => Array r ix e -> Maybe (Sz Ix1)
maxLinearSize Array r' ix e
arr))
{-# INLINE unsafeNewUpper #-}

-- | Load sequentially a pure array into the newly created mutable array.
--
-- @since 0.3.0
loadArrayS ::
     forall r ix e r' m. (Load r' ix e, Manifest r e, PrimMonad m)
  => Array r' ix e
  -> m (MArray (PrimState m) r ix e)
loadArrayS :: Array r' ix e -> m (MArray (PrimState m) r ix e)
loadArrayS Array r' ix e
arr = do
  MVector (PrimState m) r e
marr <- Array r' ix e -> m (MVector (PrimState m) r e)
forall r' ix e r (m :: * -> *).
(Load r' ix e, Manifest r e, PrimMonad m) =>
Array r' ix e -> m (MArray (PrimState m) r Ix1 e)
unsafeNewUpper Array r' ix e
arr
  ST (PrimState m) (MArray (PrimState m) r ix e)
-> m (MArray (PrimState m) r ix e)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (MArray (PrimState m) r ix e)
 -> m (MArray (PrimState m) r ix e))
-> ST (PrimState m) (MArray (PrimState m) r ix e)
-> m (MArray (PrimState m) r ix e)
forall a b. (a -> b) -> a -> b
$ MVector (PrimState m) r e
-> Array r' ix e -> ST (PrimState m) (MArray (PrimState m) r ix e)
forall r ix e r' s.
(Load r ix e, Manifest r' e) =>
MVector s r' e -> Array r ix e -> ST s (MArray s r' ix e)
unsafeLoadIntoST MVector (PrimState m) r e
marr Array r' ix e
arr
{-# INLINE loadArrayS #-}


-- | Load a pure array into the newly created mutable array, while respecting computation startegy.
--
-- @since 0.3.0
loadArray ::
     forall r ix e r' m. (Load r' ix e, Manifest r e, MonadIO m)
  => Array r' ix e
  -> m (MArray RealWorld r ix e)
loadArray :: Array r' ix e -> m (MArray RealWorld r ix e)
loadArray Array r' ix e
arr =
  IO (MArray RealWorld r ix e) -> m (MArray RealWorld r ix e)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MArray RealWorld r ix e) -> m (MArray RealWorld r ix e))
-> IO (MArray RealWorld r ix e) -> m (MArray RealWorld r ix e)
forall a b. (a -> b) -> a -> b
$ do
    MVector RealWorld r e
marr <- Array r' ix e -> IO (MArray (PrimState IO) r Ix1 e)
forall r' ix e r (m :: * -> *).
(Load r' ix e, Manifest r e, PrimMonad m) =>
Array r' ix e -> m (MArray (PrimState m) r Ix1 e)
unsafeNewUpper Array r' ix e
arr
    MVector RealWorld r e
-> Array r' ix e -> IO (MArray RealWorld r ix e)
forall r ix e r'.
(Load r ix e, Manifest r' e) =>
MVector RealWorld r' e
-> Array r ix e -> IO (MArray RealWorld r' ix e)
unsafeLoadIntoIO MVector RealWorld r e
marr Array r' ix e
arr
{-# INLINE loadArray #-}



-- | Compute an Array while loading the results into the supplied mutable target array. Number of
-- elements for arrays must agree, otherwise `SizeElementsMismatchException` exception is thrown.
--
-- @since 0.1.3
computeInto ::
     (Size r', Load r' ix' e, Manifest r e, Index ix, MonadIO m)
  => MArray RealWorld r ix e -- ^ Target Array
  -> Array r' ix' e -- ^ Array to load
  -> m ()
computeInto :: MArray RealWorld r ix e -> Array r' ix' e -> m ()
computeInto !MArray RealWorld r ix e
mArr !Array r' ix' e
arr =
  IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem (MArray RealWorld r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray RealWorld r ix e
mArr) Ix1 -> Ix1 -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix' -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem (Array r' ix' e -> Sz ix'
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r' ix' e
arr)) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
      SizeException -> IO ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (SizeException -> IO ()) -> SizeException -> IO ()
forall a b. (a -> b) -> a -> b
$ Sz ix -> Sz ix' -> SizeException
forall ix ix'.
(Index ix, Index ix') =>
Sz ix -> Sz ix' -> SizeException
SizeElementsMismatchException (MArray RealWorld r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray RealWorld r ix e
mArr) (Array r' ix' e -> Sz ix'
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r' ix' e
arr)
    Comp -> (Scheduler RealWorld () -> IO ()) -> IO ()
withMassivScheduler_ (Array r' ix' e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix' e
arr) ((Scheduler RealWorld () -> IO ()) -> IO ())
-> (Scheduler RealWorld () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld ()
scheduler ->
      ST (PrimState IO) () -> IO ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState IO) () -> IO ()) -> ST (PrimState IO) () -> IO ()
forall a b. (a -> b) -> a -> b
$ Scheduler RealWorld ()
-> Array r' ix' e
-> (Ix1 -> e -> ST RealWorld ())
-> ST RealWorld ()
forall r ix e s.
Load r ix e =>
Scheduler s () -> Array r ix e -> (Ix1 -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ Scheduler RealWorld ()
scheduler Array r' ix' e
arr (MArray (PrimState (ST RealWorld)) r ix e
-> Ix1 -> e -> ST RealWorld ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
unsafeLinearWrite MArray RealWorld r ix e
MArray (PrimState (ST RealWorld)) r ix e
mArr)
{-# INLINE computeInto #-}


-- | Create a mutable array using an index aware generating action.
--
-- @since 0.3.0
makeMArrayS ::
     forall r ix e m. (Manifest r e, Index ix, PrimMonad m)
  => Sz ix -- ^ Size of the create array
  -> (ix -> m e) -- ^ Element generating action
  -> m (MArray (PrimState m) r ix e)
makeMArrayS :: Sz ix -> (ix -> m e) -> m (MArray (PrimState m) r ix e)
makeMArrayS Sz ix
sz ix -> m e
f = Sz ix -> (Ix1 -> m e) -> m (MArray (PrimState m) r ix e)
forall r ix e (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> (Ix1 -> m e) -> m (MArray (PrimState m) r ix e)
makeMArrayLinearS Sz ix
sz (ix -> m e
f (ix -> m e) -> (Ix1 -> ix) -> Ix1 -> m e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Ix1 -> ix
forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex Sz ix
sz)
{-# INLINE makeMArrayS #-}


-- | Same as `makeMArrayS`, but index supplied to the action is row-major linear index.
--
-- @since 0.3.0
makeMArrayLinearS ::
     forall r ix e m. (Manifest r e, Index ix, PrimMonad m)
  => Sz ix
  -> (Int -> m e)
  -> m (MArray (PrimState m) r ix e)
makeMArrayLinearS :: Sz ix -> (Ix1 -> m e) -> m (MArray (PrimState m) r ix e)
makeMArrayLinearS Sz ix
sz Ix1 -> m e
f = do
  MArray (PrimState m) r ix e
marr <- Sz ix -> m (MArray (PrimState m) r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz
  Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m a) -> m ()
loopM_ Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr)) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) (\ !Ix1
i -> Ix1 -> m e
f Ix1
i m e -> (e -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix e
marr Ix1
i)
  MArray (PrimState m) r ix e -> m (MArray (PrimState m) r ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return MArray (PrimState m) r ix e
marr
{-# INLINE makeMArrayLinearS #-}

-- | Just like `makeMArrayS`, but also accepts computation strategy and runs in `IO`.
--
-- @since 0.3.0
makeMArray ::
     forall r ix e m. (MonadUnliftIO m, Manifest r e, Index ix)
  => Comp
  -> Sz ix
  -> (ix -> m e)
  -> m (MArray RealWorld r ix e)
makeMArray :: Comp -> Sz ix -> (ix -> m e) -> m (MArray RealWorld r ix e)
makeMArray Comp
comp Sz ix
sz ix -> m e
f = Comp -> Sz ix -> (Ix1 -> m e) -> m (MArray RealWorld r ix e)
forall r ix e (m :: * -> *).
(MonadUnliftIO m, Manifest r e, Index ix) =>
Comp -> Sz ix -> (Ix1 -> m e) -> m (MArray RealWorld r ix e)
makeMArrayLinear Comp
comp Sz ix
sz (ix -> m e
f (ix -> m e) -> (Ix1 -> ix) -> Ix1 -> m e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Ix1 -> ix
forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex Sz ix
sz)
{-# INLINE makeMArray #-}


-- | Just like `makeMArrayLinearS`, but also accepts computation strategy and runs in `IO`.
--
-- @since 0.3.0
makeMArrayLinear ::
     forall r ix e m. (MonadUnliftIO m, Manifest r e, Index ix)
  => Comp
  -> Sz ix
  -> (Int -> m e)
  -> m (MArray RealWorld r ix e)
makeMArrayLinear :: Comp -> Sz ix -> (Ix1 -> m e) -> m (MArray RealWorld r ix e)
makeMArrayLinear Comp
comp Sz ix
sz Ix1 -> m e
f = do
  MArray RealWorld r ix e
marr <- IO (MArray RealWorld r ix e) -> m (MArray RealWorld r ix e)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MArray RealWorld r ix e) -> m (MArray RealWorld r ix e))
-> IO (MArray RealWorld r ix e) -> m (MArray RealWorld r ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> IO (MArray (PrimState IO) r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz
  Comp -> (Scheduler RealWorld () -> m ()) -> m ()
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m ()
withScheduler_ Comp
comp ((Scheduler RealWorld () -> m ()) -> m ())
-> (Scheduler RealWorld () -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld ()
scheduler ->
    ((forall a. m a -> IO a) -> IO ()) -> m ()
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO ()) -> m ())
-> ((forall a. m a -> IO a) -> IO ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run ->
    Scheduler RealWorld ()
-> Ix1 -> (Ix1 -> IO e) -> (Ix1 -> e -> IO ()) -> IO ()
forall s (m :: * -> *) b c.
MonadPrimBase s m =>
Scheduler s () -> Ix1 -> (Ix1 -> m b) -> (Ix1 -> b -> m c) -> m ()
splitLinearlyWithM_ Scheduler RealWorld ()
scheduler (Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz) (m e -> IO e
forall a. m a -> IO a
run (m e -> IO e) -> (Ix1 -> m e) -> Ix1 -> IO e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ix1 -> m e
f) (MArray (PrimState IO) r ix e -> Ix1 -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
unsafeLinearWrite MArray RealWorld r ix e
MArray (PrimState IO) r ix e
marr)
  MArray RealWorld r ix e -> m (MArray RealWorld r ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return MArray RealWorld r ix e
marr
{-# INLINE makeMArrayLinear #-}




-- | Create a new array by supplying an action that will fill the new blank mutable array. Use
-- `createArray` if you'd like to keep the result of the filling function.
--
-- ====__Examples__
--
-- >>> :set -XTypeApplications
-- >>> import Data.Massiv.Array
-- >>> createArray_ @P @_ @Int Seq (Sz1 2) (\ s marr -> scheduleWork s (writeM marr 0 10) >> scheduleWork s (writeM marr 1 11))
-- Array P Seq (Sz1 2)
--   [ 10, 11 ]
--
-- @since 0.3.0
--
createArray_ ::
     forall r ix e a m. (Manifest r e, Index ix, MonadUnliftIO m)
  => Comp -- ^ Computation strategy to use after `MArray` gets frozen and onward.
  -> Sz ix -- ^ Size of the newly created array
  -> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m a)
  -- ^ An action that should fill all elements of the brand new mutable array
  -> m (Array r ix e)
createArray_ :: Comp
-> Sz ix
-> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m a)
-> m (Array r ix e)
createArray_ Comp
comp Sz ix
sz Scheduler RealWorld () -> MArray RealWorld r ix e -> m a
action = do
  MArray RealWorld r ix e
marr <- IO (MArray RealWorld r ix e) -> m (MArray RealWorld r ix e)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MArray RealWorld r ix e) -> m (MArray RealWorld r ix e))
-> IO (MArray RealWorld r ix e) -> m (MArray RealWorld r ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> IO (MArray (PrimState IO) r ix e)
forall r ix e (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
newMArray' Sz ix
sz
  Comp -> (Scheduler RealWorld () -> m a) -> m ()
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m ()
withScheduler_ Comp
comp (Scheduler RealWorld () -> MArray RealWorld r ix e -> m a
`action` MArray RealWorld r ix e
marr)
  IO (Array r ix e) -> m (Array r ix e)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Array r ix e) -> m (Array r ix e))
-> IO (Array r ix e) -> m (Array r ix e)
forall a b. (a -> b) -> a -> b
$ Comp -> MArray (PrimState IO) r ix e -> IO (Array r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
comp MArray RealWorld r ix e
MArray (PrimState IO) r ix e
marr
{-# INLINE createArray_ #-}

-- | Just like `createArray_`, but together with `Array` it returns results of scheduled filling
-- actions.
--
-- @since 0.3.0
--
createArray ::
     forall r ix e a m b. (Manifest r e, Index ix, MonadUnliftIO m)
  => Comp -- ^ Computation strategy to use after `MArray` gets frozen and onward.
  -> Sz ix -- ^ Size of the newly created array
  -> (Scheduler RealWorld a -> MArray RealWorld r ix e -> m b)
  -- ^ An action that should fill all elements of the brand new mutable array
  -> m ([a], Array r ix e)
createArray :: Comp
-> Sz ix
-> (Scheduler RealWorld a -> MArray RealWorld r ix e -> m b)
-> m ([a], Array r ix e)
createArray Comp
comp Sz ix
sz Scheduler RealWorld a -> MArray RealWorld r ix e -> m b
action = do
  MArray RealWorld r ix e
marr <- IO (MArray RealWorld r ix e) -> m (MArray RealWorld r ix e)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MArray RealWorld r ix e) -> m (MArray RealWorld r ix e))
-> IO (MArray RealWorld r ix e) -> m (MArray RealWorld r ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> IO (MArray (PrimState IO) r ix e)
forall r ix e (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
newMArray' Sz ix
sz
  [a]
a <- Comp -> (Scheduler RealWorld a -> m b) -> m [a]
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m [a]
withScheduler Comp
comp (Scheduler RealWorld a -> MArray RealWorld r ix e -> m b
`action` MArray RealWorld r ix e
marr)
  Array r ix e
arr <- IO (Array r ix e) -> m (Array r ix e)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Array r ix e) -> m (Array r ix e))
-> IO (Array r ix e) -> m (Array r ix e)
forall a b. (a -> b) -> a -> b
$ Comp -> MArray (PrimState IO) r ix e -> IO (Array r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
comp MArray RealWorld r ix e
MArray (PrimState IO) r ix e
marr
  ([a], Array r ix e) -> m ([a], Array r ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return ([a]
a, Array r ix e
arr)
{-# INLINE createArray #-}


-- | Create a new array by supplying an action that will fill the new blank mutable array. Use
-- `createArrayS` if you'd like to keep the result of the filling function.
--
-- ====__Examples__
--
-- >>> :set -XTypeApplications
-- >>> import Data.Massiv.Array
-- >>> createArrayS_ @P @_ @Int (Sz1 2) (\ marr -> write marr 0 10 >> write marr 1 12)
-- Array P Seq (Sz1 2)
--   [ 10, 12 ]
--
-- @since 0.3.0
createArrayS_ ::
     forall r ix e a m. (Manifest r e, Index ix, PrimMonad m)
  => Sz ix -- ^ Size of the newly created array
  -> (MArray (PrimState m) r ix e -> m a)
  -- ^ An action that should fill all elements of the brand new mutable array
  -> m (Array r ix e)
createArrayS_ :: Sz ix -> (MArray (PrimState m) r ix e -> m a) -> m (Array r ix e)
createArrayS_ Sz ix
sz MArray (PrimState m) r ix e -> m a
action = (a, Array r ix e) -> Array r ix e
forall a b. (a, b) -> b
snd ((a, Array r ix e) -> Array r ix e)
-> m (a, Array r ix e) -> m (Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sz ix
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
forall r ix e a (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
createArrayS Sz ix
sz MArray (PrimState m) r ix e -> m a
action
{-# INLINE createArrayS_ #-}

-- | Just like `createArray_`, but together with `Array` it returns the result of the filling action.
--
-- @since 0.3.0
createArrayS ::
     forall r ix e a m. (Manifest r e, Index ix, PrimMonad m)
  => Sz ix -- ^ Size of the newly created array
  -> (MArray (PrimState m) r ix e -> m a)
  -- ^ An action that should fill all elements of the brand new mutable array
  -> m (a, Array r ix e)
createArrayS :: Sz ix
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
createArrayS Sz ix
sz MArray (PrimState m) r ix e -> m a
action = do
  MArray (PrimState m) r ix e
marr <- Sz ix -> m (MArray (PrimState m) r ix e)
forall r ix e (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
newMArray' Sz ix
sz
  a
a <- MArray (PrimState m) r ix e -> m a
action MArray (PrimState m) r ix e
marr
  Array r ix e
arr <- Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
Seq MArray (PrimState m) r ix e
marr
  (a, Array r ix e) -> m (a, Array r ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a, Array r ix e
arr)
{-# INLINE createArrayS #-}

-- | Just like `createArrayS_`, but restricted to `ST`.
--
-- @since 0.3.0
createArrayST_ ::
     forall r ix e a. (Manifest r e, Index ix)
  => Sz ix
  -> (forall s. MArray s r ix e -> ST s a)
  -> Array r ix e
createArrayST_ :: Sz ix -> (forall s. MArray s r ix e -> ST s a) -> Array r ix e
createArrayST_ Sz ix
sz forall s. MArray s r ix e -> ST s a
action = (forall s. ST s (Array r ix e)) -> Array r ix e
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Array r ix e)) -> Array r ix e)
-> (forall s. ST s (Array r ix e)) -> Array r ix e
forall a b. (a -> b) -> a -> b
$ Sz ix
-> (MArray (PrimState (ST s)) r ix e -> ST s a)
-> ST s (Array r ix e)
forall r ix e a (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> (MArray (PrimState m) r ix e -> m a) -> m (Array r ix e)
createArrayS_ Sz ix
sz MArray (PrimState (ST s)) r ix e -> ST s a
forall s. MArray s r ix e -> ST s a
action
{-# INLINE createArrayST_ #-}


-- | Just like `createArrayS`, but restricted to `ST`.
--
-- @since 0.2.6
createArrayST ::
     forall r ix e a. (Manifest r e, Index ix)
  => Sz ix
  -> (forall s. MArray s r ix e -> ST s a)
  -> (a, Array r ix e)
createArrayST :: Sz ix -> (forall s. MArray s r ix e -> ST s a) -> (a, Array r ix e)
createArrayST Sz ix
sz forall s. MArray s r ix e -> ST s a
action = (forall s. ST s (a, Array r ix e)) -> (a, Array r ix e)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (a, Array r ix e)) -> (a, Array r ix e))
-> (forall s. ST s (a, Array r ix e)) -> (a, Array r ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix
-> (MArray (PrimState (ST s)) r ix e -> ST s a)
-> ST s (a, Array r ix e)
forall r ix e a (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
createArrayS Sz ix
sz MArray (PrimState (ST s)) r ix e -> ST s a
forall s. MArray s r ix e -> ST s a
action
{-# INLINE createArrayST #-}


-- | Sequentially generate a pure array. Much like `makeArray` creates a pure array this
-- function will use `Manifest` interface to generate a pure `Array` in the end, except that
-- computation strategy is set to `Seq`. Element producing function no longer has to be pure
-- but is a stateful action, becuase it is restricted to `PrimMonad` thus allows for sharing
-- the state between computation of each element.
--
-- ====__Examples__
--
-- >>> import Data.Massiv.Array
-- >>> import Data.IORef
-- >>> ref <- newIORef (0 :: Int)
-- >>> generateArrayS (Sz1 6) (\ i -> modifyIORef' ref (+i) >> print i >> pure i) :: IO (Array U Ix1 Int)
-- 0
-- 1
-- 2
-- 3
-- 4
-- 5
-- Array U Seq (Sz1 6)
--   [ 0, 1, 2, 3, 4, 5 ]
-- >>> readIORef ref
-- 15
--
-- @since 0.2.6
generateArrayS ::
     forall r ix e m. (Manifest r e, Index ix, PrimMonad m)
  => Sz ix -- ^ Resulting size of the array
  -> (ix -> m e) -- ^ Element producing generator
  -> m (Array r ix e)
generateArrayS :: Sz ix -> (ix -> m e) -> m (Array r ix e)
generateArrayS Sz ix
sz ix -> m e
gen = Sz ix -> (Ix1 -> m e) -> m (Array r ix e)
forall r ix e (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> (Ix1 -> m e) -> m (Array r ix e)
generateArrayLinearS Sz ix
sz (ix -> m e
gen (ix -> m e) -> (Ix1 -> ix) -> Ix1 -> m e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Ix1 -> ix
forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex Sz ix
sz)
{-# INLINE generateArrayS #-}

-- | Same as `generateArray` but with action that accepts row-major linear index.
--
-- @since 0.3.0
generateArrayLinearS ::
     forall r ix e m. (Manifest r e, Index ix, PrimMonad m)
  => Sz ix -- ^ Resulting size of the array
  -> (Int -> m e) -- ^ Element producing generator
  -> m (Array r ix e)
generateArrayLinearS :: Sz ix -> (Ix1 -> m e) -> m (Array r ix e)
generateArrayLinearS Sz ix
sz Ix1 -> m e
gen = do
  MArray (PrimState m) r ix e
marr <- Sz ix -> m (MArray (PrimState m) r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz
  Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m a) -> m ()
loopM_ Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr)) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) ((Ix1 -> m ()) -> m ()) -> (Ix1 -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Ix1
i -> Ix1 -> m e
gen Ix1
i m e -> (e -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix e
marr Ix1
i
  Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
Seq MArray (PrimState m) r ix e
marr
{-# INLINE generateArrayLinearS #-}


-- | Just like `generateArrayS`, except this generator __will__ respect the supplied computation
-- strategy, and for that reason it is restricted to `IO`.
--
-- @since 0.2.6
generateArray ::
     forall r ix e m. (MonadUnliftIO m, Manifest r e, Index ix)
  => Comp
  -> Sz ix
  -> (ix -> m e)
  -> m (Array r ix e)
generateArray :: Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
generateArray Comp
comp Sz ix
sz ix -> m e
f = Comp -> Sz ix -> (Ix1 -> m e) -> m (Array r ix e)
forall r ix e (m :: * -> *).
(MonadUnliftIO m, Manifest r e, Index ix) =>
Comp -> Sz ix -> (Ix1 -> m e) -> m (Array r ix e)
generateArrayLinear Comp
comp Sz ix
sz (ix -> m e
f (ix -> m e) -> (Ix1 -> ix) -> Ix1 -> m e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Ix1 -> ix
forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex Sz ix
sz)
{-# INLINE generateArray #-}

-- | Just like `generateArray`, except generating action will receive a row-major linear
-- index.
--
-- @since 0.3.0
generateArrayLinear ::
     forall r ix e m. (MonadUnliftIO m, Manifest r e, Index ix)
  => Comp
  -> Sz ix
  -> (Int -> m e)
  -> m (Array r ix e)
generateArrayLinear :: Comp -> Sz ix -> (Ix1 -> m e) -> m (Array r ix e)
generateArrayLinear Comp
comp Sz ix
sz Ix1 -> m e
f = Comp -> Sz ix -> (Ix1 -> m e) -> m (MArray RealWorld r ix e)
forall r ix e (m :: * -> *).
(MonadUnliftIO m, Manifest r e, Index ix) =>
Comp -> Sz ix -> (Ix1 -> m e) -> m (MArray RealWorld r ix e)
makeMArrayLinear Comp
comp Sz ix
sz Ix1 -> m e
f m (MArray RealWorld r ix e)
-> (MArray RealWorld r ix e -> m (Array r ix e))
-> m (Array r ix e)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO (Array r ix e) -> m (Array r ix e)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Array r ix e) -> m (Array r ix e))
-> (MArray RealWorld r ix e -> IO (Array r ix e))
-> MArray RealWorld r ix e
-> m (Array r ix e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Comp -> MArray (PrimState IO) r ix e -> IO (Array r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
comp
{-# INLINE generateArrayLinear #-}


-- | Same as `generateArrayWS`, but use linear indexing instead.
--
-- @since 0.3.4
generateArrayLinearWS ::
     forall r ix e s m. (Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m)
  => WorkerStates s
  -> Sz ix
  -> (Int -> s -> m e)
  -> m (Array r ix e)
generateArrayLinearWS :: WorkerStates s -> Sz ix -> (Ix1 -> s -> m e) -> m (Array r ix e)
generateArrayLinearWS WorkerStates s
states Sz ix
sz Ix1 -> s -> m e
make = do
  MArray (PrimState m) r ix e
marr <- Sz ix -> m (MArray (PrimState m) r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz
  WorkerStates s -> (SchedulerWS s () -> m ()) -> m ()
forall (m :: * -> *) ws b.
MonadUnliftIO m =>
WorkerStates ws -> (SchedulerWS ws () -> m b) -> m ()
withSchedulerWS_ WorkerStates s
states ((SchedulerWS s () -> m ()) -> m ())
-> (SchedulerWS s () -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \SchedulerWS s ()
schedulerWS ->
    SchedulerWS s ()
-> Ix1 -> (Ix1 -> s -> m e) -> (Ix1 -> e -> m ()) -> m ()
forall (m :: * -> *) ws b c.
MonadUnliftIO m =>
SchedulerWS ws ()
-> Ix1 -> (Ix1 -> ws -> m b) -> (Ix1 -> b -> m c) -> m ()
splitLinearlyWithStatefulM_
      SchedulerWS s ()
schedulerWS
      (Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz)
      Ix1 -> s -> m e
make
      (MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix e
marr)
  Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (WorkerStates s -> Comp
forall ws. WorkerStates ws -> Comp
workerStatesComp WorkerStates s
states) MArray (PrimState m) r ix e
marr
{-# INLINE generateArrayLinearWS #-}

-- | Use per worker thread state while generating elements of the array. Very useful for
-- things that are not thread safe.
--
-- @since 0.3.4
generateArrayWS ::
     forall r ix e s m. (Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m)
  => WorkerStates s
  -> Sz ix
  -> (ix -> s -> m e)
  -> m (Array r ix e)
generateArrayWS :: WorkerStates s -> Sz ix -> (ix -> s -> m e) -> m (Array r ix e)
generateArrayWS WorkerStates s
states Sz ix
sz ix -> s -> m e
make = WorkerStates s -> Sz ix -> (Ix1 -> s -> m e) -> m (Array r ix e)
forall r ix e s (m :: * -> *).
(Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m) =>
WorkerStates s -> Sz ix -> (Ix1 -> s -> m e) -> m (Array r ix e)
generateArrayLinearWS WorkerStates s
states Sz ix
sz (ix -> s -> m e
make (ix -> s -> m e) -> (Ix1 -> ix) -> Ix1 -> s -> m e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Ix1 -> ix
forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex Sz ix
sz)
{-# INLINE generateArrayWS #-}


-- | Sequentially unfold an array from the left.
--
-- ====__Examples__
--
-- Create an array with Fibonacci numbers while performing and `IO` action on the accumulator for
-- each element of the array.
--
-- >>> import Data.Massiv.Array
-- >>> unfoldrPrimM_ (Sz1 10) (\a@(f0, f1) -> let fn = f0 + f1 in print a >> return (f0, (f1, fn))) (0, 1) :: IO (Array P Ix1 Int)
-- (0,1)
-- (1,1)
-- (1,2)
-- (2,3)
-- (3,5)
-- (5,8)
-- (8,13)
-- (13,21)
-- (21,34)
-- (34,55)
-- Array P Seq (Sz1 10)
--   [ 0, 1, 1, 2, 3, 5, 8, 13, 21, 34 ]
--
-- @since 0.3.0
unfoldrPrimM_ ::
     forall r ix e a m. (Manifest r e, Index ix, PrimMonad m)
  => Sz ix -- ^ Size of the desired array
  -> (a -> m (e, a)) -- ^ Unfolding action
  -> a -- ^ Initial accumulator
  -> m (Array r ix e)
unfoldrPrimM_ :: Sz ix -> (a -> m (e, a)) -> a -> m (Array r ix e)
unfoldrPrimM_ Sz ix
sz a -> m (e, a)
gen a
acc0 = (a, Array r ix e) -> Array r ix e
forall a b. (a, b) -> b
snd ((a, Array r ix e) -> Array r ix e)
-> m (a, Array r ix e) -> m (Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sz ix -> (a -> m (e, a)) -> a -> m (a, Array r ix e)
forall r ix e a (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> (a -> m (e, a)) -> a -> m (a, Array r ix e)
unfoldrPrimM Sz ix
sz a -> m (e, a)
gen a
acc0
{-# INLINE unfoldrPrimM_ #-}

-- | Same as `unfoldrPrimM_` but do the unfolding with index aware function.
--
-- @since 0.3.0
iunfoldrPrimM_ ::
     forall r ix e a m. (Manifest r e, Index ix, PrimMonad m)
  => Sz ix -- ^ Size of the desired array
  -> (a -> ix -> m (e, a)) -- ^ Unfolding action
  -> a -- ^ Initial accumulator
  -> m (Array r ix e)
iunfoldrPrimM_ :: Sz ix -> (a -> ix -> m (e, a)) -> a -> m (Array r ix e)
iunfoldrPrimM_ Sz ix
sz a -> ix -> m (e, a)
gen a
acc0 = (a, Array r ix e) -> Array r ix e
forall a b. (a, b) -> b
snd ((a, Array r ix e) -> Array r ix e)
-> m (a, Array r ix e) -> m (Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sz ix -> (a -> ix -> m (e, a)) -> a -> m (a, Array r ix e)
forall r ix e a (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> (a -> ix -> m (e, a)) -> a -> m (a, Array r ix e)
iunfoldrPrimM Sz ix
sz a -> ix -> m (e, a)
gen a
acc0
{-# INLINE iunfoldrPrimM_ #-}


-- | Just like `iunfoldrPrimM_`, but also returns the final value of the accumulator.
--
-- @since 0.3.0
iunfoldrPrimM ::
     forall r ix e a m. (Manifest r e, Index ix, PrimMonad m)
  => Sz ix -- ^ Size of the desired array
  -> (a -> ix -> m (e, a)) -- ^ Unfolding action
  -> a -- ^ Initial accumulator
  -> m (a, Array r ix e)
iunfoldrPrimM :: Sz ix -> (a -> ix -> m (e, a)) -> a -> m (a, Array r ix e)
iunfoldrPrimM Sz ix
sz a -> ix -> m (e, a)
gen a
acc0 =
  Sz ix
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
forall r ix e a (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
unsafeCreateArrayS Sz ix
sz ((MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e))
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
forall a b. (a -> b) -> a -> b
$ \MArray (PrimState m) r ix e
marr ->
    let sz' :: Sz ix
sz' = MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr
     in Sz ix
-> Ix1
-> Ix1
-> Ix1
-> (Ix1 -> Ix1 -> Bool)
-> a
-> (Ix1 -> ix -> a -> m a)
-> m a
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
Sz ix
-> Ix1
-> Ix1
-> Ix1
-> (Ix1 -> Ix1 -> Bool)
-> a
-> (Ix1 -> ix -> a -> m a)
-> m a
iterLinearM Sz ix
sz' Ix1
0 (Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz') Ix1
1 Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
(<) a
acc0 ((Ix1 -> ix -> a -> m a) -> m a) -> (Ix1 -> ix -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \ !Ix1
i ix
ix !a
acc -> do
          (e
e, a
acc') <- a -> ix -> m (e, a)
gen a
acc ix
ix
          MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix e
marr Ix1
i e
e
          a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc'
{-# INLINE iunfoldrPrimM #-}

-- | Just like `iunfoldrPrimM`, but do the unfolding with index aware function.
--
-- @since 0.3.0
unfoldrPrimM ::
     forall r ix e a m. (Manifest r e, Index ix, PrimMonad m)
  => Sz ix -- ^ Size of the desired array
  -> (a -> m (e, a)) -- ^ Unfolding action
  -> a -- ^ Initial accumulator
  -> m (a, Array r ix e)
unfoldrPrimM :: Sz ix -> (a -> m (e, a)) -> a -> m (a, Array r ix e)
unfoldrPrimM Sz ix
sz a -> m (e, a)
gen a
acc0 =
  Sz ix
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
forall r ix e a (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
unsafeCreateArrayS Sz ix
sz ((MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e))
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
forall a b. (a -> b) -> a -> b
$ \MArray (PrimState m) r ix e
marr ->
    let sz' :: Sz ix
sz' = MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr
     in Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
forall (m :: * -> *) a.
Monad m =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
loopM Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz') (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) a
acc0 ((Ix1 -> a -> m a) -> m a) -> (Ix1 -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \ !Ix1
i !a
acc -> do
          (e
e, a
acc') <- a -> m (e, a)
gen a
acc
          MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix e
marr Ix1
i e
e
          a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc'
{-# INLINE unfoldrPrimM #-}

-- | Sequentially unfold an array from the left.
--
-- ====__Examples__
--
-- Create an array with Fibonacci numbers starting at the end while performing and `IO` action on
-- the accumulator for each element of the array.
--
-- >>> import Data.Massiv.Array
-- >>> unfoldlPrimM_ (Sz1 10) (\a@(f0, f1) -> let fn = f0 + f1 in print a >> return ((f1, fn), f0)) (0, 1) :: IO (Array P Ix1 Int)
-- (0,1)
-- (1,1)
-- (1,2)
-- (2,3)
-- (3,5)
-- (5,8)
-- (8,13)
-- (13,21)
-- (21,34)
-- (34,55)
-- Array P Seq (Sz1 10)
--   [ 34, 21, 13, 8, 5, 3, 2, 1, 1, 0 ]
--
-- @since 0.3.0
unfoldlPrimM_ ::
     forall r ix e a m. (Manifest r e, Index ix, PrimMonad m)
  => Sz ix -- ^ Size of the desired array
  -> (a -> m (a, e)) -- ^ Unfolding action
  -> a -- ^ Initial accumulator
  -> m (Array r ix e)
unfoldlPrimM_ :: Sz ix -> (a -> m (a, e)) -> a -> m (Array r ix e)
unfoldlPrimM_ Sz ix
sz a -> m (a, e)
gen a
acc0 = (a, Array r ix e) -> Array r ix e
forall a b. (a, b) -> b
snd ((a, Array r ix e) -> Array r ix e)
-> m (a, Array r ix e) -> m (Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sz ix -> (a -> m (a, e)) -> a -> m (a, Array r ix e)
forall r ix e a (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> (a -> m (a, e)) -> a -> m (a, Array r ix e)
unfoldlPrimM Sz ix
sz a -> m (a, e)
gen a
acc0
{-# INLINE unfoldlPrimM_ #-}

-- | Same as `unfoldlPrimM_` but do the unfolding with index aware function.
--
-- @since 0.3.0
iunfoldlPrimM_ ::
     forall r ix e a m. (Manifest r e, Index ix, PrimMonad m)
  => Sz ix -- ^ Size of the desired array
  -> (a -> ix -> m (a, e)) -- ^ Unfolding action
  -> a -- ^ Initial accumulator
  -> m (Array r ix e)
iunfoldlPrimM_ :: Sz ix -> (a -> ix -> m (a, e)) -> a -> m (Array r ix e)
iunfoldlPrimM_ Sz ix
sz a -> ix -> m (a, e)
gen a
acc0 = (a, Array r ix e) -> Array r ix e
forall a b. (a, b) -> b
snd ((a, Array r ix e) -> Array r ix e)
-> m (a, Array r ix e) -> m (Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sz ix -> (a -> ix -> m (a, e)) -> a -> m (a, Array r ix e)
forall r ix e a (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> (a -> ix -> m (a, e)) -> a -> m (a, Array r ix e)
iunfoldlPrimM Sz ix
sz a -> ix -> m (a, e)
gen a
acc0
{-# INLINE iunfoldlPrimM_ #-}


-- | Just like `iunfoldlPrimM_`, but also returns the final value of the accumulator.
--
-- @since 0.3.0
iunfoldlPrimM ::
     forall r ix e a m. (Manifest r e, Index ix, PrimMonad m)
  => Sz ix -- ^ Size of the desired array
  -> (a -> ix -> m (a, e)) -- ^ Unfolding action
  -> a -- ^ Initial accumulator
  -> m (a, Array r ix e)
iunfoldlPrimM :: Sz ix -> (a -> ix -> m (a, e)) -> a -> m (a, Array r ix e)
iunfoldlPrimM Sz ix
sz a -> ix -> m (a, e)
gen a
acc0 =
  Sz ix
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
forall r ix e a (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
unsafeCreateArrayS Sz ix
sz ((MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e))
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
forall a b. (a -> b) -> a -> b
$ \MArray (PrimState m) r ix e
marr ->
    let sz' :: Sz ix
sz' = MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr
     in Sz ix
-> Ix1
-> Ix1
-> Ix1
-> (Ix1 -> Ix1 -> Bool)
-> a
-> (Ix1 -> ix -> a -> m a)
-> m a
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
Sz ix
-> Ix1
-> Ix1
-> Ix1
-> (Ix1 -> Ix1 -> Bool)
-> a
-> (Ix1 -> ix -> a -> m a)
-> m a
iterLinearM Sz ix
sz' (Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz' Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
- Ix1
1) Ix1
0 (Ix1 -> Ix1
forall a. Num a => a -> a
negate Ix1
1) Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
(>=) a
acc0 ((Ix1 -> ix -> a -> m a) -> m a) -> (Ix1 -> ix -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \ !Ix1
i ix
ix !a
acc -> do
          (a
acc', e
e) <- a -> ix -> m (a, e)
gen a
acc ix
ix
          MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix e
marr Ix1
i e
e
          a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc'
{-# INLINE iunfoldlPrimM #-}

-- | Just like `iunfoldlPrimM`, but do the unfolding with index aware function.
--
-- @since 0.3.0
unfoldlPrimM ::
     forall r ix e a m. (Manifest r e, Index ix, PrimMonad m)
  => Sz ix -- ^ Size of the desired array
  -> (a -> m (a, e)) -- ^ Unfolding action
  -> a -- ^ Initial accumulator
  -> m (a, Array r ix e)
unfoldlPrimM :: Sz ix -> (a -> m (a, e)) -> a -> m (a, Array r ix e)
unfoldlPrimM Sz ix
sz a -> m (a, e)
gen a
acc0 =
  Sz ix
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
forall r ix e a (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
unsafeCreateArrayS Sz ix
sz ((MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e))
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
forall a b. (a -> b) -> a -> b
$ \MArray (PrimState m) r ix e
marr ->
    let sz' :: Sz ix
sz' = MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr
     in Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
forall (m :: * -> *) a.
Monad m =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
loopDeepM Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz') (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+Ix1
1) a
acc0 ((Ix1 -> a -> m a) -> m a) -> (Ix1 -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \ !Ix1
i !a
acc -> do
          (a
acc', e
e) <- a -> m (a, e)
gen a
acc
          MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix e
marr Ix1
i e
e
          a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc'
{-# INLINE unfoldlPrimM #-}

-- | Sequentially loop over a mutable array while reading each element and applying an
-- action to it. There is no mutation to the array, unless the action itself modifies it.
--
-- @since 0.4.0
forPrimM_ :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (e -> m ()) -> m ()
forPrimM_ :: MArray (PrimState m) r ix e -> (e -> m ()) -> m ()
forPrimM_ MArray (PrimState m) r ix e
marr e -> m ()
f =
  Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m a) -> m ()
loopM_ Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr)) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+Ix1
1) (MArray (PrimState m) r ix e -> Ix1 -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Ix1 -> m e
unsafeLinearRead MArray (PrimState m) r ix e
marr (Ix1 -> m e) -> (e -> m ()) -> Ix1 -> m ()
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> e -> m ()
f)
{-# INLINE forPrimM_ #-}

-- | Sequentially loop over a mutable array while modifying each element with an action.
--
-- @since 0.4.0
forPrimM :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (e -> m e) -> m ()
forPrimM :: MArray (PrimState m) r ix e -> (e -> m e) -> m ()
forPrimM MArray (PrimState m) r ix e
marr e -> m e
f =
  Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m e) -> m ()
forall (m :: * -> *) a.
Monad m =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m a) -> m ()
loopM_ Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr)) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+Ix1
1) (MArray (PrimState m) r ix e -> (e -> m e) -> Ix1 -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> (e -> m e) -> Ix1 -> m e
unsafeLinearModify MArray (PrimState m) r ix e
marr e -> m e
f)
{-# INLINE forPrimM #-}


-- | Sequentially loop over a mutable array while reading each element and applying an
-- index aware action to it. There is no mutation to the array, unless the
-- action itself modifies it.
--
-- @since 0.4.0
iforPrimM_ ::
     (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (ix -> e -> m ()) -> m ()
iforPrimM_ :: MArray (PrimState m) r ix e -> (ix -> e -> m ()) -> m ()
iforPrimM_ MArray (PrimState m) r ix e
marr ix -> e -> m ()
f = MArray (PrimState m) r ix e -> (Ix1 -> e -> m ()) -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> (Ix1 -> e -> m ()) -> m ()
iforLinearPrimM_ MArray (PrimState m) r ix e
marr (ix -> e -> m ()
f (ix -> e -> m ()) -> (Ix1 -> ix) -> Ix1 -> e -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Ix1 -> ix
forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr))
{-# INLINE iforPrimM_ #-}

-- | Sequentially loop over a mutable array while modifying each element with an index aware action.
--
-- @since 0.4.0
iforPrimM ::
     (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (ix -> e -> m e) -> m ()
iforPrimM :: MArray (PrimState m) r ix e -> (ix -> e -> m e) -> m ()
iforPrimM MArray (PrimState m) r ix e
marr ix -> e -> m e
f = MArray (PrimState m) r ix e -> (Ix1 -> e -> m e) -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> (Ix1 -> e -> m e) -> m ()
iforLinearPrimM MArray (PrimState m) r ix e
marr (ix -> e -> m e
f (ix -> e -> m e) -> (Ix1 -> ix) -> Ix1 -> e -> m e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Ix1 -> ix
forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr))
{-# INLINE iforPrimM #-}


-- | Sequentially loop over a mutable array while reading each element and applying a
-- linear index aware action to it. There is no mutation to the array, unless the action
-- itself modifies it.
--
-- @since 0.4.0
iforLinearPrimM_ ::
     (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (Int -> e -> m ()) -> m ()
iforLinearPrimM_ :: MArray (PrimState m) r ix e -> (Ix1 -> e -> m ()) -> m ()
iforLinearPrimM_ MArray (PrimState m) r ix e
marr Ix1 -> e -> m ()
f =
  Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m a) -> m ()
loopM_ Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr)) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) (\Ix1
i -> MArray (PrimState m) r ix e -> Ix1 -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Ix1 -> m e
unsafeLinearRead MArray (PrimState m) r ix e
marr Ix1
i m e -> (e -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Ix1 -> e -> m ()
f Ix1
i)
{-# INLINE iforLinearPrimM_ #-}

-- | Sequentially loop over a mutable array while modifying each element with an index aware action.
--
-- @since 0.4.0
iforLinearPrimM ::
     (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (Int -> e -> m e) -> m ()
iforLinearPrimM :: MArray (PrimState m) r ix e -> (Ix1 -> e -> m e) -> m ()
iforLinearPrimM MArray (PrimState m) r ix e
marr Ix1 -> e -> m e
f =
  Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m e) -> m ()
forall (m :: * -> *) a.
Monad m =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> m a) -> m ()
loopM_ Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr)) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) (\Ix1
i -> MArray (PrimState m) r ix e -> (e -> m e) -> Ix1 -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> (e -> m e) -> Ix1 -> m e
unsafeLinearModify MArray (PrimState m) r ix e
marr (Ix1 -> e -> m e
f Ix1
i) Ix1
i)
{-# INLINE iforLinearPrimM #-}



-- | Sequentially loop over the intersection of two mutable arrays while reading
-- elements from both and applying an action to it. There is no mutation to the
-- actual arrays, unless the action itself modifies either one of them.
--
-- @since 1.0.0
for2PrimM_ ::
     forall r1 r2 e1 e2 ix m. (PrimMonad m, Index ix, Manifest r1 e1, Manifest r2 e2)
  => MArray (PrimState m) r1 ix e1
  -> MArray (PrimState m) r2 ix e2
  -> (e1 -> e2 -> m ())
  -> m ()
for2PrimM_ :: MArray (PrimState m) r1 ix e1
-> MArray (PrimState m) r2 ix e2 -> (e1 -> e2 -> m ()) -> m ()
for2PrimM_ MArray (PrimState m) r1 ix e1
m1 MArray (PrimState m) r2 ix e2
m2 e1 -> e2 -> m ()
f = MArray (PrimState m) r1 ix e1
-> MArray (PrimState m) r2 ix e2
-> (ix -> e1 -> e2 -> m ())
-> m ()
forall r1 r2 e1 e2 ix (m :: * -> *).
(PrimMonad m, Index ix, Manifest r1 e1, Manifest r2 e2) =>
MArray (PrimState m) r1 ix e1
-> MArray (PrimState m) r2 ix e2
-> (ix -> e1 -> e2 -> m ())
-> m ()
ifor2PrimM_ MArray (PrimState m) r1 ix e1
m1 MArray (PrimState m) r2 ix e2
m2 ((e1 -> e2 -> m ()) -> ix -> e1 -> e2 -> m ()
forall a b. a -> b -> a
const e1 -> e2 -> m ()
f)
{-# INLINE for2PrimM_ #-}

-- | Same as `for2PrimM_`, but with index aware action.
--
-- @since 1.0.0
ifor2PrimM_ ::
     forall r1 r2 e1 e2 ix m. (PrimMonad m, Index ix, Manifest r1 e1, Manifest r2 e2)
  => MArray (PrimState m) r1 ix e1
  -> MArray (PrimState m) r2 ix e2
  -> (ix -> e1 -> e2 -> m ())
  -> m ()
ifor2PrimM_ :: MArray (PrimState m) r1 ix e1
-> MArray (PrimState m) r2 ix e2
-> (ix -> e1 -> e2 -> m ())
-> m ()
ifor2PrimM_ MArray (PrimState m) r1 ix e1
m1 MArray (PrimState m) r2 ix e2
m2 ix -> e1 -> e2 -> m ()
f = do
  let sz :: ix
sz = (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
forall ix. Index ix => (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
liftIndex2 Ix1 -> Ix1 -> Ix1
forall a. Ord a => a -> a -> a
min (Sz ix -> ix
forall ix. Sz ix -> ix
unSz (MArray (PrimState m) r1 ix e1 -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r1 ix e1
m1)) (Sz ix -> ix
forall ix. Sz ix -> ix
unSz (MArray (PrimState m) r2 ix e2 -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r2 ix e2
m2))
  ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> m ()) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> m a) -> m ()
iterM_ ix
forall ix. Index ix => ix
zeroIndex ix
sz ix
forall ix. Index ix => ix
oneIndex Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
(<) ((ix -> m ()) -> m ()) -> (ix -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ix
ix -> do
    e1
e1 <- MArray (PrimState m) r1 ix e1 -> ix -> m e1
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> m e
unsafeRead MArray (PrimState m) r1 ix e1
m1 ix
ix
    e2
e2 <- MArray (PrimState m) r2 ix e2 -> ix -> m e2
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> m e
unsafeRead MArray (PrimState m) r2 ix e2
m2 ix
ix
    ix -> e1 -> e2 -> m ()
f ix
ix e1
e1 e2
e2
{-# INLINE ifor2PrimM_ #-}


-- | Same as `withMArray_`, but allows to keep artifacts of scheduled tasks.
--
-- @since 0.5.0
withMArray ::
     (Manifest r e, Index ix, MonadUnliftIO m)
  => Array r ix e
  -> (Scheduler RealWorld a -> MArray RealWorld r ix e -> m b)
  -> m ([a], Array r ix e)
withMArray :: Array r ix e
-> (Scheduler RealWorld a -> MArray RealWorld r ix e -> m b)
-> m ([a], Array r ix e)
withMArray Array r ix e
arr Scheduler RealWorld a -> MArray RealWorld r ix e -> m b
action = do
  MArray RealWorld r ix e
marr <- Array r ix e -> m (MArray RealWorld r ix e)
forall r ix e (m :: * -> *).
(Manifest r e, Index ix, MonadIO m) =>
Array r ix e -> m (MArray RealWorld r ix e)
thaw Array r ix e
arr
  [a]
xs <- Comp -> (Scheduler RealWorld a -> m b) -> m [a]
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m [a]
withScheduler (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) (Scheduler RealWorld a -> MArray RealWorld r ix e -> m b
`action` MArray RealWorld r ix e
marr)
  IO ([a], Array r ix e) -> m ([a], Array r ix e)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO ((,) [a]
xs (Array r ix e -> ([a], Array r ix e))
-> IO (Array r ix e) -> IO ([a], Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Comp -> MArray (PrimState IO) r ix e -> IO (Array r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) MArray RealWorld r ix e
MArray (PrimState IO) r ix e
marr)
{-# INLINE withMArray #-}

-- | Create a copy of a pure array, mutate it in place and return its frozen version. The big
-- difference between `withMArrayS` is that it's not only gonna respect the computation strategy
-- supplied to it while making a copy, but it will also pass extra argumens to the action that
-- suppose to modify the mutable copy of the source array. These two extra arguments are:
--
-- * Number of capabilities derived from the `Comp`utation strategy of the array.
--
-- * An action that can be used to schedule arbitrary number of jobs that will be executed in
--   parallel.
--
-- * And, of course, the mutable array itself.
--
-- @since 0.5.0
withMArray_ ::
     (Manifest r e, Index ix, MonadUnliftIO m)
  => Array r ix e
  -> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m a)
  -> m (Array r ix e)
withMArray_ :: Array r ix e
-> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m a)
-> m (Array r ix e)
withMArray_ Array r ix e
arr Scheduler RealWorld () -> MArray RealWorld r ix e -> m a
action = do
  MArray RealWorld r ix e
marr <- Array r ix e -> m (MArray RealWorld r ix e)
forall r ix e (m :: * -> *).
(Manifest r e, Index ix, MonadIO m) =>
Array r ix e -> m (MArray RealWorld r ix e)
thaw Array r ix e
arr
  Comp -> (Scheduler RealWorld () -> m a) -> m ()
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m ()
withScheduler_ (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) (Scheduler RealWorld () -> MArray RealWorld r ix e -> m a
`action` MArray RealWorld r ix e
marr)
  IO (Array r ix e) -> m (Array r ix e)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Array r ix e) -> m (Array r ix e))
-> IO (Array r ix e) -> m (Array r ix e)
forall a b. (a -> b) -> a -> b
$ Comp -> MArray (PrimState IO) r ix e -> IO (Array r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) MArray RealWorld r ix e
MArray (PrimState IO) r ix e
marr
{-# INLINE withMArray_ #-}


-- | Same as `withMArray_`, but the array supplied to this function can be any loadable
-- array. For that reason it will be faster if supplied array is delayed.
--
-- @since 0.6.1
withLoadMArray_ ::
     forall r ix e r' m b. (Load r' ix e, Manifest r e, MonadUnliftIO m)
  => Array r' ix e
  -> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m b)
  -> m (Array r ix e)
withLoadMArray_ :: Array r' ix e
-> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m b)
-> m (Array r ix e)
withLoadMArray_ Array r' ix e
arr Scheduler RealWorld () -> MArray RealWorld r ix e -> m b
action = do
  MArray RealWorld r ix e
marr <- Array r' ix e -> m (MArray RealWorld r ix e)
forall r ix e r' (m :: * -> *).
(Load r' ix e, Manifest r e, MonadIO m) =>
Array r' ix e -> m (MArray RealWorld r ix e)
loadArray Array r' ix e
arr
  Comp -> (Scheduler RealWorld () -> m b) -> m ()
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m ()
withScheduler_ (Array r' ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix e
arr) (Scheduler RealWorld () -> MArray RealWorld r ix e -> m b
`action` MArray RealWorld r ix e
marr)
  IO (Array r ix e) -> m (Array r ix e)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Array r ix e) -> m (Array r ix e))
-> IO (Array r ix e) -> m (Array r ix e)
forall a b. (a -> b) -> a -> b
$ Comp -> MArray (PrimState IO) r ix e -> IO (Array r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (Array r' ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix e
arr) MArray RealWorld r ix e
MArray (PrimState IO) r ix e
marr
{-# INLINE[2] withLoadMArray_ #-}
{-# RULES
"withLoadMArray_/withMArray_" [~2] withLoadMArray_ = withMArray_
"withLoadMArrayS/withMArrayS" [~2] withLoadMArrayS = withMArrayS
"withLoadMArrayS_/withMArrayS_" [~2] withLoadMArrayS_ = withMArrayS_
#-}

-- | Create a copy of a pure array, mutate it in place and return its frozen version. The important
-- benefit over doing a manual `thawS` followed by a `freezeS` is that an array will only be copied
-- once.
--
-- @since 0.5.0
withMArrayS ::
     (Manifest r e, Index ix, PrimMonad m)
  => Array r ix e
  -> (MArray (PrimState m) r ix e -> m a)
  -> m (a, Array r ix e)
withMArrayS :: Array r ix e
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
withMArrayS Array r ix e
arr MArray (PrimState m) r ix e -> m a
action = do
  MArray (PrimState m) r ix e
marr <- Array r ix e -> m (MArray (PrimState m) r ix e)
forall r ix e (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Array r ix e -> m (MArray (PrimState m) r ix e)
thawS Array r ix e
arr
  a
a <- MArray (PrimState m) r ix e -> m a
action MArray (PrimState m) r ix e
marr
  (,) a
a (Array r ix e -> (a, Array r ix e))
-> m (Array r ix e) -> m (a, Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) MArray (PrimState m) r ix e
marr
{-# INLINE withMArrayS #-}


-- | Same as `withMArrayS`, except it discards the value produced by the supplied action
--
-- @since 0.5.0
withMArrayS_ ::
     (Manifest r e, Index ix, PrimMonad m)
  => Array r ix e
  -> (MArray (PrimState m) r ix e -> m a)
  -> m (Array r ix e)
withMArrayS_ :: Array r ix e
-> (MArray (PrimState m) r ix e -> m a) -> m (Array r ix e)
withMArrayS_ Array r ix e
arr MArray (PrimState m) r ix e -> m a
action = (a, Array r ix e) -> Array r ix e
forall a b. (a, b) -> b
snd ((a, Array r ix e) -> Array r ix e)
-> m (a, Array r ix e) -> m (Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Array r ix e
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
forall r e ix (m :: * -> *) a.
(Manifest r e, Index ix, PrimMonad m) =>
Array r ix e
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
withMArrayS Array r ix e
arr MArray (PrimState m) r ix e -> m a
action
{-# INLINE withMArrayS_ #-}


-- | Same as `withMArrayS`, but will work with any loadable array.
--
-- @since 0.6.1
withLoadMArrayS ::
     forall r ix e r' m a. (Load r' ix e, Manifest r e, PrimMonad m)
  => Array r' ix e
  -> (MArray (PrimState m) r ix e -> m a)
  -> m (a, Array r ix e)
withLoadMArrayS :: Array r' ix e
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
withLoadMArrayS Array r' ix e
arr MArray (PrimState m) r ix e -> m a
action = do
  MArray (PrimState m) r ix e
marr <- Array r' ix e -> m (MArray (PrimState m) r ix e)
forall r ix e r' (m :: * -> *).
(Load r' ix e, Manifest r e, PrimMonad m) =>
Array r' ix e -> m (MArray (PrimState m) r ix e)
loadArrayS Array r' ix e
arr
  a
a <- MArray (PrimState m) r ix e -> m a
action MArray (PrimState m) r ix e
marr
  (,) a
a (Array r ix e -> (a, Array r ix e))
-> m (Array r ix e) -> m (a, Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (Array r' ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix e
arr) MArray (PrimState m) r ix e
marr
{-# INLINE[2] withLoadMArrayS #-}

-- | Same as `withMArrayS_`, but will work with any loadable array.
--
-- @since 0.6.1
withLoadMArrayS_ ::
     forall r ix e r' m a. (Load r' ix e, Manifest r e, PrimMonad m)
  => Array r' ix e
  -> (MArray (PrimState m) r ix e -> m a)
  -> m (Array r ix e)
withLoadMArrayS_ :: Array r' ix e
-> (MArray (PrimState m) r ix e -> m a) -> m (Array r ix e)
withLoadMArrayS_ Array r' ix e
arr MArray (PrimState m) r ix e -> m a
action = (a, Array r ix e) -> Array r ix e
forall a b. (a, b) -> b
snd ((a, Array r ix e) -> Array r ix e)
-> m (a, Array r ix e) -> m (Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Array r' ix e
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
forall r ix e r' (m :: * -> *) a.
(Load r' ix e, Manifest r e, PrimMonad m) =>
Array r' ix e
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
withLoadMArrayS Array r' ix e
arr MArray (PrimState m) r ix e -> m a
action
{-# INLINE[2] withLoadMArrayS_ #-}


-- | Same as `withMArrayS` but in `ST`. This is not only pure, but also the safest way to do
-- mutation to the array.
--
-- @since 0.5.0
withMArrayST ::
     (Manifest r e, Index ix)
  => Array r ix e
  -> (forall s . MArray s r ix e -> ST s a)
  -> (a, Array r ix e)
withMArrayST :: Array r ix e
-> (forall s. MArray s r ix e -> ST s a) -> (a, Array r ix e)
withMArrayST Array r ix e
arr forall s. MArray s r ix e -> ST s a
f = (forall s. ST s (a, Array r ix e)) -> (a, Array r ix e)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (a, Array r ix e)) -> (a, Array r ix e))
-> (forall s. ST s (a, Array r ix e)) -> (a, Array r ix e)
forall a b. (a -> b) -> a -> b
$ Array r ix e
-> (MArray (PrimState (ST s)) r ix e -> ST s a)
-> ST s (a, Array r ix e)
forall r e ix (m :: * -> *) a.
(Manifest r e, Index ix, PrimMonad m) =>
Array r ix e
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
withMArrayS Array r ix e
arr MArray (PrimState (ST s)) r ix e -> ST s a
forall s. MArray s r ix e -> ST s a
f
{-# INLINE withMArrayST #-}

-- | Same as `withMArrayS` but in `ST`. This is not only pure, but also the safest way to do
-- mutation to the array.
--
-- @since 0.5.0
withMArrayST_ ::
     (Manifest r e, Index ix) => Array r ix e -> (forall s. MArray s r ix e -> ST s a) -> Array r ix e
withMArrayST_ :: Array r ix e
-> (forall s. MArray s r ix e -> ST s a) -> Array r ix e
withMArrayST_ Array r ix e
arr forall s. MArray s r ix e -> ST s a
f = (forall s. ST s (Array r ix e)) -> Array r ix e
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Array r ix e)) -> Array r ix e)
-> (forall s. ST s (Array r ix e)) -> Array r ix e
forall a b. (a -> b) -> a -> b
$ Array r ix e
-> (MArray (PrimState (ST s)) r ix e -> ST s a)
-> ST s (Array r ix e)
forall r e ix (m :: * -> *) a.
(Manifest r e, Index ix, PrimMonad m) =>
Array r ix e
-> (MArray (PrimState m) r ix e -> m a) -> m (Array r ix e)
withMArrayS_ Array r ix e
arr MArray (PrimState (ST s)) r ix e -> ST s a
forall s. MArray s r ix e -> ST s a
f
{-# INLINE withMArrayST_ #-}


-- | Same as `withMArrayST`, but works with any loadable array.
--
-- @since 0.6.1
withLoadMArrayST ::
     forall r ix e r' a. (Load r' ix e, Manifest r e)
  => Array r' ix e
  -> (forall s. MArray s r ix e -> ST s a)
  -> (a, Array r ix e)
withLoadMArrayST :: Array r' ix e
-> (forall s. MArray s r ix e -> ST s a) -> (a, Array r ix e)
withLoadMArrayST Array r' ix e
arr forall s. MArray s r ix e -> ST s a
f = (forall s. ST s (a, Array r ix e)) -> (a, Array r ix e)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (a, Array r ix e)) -> (a, Array r ix e))
-> (forall s. ST s (a, Array r ix e)) -> (a, Array r ix e)
forall a b. (a -> b) -> a -> b
$ Array r' ix e
-> (MArray (PrimState (ST s)) r ix e -> ST s a)
-> ST s (a, Array r ix e)
forall r ix e r' (m :: * -> *) a.
(Load r' ix e, Manifest r e, PrimMonad m) =>
Array r' ix e
-> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e)
withLoadMArrayS Array r' ix e
arr MArray (PrimState (ST s)) r ix e -> ST s a
forall s. MArray s r ix e -> ST s a
f
{-# INLINE[2] withLoadMArrayST #-}

-- | Same as `withMArrayST_`, but works with any loadable array.
--
-- @since 0.6.1
withLoadMArrayST_ ::
     forall r ix e r' a. (Load r' ix e, Manifest r e)
  => Array r' ix e
  -> (forall s. MArray s r ix e -> ST s a)
  -> Array r ix e
withLoadMArrayST_ :: Array r' ix e
-> (forall s. MArray s r ix e -> ST s a) -> Array r ix e
withLoadMArrayST_ Array r' ix e
arr forall s. MArray s r ix e -> ST s a
f = (forall s. ST s (Array r ix e)) -> Array r ix e
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Array r ix e)) -> Array r ix e)
-> (forall s. ST s (Array r ix e)) -> Array r ix e
forall a b. (a -> b) -> a -> b
$ Array r' ix e
-> (MArray (PrimState (ST s)) r ix e -> ST s a)
-> ST s (Array r ix e)
forall r ix e r' (m :: * -> *) a.
(Load r' ix e, Manifest r e, PrimMonad m) =>
Array r' ix e
-> (MArray (PrimState m) r ix e -> m a) -> m (Array r ix e)
withLoadMArrayS_ Array r' ix e
arr MArray (PrimState (ST s)) r ix e -> ST s a
forall s. MArray s r ix e -> ST s a
f
{-# INLINE[2] withLoadMArrayST_ #-}


-- | /O(1)/ - Lookup an element in the mutable array. Returns `Nothing` when index is out of bounds.
--
-- @since 0.1.0
read :: (Manifest r e, Index ix, PrimMonad m) =>
        MArray (PrimState m) r ix e -> ix -> m (Maybe e)
read :: MArray (PrimState m) r ix e -> ix -> m (Maybe e)
read MArray (PrimState m) r ix e
marr ix
ix =
  if Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr) ix
ix
    then e -> Maybe e
forall a. a -> Maybe a
Just (e -> Maybe e) -> m e -> m (Maybe e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MArray (PrimState m) r ix e -> ix -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> m e
unsafeRead MArray (PrimState m) r ix e
marr ix
ix
    else Maybe e -> m (Maybe e)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe e
forall a. Maybe a
Nothing
{-# INLINE read #-}


-- | /O(1)/ - Same as `read`, but throws `IndexOutOfBoundsException` on an invalid index.
--
-- @since 0.4.0
readM :: (Manifest r e, Index ix, PrimMonad m, MonadThrow m) =>
        MArray (PrimState m) r ix e -> ix -> m e
readM :: MArray (PrimState m) r ix e -> ix -> m e
readM MArray (PrimState m) r ix e
marr ix
ix =
  MArray (PrimState m) r ix e -> ix -> m (Maybe e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> m (Maybe e)
read MArray (PrimState m) r ix e
marr ix
ix m (Maybe e) -> (Maybe e -> m e) -> m e
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just e
e -> e -> m e
forall (f :: * -> *) a. Applicative f => a -> f a
pure e
e
    Maybe e
Nothing -> IndexException -> m e
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (IndexException -> m e) -> IndexException -> m e
forall a b. (a -> b) -> a -> b
$ Sz ix -> ix -> IndexException
forall ix. Index ix => Sz ix -> ix -> IndexException
IndexOutOfBoundsException (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr) ix
ix
{-# INLINE readM #-}


-- | /O(1)/ - Write an element into the cell of a mutable array. Returns `False` when index is out
-- of bounds.
--
-- @since 0.1.0
write :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> e -> m Bool
write :: MArray (PrimState m) r ix e -> ix -> e -> m Bool
write MArray (PrimState m) r ix e
marr ix
ix e
e =
  if Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr) ix
ix
  then MArray (PrimState m) r ix e -> ix -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> e -> m ()
unsafeWrite MArray (PrimState m) r ix e
marr ix
ix e
e m () -> m Bool -> m Bool
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> m Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
  else Bool -> m Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
{-# INLINE write #-}

-- | /O(1)/ - Write an element into the cell of a mutable array. Same as `write` function
-- in case of an out of bounds index it is noop, but unlike `write`, there is no
-- information is returned about was the writing of element successful or not.  In other
-- words, just like `writeM`, but doesn't throw an exception.
--
-- @since 0.4.4
write_ :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> e -> m ()
write_ :: MArray (PrimState m) r ix e -> ix -> e -> m ()
write_ MArray (PrimState m) r ix e
marr ix
ix = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr) ix
ix) (m () -> m ()) -> (e -> m ()) -> e -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MArray (PrimState m) r ix e -> ix -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> e -> m ()
unsafeWrite MArray (PrimState m) r ix e
marr ix
ix
{-# INLINE write_ #-}

-- | /O(1)/ - Same as `write`, but throws `IndexOutOfBoundsException` on an invalid index.
--
-- @since 0.4.0
writeM ::
     (Manifest r e, Index ix, PrimMonad m, MonadThrow m) => MArray (PrimState m) r ix e -> ix -> e -> m ()
writeM :: MArray (PrimState m) r ix e -> ix -> e -> m ()
writeM MArray (PrimState m) r ix e
marr ix
ix e
e =
  MArray (PrimState m) r ix e -> ix -> e -> m Bool
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> e -> m Bool
write MArray (PrimState m) r ix e
marr ix
ix e
e m Bool -> (Bool -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
`unless` IndexException -> m ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (Sz ix -> ix -> IndexException
forall ix. Index ix => Sz ix -> ix -> IndexException
IndexOutOfBoundsException (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr) ix
ix))
{-# INLINE writeM #-}


-- | /O(1)/ - Modify an element in the cell of a mutable array with a supplied
-- action. Returns the previous value, if index was not out of bounds.
--
-- @since 0.1.0
modify ::
     (Manifest r e, Index ix, PrimMonad m)
  => MArray (PrimState m) r ix e -- ^ Array to mutate.
  -> (e -> m e) -- ^ Monadic action that modifies the element
  -> ix -- ^ Index at which to perform modification.
  -> m (Maybe e)
modify :: MArray (PrimState m) r ix e -> (e -> m e) -> ix -> m (Maybe e)
modify MArray (PrimState m) r ix e
marr e -> m e
f ix
ix =
  if Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr) ix
ix
    then e -> Maybe e
forall a. a -> Maybe a
Just (e -> Maybe e) -> m e -> m (Maybe e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MArray (PrimState m) r ix e -> (e -> m e) -> ix -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> (e -> m e) -> ix -> m e
unsafeModify MArray (PrimState m) r ix e
marr e -> m e
f ix
ix
    else Maybe e -> m (Maybe e)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe e
forall a. Maybe a
Nothing
{-# INLINE modify #-}

-- | /O(1)/ - Same as `modify`, except that neither the previous value, nor any
-- information on whether the modification was successful are returned. In other words,
-- just like `modifyM_`, but doesn't throw an exception.
--
-- @since 0.4.4
modify_ ::
     (Manifest r e, Index ix, PrimMonad m)
  => MArray (PrimState m) r ix e -- ^ Array to mutate.
  -> (e -> m e) -- ^ Monadic action that modifies the element
  -> ix -- ^ Index at which to perform modification.
  -> m ()
modify_ :: MArray (PrimState m) r ix e -> (e -> m e) -> ix -> m ()
modify_ MArray (PrimState m) r ix e
marr e -> m e
f ix
ix = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr) ix
ix) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ m e -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m e -> m ()) -> m e -> m ()
forall a b. (a -> b) -> a -> b
$ MArray (PrimState m) r ix e -> (e -> m e) -> ix -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> (e -> m e) -> ix -> m e
unsafeModify MArray (PrimState m) r ix e
marr e -> m e
f ix
ix
{-# INLINE modify_ #-}

-- | /O(1)/ - Modify an element in the cell of a mutable array with a supplied
-- action. Throws an `IndexOutOfBoundsException` exception for invalid index and returns
-- the previous value otherwise.
--
-- @since 0.4.0
modifyM ::
     (Manifest r e, Index ix, PrimMonad m, MonadThrow m)
  => MArray (PrimState m) r ix e -- ^ Array to mutate.
  -> (e -> m e) -- ^ Monadic action that modifies the element
  -> ix -- ^ Index at which to perform modification.
  -> m e
modifyM :: MArray (PrimState m) r ix e -> (e -> m e) -> ix -> m e
modifyM MArray (PrimState m) r ix e
marr e -> m e
f ix
ix
  | Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr) ix
ix = MArray (PrimState m) r ix e -> (e -> m e) -> ix -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> (e -> m e) -> ix -> m e
unsafeModify MArray (PrimState m) r ix e
marr e -> m e
f ix
ix
  | Bool
otherwise = IndexException -> m e
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (Sz ix -> ix -> IndexException
forall ix. Index ix => Sz ix -> ix -> IndexException
IndexOutOfBoundsException (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr) ix
ix)
{-# INLINE modifyM #-}

-- | /O(1)/ - Same as `modifyM`, but discard the returned element
--
-- ====__Examples__
--
-- >>> :set -XTypeApplications
-- >>> import Control.Monad.ST
-- >>> import Data.Massiv.Array
-- >>> runST $ newMArray' @P @Ix1 @Int (Sz1 3) >>= (\ma -> modifyM_ ma (pure . (+10)) 1 >> freezeS ma)
-- Array P Seq (Sz1 3)
--   [ 0, 10, 0 ]
--
-- @since 0.4.0
modifyM_ ::
     (Manifest r e, Index ix, PrimMonad m, MonadThrow m)
  => MArray (PrimState m) r ix e -- ^ Array to mutate.
  -> (e -> m e) -- ^ Monadic action that modifies the element
  -> ix -- ^ Index at which to perform modification.
  -> m ()
modifyM_ :: MArray (PrimState m) r ix e -> (e -> m e) -> ix -> m ()
modifyM_ MArray (PrimState m) r ix e
marr e -> m e
f ix
ix = m e -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m e -> m ()) -> m e -> m ()
forall a b. (a -> b) -> a -> b
$ MArray (PrimState m) r ix e -> (e -> m e) -> ix -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m, MonadThrow m) =>
MArray (PrimState m) r ix e -> (e -> m e) -> ix -> m e
modifyM MArray (PrimState m) r ix e
marr e -> m e
f ix
ix
{-# INLINE modifyM_ #-}


-- | /O(1)/ - Same as `swapM`, but instead of throwing an exception returns `Nothing` when
-- either one of the indices is out of bounds and `Just` elements under those indices
-- otherwise.
--
-- @since 0.1.0
swap :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> ix -> m (Maybe (e, e))
swap :: MArray (PrimState m) r ix e -> ix -> ix -> m (Maybe (e, e))
swap MArray (PrimState m) r ix e
marr ix
ix1 ix
ix2 =
  let !sz :: Sz ix
sz = MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr
   in if Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz ix
ix1 Bool -> Bool -> Bool
&& Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz ix
ix2
        then (e, e) -> Maybe (e, e)
forall a. a -> Maybe a
Just ((e, e) -> Maybe (e, e)) -> m (e, e) -> m (Maybe (e, e))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MArray (PrimState m) r ix e -> ix -> ix -> m (e, e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> ix -> m (e, e)
unsafeSwap MArray (PrimState m) r ix e
marr ix
ix1 ix
ix2
        else Maybe (e, e) -> m (Maybe (e, e))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (e, e)
forall a. Maybe a
Nothing
{-# INLINE swap #-}


-- | /O(1)/ - Same as `swap`, but instead of returning `Nothing` it does nothing. In other
-- words, it is similar to `swapM_`, but does not throw any exceptions.
--
-- @since 0.4.4
swap_ :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> ix -> m ()
swap_ :: MArray (PrimState m) r ix e -> ix -> ix -> m ()
swap_ MArray (PrimState m) r ix e
marr ix
ix1 ix
ix2 =
  let !sz :: Sz ix
sz = MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr
   in Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz ix
ix1 Bool -> Bool -> Bool
&& Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz ix
ix2) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ m (e, e) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m (e, e) -> m ()) -> m (e, e) -> m ()
forall a b. (a -> b) -> a -> b
$ MArray (PrimState m) r ix e -> ix -> ix -> m (e, e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> ix -> m (e, e)
unsafeSwap MArray (PrimState m) r ix e
marr ix
ix1 ix
ix2
{-# INLINE swap_ #-}

-- | /O(1)/ - Swap two elements in a mutable array under the supplied indices. Throws an
-- `IndexOutOfBoundsException` when either one of the indices is out of bounds and
-- elements under those indices otherwise.
--
-- @since 0.4.0
swapM ::
     (Manifest r e, Index ix, PrimMonad m, MonadThrow m)
  => MArray (PrimState m) r ix e
  -> ix -- ^ Index for the first element, which will be returned as the first element in the
        -- tuple.
  -> ix -- ^ Index for the second element, which will be returned as the second element in
        -- the tuple.
  -> m (e, e)
swapM :: MArray (PrimState m) r ix e -> ix -> ix -> m (e, e)
swapM MArray (PrimState m) r ix e
marr ix
ix1 ix
ix2
  | Bool -> Bool
not (Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz ix
ix1) = IndexException -> m (e, e)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (IndexException -> m (e, e)) -> IndexException -> m (e, e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> ix -> IndexException
forall ix. Index ix => Sz ix -> ix -> IndexException
IndexOutOfBoundsException (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr) ix
ix1
  | Bool -> Bool
not (Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz ix
ix2) = IndexException -> m (e, e)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (IndexException -> m (e, e)) -> IndexException -> m (e, e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> ix -> IndexException
forall ix. Index ix => Sz ix -> ix -> IndexException
IndexOutOfBoundsException (MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr) ix
ix2
  | Bool
otherwise = MArray (PrimState m) r ix e -> ix -> ix -> m (e, e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> ix -> m (e, e)
unsafeSwap MArray (PrimState m) r ix e
marr ix
ix1 ix
ix2
  where
    !sz :: Sz ix
sz = MArray (PrimState m) r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr
{-# INLINE swapM #-}


-- | /O(1)/ - Same as `swapM`, but discard the returned elements
--
-- @since 0.4.0
swapM_ ::
     (Manifest r e, Index ix, PrimMonad m, MonadThrow m)
  => MArray (PrimState m) r ix e
  -> ix
  -> ix
  -> m ()
swapM_ :: MArray (PrimState m) r ix e -> ix -> ix -> m ()
swapM_ MArray (PrimState m) r ix e
marr ix
ix1 ix
ix2 = m (e, e) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m (e, e) -> m ()) -> m (e, e) -> m ()
forall a b. (a -> b) -> a -> b
$ MArray (PrimState m) r ix e -> ix -> ix -> m (e, e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m, MonadThrow m) =>
MArray (PrimState m) r ix e -> ix -> ix -> m (e, e)
swapM MArray (PrimState m) r ix e
marr ix
ix1 ix
ix2
{-# INLINE swapM_ #-}

-- | Swap elements in the intersection of two mutable arrays starting at the
-- initial index.
--
-- @since 1.0.0
zipSwapM_ ::
     forall r1 r2 ix e m s. (MonadPrim s m, Manifest r2 e, Manifest r1 e, Index ix)
  => ix
  -> MArray s r1 ix e
  -> MArray s r2 ix e
  -> m ()
zipSwapM_ :: ix -> MArray s r1 ix e -> MArray s r2 ix e -> m ()
zipSwapM_ ix
startIx MArray s r1 ix e
m1 MArray s r2 ix e
m2 = do
  let sz1 :: Sz ix
sz1 = MArray s r1 ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray s r1 ix e
m1
      sz2 :: Sz ix
sz2 = MArray s r2 ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray s r2 ix e
m2
      sz :: ix
sz = (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
forall ix. Index ix => (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
liftIndex2 Ix1 -> Ix1 -> Ix1
forall a. Ord a => a -> a -> a
min (Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
sz1) (Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
sz2)
  ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> m ()) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> m a) -> m ()
iterM_ ix
startIx ix
sz ix
forall ix. Index ix => ix
oneIndex Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
(<) ((ix -> m ()) -> m ()) -> (ix -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ix
ix -> do
    let i1 :: Ix1
i1 = Sz ix -> ix -> Ix1
forall ix. Index ix => Sz ix -> ix -> Ix1
toLinearIndex Sz ix
sz1 ix
ix
        i2 :: Ix1
i2 = Sz ix -> ix -> Ix1
forall ix. Index ix => Sz ix -> ix -> Ix1
toLinearIndex Sz ix
sz2 ix
ix
    e
e1 <- MArray (PrimState m) r1 ix e -> Ix1 -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Ix1 -> m e
unsafeLinearRead MArray s r1 ix e
MArray (PrimState m) r1 ix e
m1 Ix1
i1
    e
e2 <- MArray (PrimState m) r2 ix e -> Ix1 -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Ix1 -> m e
unsafeLinearRead MArray s r2 ix e
MArray (PrimState m) r2 ix e
m2 Ix1
i2
    MArray (PrimState m) r2 ix e -> Ix1 -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
unsafeLinearWrite MArray s r2 ix e
MArray (PrimState m) r2 ix e
m2 Ix1
i2 e
e1
    MArray (PrimState m) r1 ix e -> Ix1 -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Ix1 -> e -> m ()
unsafeLinearWrite MArray s r1 ix e
MArray (PrimState m) r1 ix e
m1 Ix1
i1 e
e2
{-# INLINE zipSwapM_ #-}

-- | Get the size of a mutable array.
--
-- @since 0.1.0
msize :: (Manifest r e, Index ix) => MArray s r ix e -> Sz ix
msize :: MArray s r ix e -> Sz ix
msize = MArray s r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray
{-# DEPRECATED msize "In favor of `sizeOfMArray`" #-}