{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module      : Data.Massiv.Array.Manifest.Internal
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
module Data.Massiv.Array.Manifest.Internal (
  Manifest (..),
  Array (..),
  flattenMArray,
  compute,
  computeS,
  computeP,
  computeIO,
  computePrimM,
  computeAs,
  computeProxy,
  computeSource,
  computeWithStride,
  computeWithStrideAs,
  clone,
  convert,
  convertAs,
  convertProxy,
  gcastArr,
  fromRaggedArrayM,
  fromRaggedArray',
  unsafeLoadIntoS,
  unsafeLoadIntoM,
  iterateUntil,
  iterateUntilM,
) where

import Control.DeepSeq
import Control.Exception (try)
import Control.Monad.Primitive
import Control.Monad.ST
import Control.Scheduler
import Data.Massiv.Array.Delayed.Pull
import Data.Massiv.Array.Mutable
import Data.Massiv.Array.Mutable.Internal (unsafeCreateArray_)
import Data.Massiv.Core.Common
import Data.Massiv.Core.List
import Data.Maybe (fromMaybe)
import Data.Typeable
import System.IO.Unsafe (unsafePerformIO)

-- | Ensure that Array is computed, i.e. represented with concrete elements in memory, hence is the
-- `Mutable` type class restriction. Use `setComp` if you'd like to change computation strategy
-- before calling @compute@
--
-- @since 0.1.0
compute :: forall r ix e r'. (Manifest r e, Load r' ix e) => Array r' ix e -> Array r ix e
compute :: forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute !Array r' ix e
arr = IO (Array r ix e) -> Array r ix e
forall a. IO a -> a
unsafePerformIO (IO (Array r ix e) -> Array r ix e)
-> IO (Array r ix e) -> Array r ix e
forall a b. (a -> b) -> a -> b
$ Array r' ix e -> IO (Array r ix e)
forall r ix e r' (m :: * -> *).
(Manifest r e, Load r' ix e, MonadIO m) =>
Array r' ix e -> m (Array r ix e)
computeIO Array r' ix e
arr
{-# INLINE compute #-}

-- | Compute array sequentially disregarding predefined computation strategy. Very much
-- the same as `computePrimM`, but executed in `ST`, thus pure.
--
-- @since 0.1.0
computeS :: forall r ix e r'. (Manifest r e, Load r' ix e) => Array r' ix e -> Array r ix e
computeS :: forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
computeS !Array r' ix e
arr = (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 -> ST s (Array r ix e)
forall r ix e r' (m :: * -> *).
(Manifest r e, Load r' ix e, PrimMonad m) =>
Array r' ix e -> m (Array r ix e)
computePrimM Array r' ix e
arr
{-# INLINE computeS #-}

-- | Compute array in parallel using all cores disregarding predefined computation
-- strategy. Computation stategy of the resulting array will match the source, despite
-- that it is diregarded.
--
-- @since 0.5.4
computeP
  :: forall r ix e r'
   . (Manifest r e, Load r' ix e)
  => Array r' ix e
  -> Array r ix e
computeP :: forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
computeP Array r' ix e
arr = Comp -> Array r ix e -> Array r ix e
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
forall ix e. Comp -> Array r ix e -> Array r ix e
setComp (Array r' ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r' ix e -> Comp
getComp Array r' ix e
arr) (Array r ix e -> Array r ix e) -> Array r ix e -> Array r ix e
forall a b. (a -> b) -> a -> b
$ Array r' ix e -> Array r ix e
forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute (Comp -> Array r' ix e -> Array r' ix e
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
forall ix e. Comp -> Array r' ix e -> Array r' ix e
setComp Comp
Par Array r' ix e
arr)
{-# INLINE computeP #-}

-- | Very similar to `compute`, but computes an array inside the `IO` monad. Despite being
-- deterministic and referentially transparent, because this is an `IO` action it
-- can be very useful for enforcing the order of evaluation. Should be a prefered way of
-- computing an array during benchmarking.
--
-- @since 0.4.5
computeIO
  :: forall r ix e r' m
   . (Manifest r e, Load r' ix e, MonadIO m)
  => Array r' ix e
  -> m (Array r ix e)
computeIO :: forall r ix e r' (m :: * -> *).
(Manifest r e, Load r' ix e, MonadIO m) =>
Array r' ix e -> m (Array r ix e)
computeIO Array r' ix e
arr = IO (Array r ix e) -> m (Array r ix e)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Array r' ix e -> IO (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 IO (MArray RealWorld r ix e)
-> (MArray RealWorld r ix e -> IO (Array r ix e))
-> IO (Array r ix e)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m 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)
forall ix (m :: * -> *).
(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
forall ix e. Array r' ix e -> Comp
getComp Array r' ix e
arr))
{-# INLINE computeIO #-}

-- | Compute an array in `PrimMonad` sequentially disregarding predefined computation
-- strategy.
--
-- @since 0.4.5
computePrimM
  :: forall r ix e r' m
   . (Manifest r e, Load r' ix e, PrimMonad m)
  => Array r' ix e
  -> m (Array r ix e)
computePrimM :: forall r ix e r' (m :: * -> *).
(Manifest r e, Load r' ix e, PrimMonad m) =>
Array r' ix e -> m (Array r ix e)
computePrimM Array r' ix e
arr = 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 m (MArray (PrimState m) r ix e)
-> (MArray (PrimState m) r ix e -> m (Array r ix e))
-> m (Array r ix e)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m 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)
forall ix (m :: * -> *).
(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
forall ix e. Array r' ix e -> Comp
getComp Array r' ix e
arr)
{-# INLINE computePrimM #-}

-- | Just as `compute`, but let's you supply resulting representation type as an argument.
--
-- ====__Examples__
--
-- >>> import Data.Massiv.Array
-- >>> computeAs P $ range Seq (Ix1 0) 10
-- Array P Seq (Sz1 10)
--   [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 ]
computeAs :: (Manifest r e, Load r' ix e) => r -> Array r' ix e -> Array r ix e
computeAs :: forall r e r' ix.
(Manifest r e, Load r' ix e) =>
r -> Array r' ix e -> Array r ix e
computeAs r
_ = Array r' ix e -> Array r ix e
forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute
{-# INLINE computeAs #-}

-- | Same as `compute` and `computeAs`, but let's you supply resulting representation type as a proxy
-- argument.
--
-- ==== __Examples__
--
-- Useful only really for cases when representation constructor or @TypeApplications@ extension
-- aren't desireable for some reason:
--
-- >>> import Data.Proxy
-- >>> import Data.Massiv.Array
-- >>> computeProxy (Proxy :: Proxy P) $ (^ (2 :: Int)) <$> range Seq (Ix1 0) 10
-- Array P Seq (Sz1 10)
--   [ 0, 1, 4, 9, 16, 25, 36, 49, 64, 81 ]
--
-- @since 0.1.1
computeProxy :: (Manifest r e, Load r' ix e) => proxy r -> Array r' ix e -> Array r ix e
computeProxy :: forall r e r' ix (proxy :: * -> *).
(Manifest r e, Load r' ix e) =>
proxy r -> Array r' ix e -> Array r ix e
computeProxy proxy r
_ = Array r' ix e -> Array r ix e
forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute
{-# INLINE computeProxy #-}

-- | This is just like `convert`, but restricted to `Source` arrays. Will be a noop if
-- resulting type is the same as the input.
--
-- @since 0.1.0
computeSource
  :: forall r ix e r'
   . (Manifest r e, Source r' e, Index ix)
  => Array r' ix e
  -> Array r ix e
computeSource :: forall r ix e r'.
(Manifest r e, Source r' e, Index ix) =>
Array r' ix e -> Array r ix e
computeSource Array r' ix e
arr = Array r ix e
-> ((r' :~: r) -> Array r ix e) -> Maybe (r' :~: r) -> Array r ix e
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Array D ix e -> Array r ix e
forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute (Array D ix e -> Array r ix e) -> Array D ix e -> Array r ix e
forall a b. (a -> b) -> a -> b
$ Array r' ix e -> Array D ix e
forall ix r e.
(Index ix, Source r e) =>
Array r ix e -> Array D ix e
delay Array r' ix e
arr) (\r' :~: r
Refl -> Array r ix e
Array r' ix e
arr) (Maybe (r' :~: r)
forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT :: Maybe (r' :~: r))
{-# INLINE computeSource #-}

-- | /O(n)/ - Make an exact immutable copy of an Array.
--
-- @since 0.1.0
clone :: (Manifest r e, Index ix) => Array r ix e -> Array r ix e
clone :: forall r e ix.
(Manifest r e, Index ix) =>
Array r ix e -> Array r ix e
clone Array r ix e
arr = IO (Array r ix e) -> Array r ix e
forall a. IO a -> a
unsafePerformIO (IO (Array r ix e) -> Array r ix e)
-> IO (Array r ix e) -> Array r ix e
forall a b. (a -> b) -> a -> b
$ Array r ix e -> IO (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 IO (MArray RealWorld r ix e)
-> (MArray RealWorld r ix e -> IO (Array r ix e))
-> IO (Array r ix e)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m 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)
forall ix (m :: * -> *).
(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
forall ix e. Array r ix e -> Comp
getComp Array r ix e
arr)
{-# INLINE clone #-}

-- | /O(1)/ - Cast over Array representation
gcastArr
  :: forall r ix e r'
   . (Typeable r, Typeable r')
  => Array r' ix e
  -> Maybe (Array r ix e)
gcastArr :: forall r ix e r'.
(Typeable r, Typeable r') =>
Array r' ix e -> Maybe (Array r ix e)
gcastArr Array r' ix e
arr = ((r :~: r') -> Array r ix e)
-> Maybe (r :~: r') -> Maybe (Array r ix e)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\r :~: r'
Refl -> Array r ix e
Array r' ix e
arr) (Maybe (r :~: r')
forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT :: Maybe (r :~: r'))

-- | /O(n)/ - conversion between array types. A full copy will occur, unless when the source and
-- result arrays are of the same representation, in which case it is an /O(1)/ operation.
--
-- @since 0.1.0
convert
  :: forall r ix e r'
   . (Manifest r e, Load r' ix e)
  => Array r' ix e
  -> Array r ix e
convert :: forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
convert Array r' ix e
arr = Array r ix e -> Maybe (Array r ix e) -> Array r ix e
forall a. a -> Maybe a -> a
fromMaybe (Array r' ix e -> Array r ix e
forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute Array r' ix e
arr) (Array r' ix e -> Maybe (Array r ix e)
forall r ix e r'.
(Typeable r, Typeable r') =>
Array r' ix e -> Maybe (Array r ix e)
gcastArr Array r' ix e
arr)
{-# INLINE convert #-}

-- | Same as `convert`, but let's you supply resulting representation type as an argument.
--
-- @since 0.1.0
convertAs
  :: (Manifest r e, Load r' ix e)
  => r
  -> Array r' ix e
  -> Array r ix e
convertAs :: forall r e r' ix.
(Manifest r e, Load r' ix e) =>
r -> Array r' ix e -> Array r ix e
convertAs r
_ = Array r' ix e -> Array r ix e
forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
convert
{-# INLINE convertAs #-}

-- | Same as `convert` and `convertAs`, but let's you supply resulting representation type as a
-- proxy argument.
--
-- @since 0.1.1
convertProxy
  :: (Manifest r e, Load r' ix e)
  => proxy r
  -> Array r' ix e
  -> Array r ix e
convertProxy :: forall r e r' ix (proxy :: * -> *).
(Manifest r e, Load r' ix e) =>
proxy r -> Array r' ix e -> Array r ix e
convertProxy proxy r
_ = Array r' ix e -> Array r ix e
forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
convert
{-# INLINE convertProxy #-}

-- | Convert a ragged array into a common array with rectangular shape. Throws `ShapeException`
-- whenever supplied ragged array does not have a rectangular shape.
--
-- @since 0.4.0
fromRaggedArrayM
  :: forall r ix e r' m
   . (Manifest r e, Ragged r' ix e, MonadThrow m)
  => Array r' ix e
  -> m (Array r ix e)
fromRaggedArrayM :: forall r ix e r' (m :: * -> *).
(Manifest r e, Ragged r' ix e, MonadThrow m) =>
Array r' ix e -> m (Array r ix e)
fromRaggedArrayM Array r' ix e
arr =
  let sz :: Sz ix
sz = Array r' ix e -> Sz ix
forall e. Array r' ix e -> Sz ix
forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array r' ix e
arr
   in (ShapeException -> m (Array r ix e))
-> (Array r ix e -> m (Array r ix e))
-> Either ShapeException (Array r ix e)
-> m (Array r ix e)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (\(ShapeException
e :: ShapeException) -> ShapeException -> m (Array r ix e)
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM ShapeException
e) Array r ix e -> m (Array r ix e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either ShapeException (Array r ix e) -> m (Array r ix e))
-> Either ShapeException (Array r ix e) -> m (Array r ix e)
forall a b. (a -> b) -> a -> b
$
        IO (Either ShapeException (Array r ix e))
-> Either ShapeException (Array r ix e)
forall a. IO a -> a
unsafePerformIO (IO (Either ShapeException (Array r ix e))
 -> Either ShapeException (Array r ix e))
-> IO (Either ShapeException (Array r ix e))
-> Either ShapeException (Array r ix e)
forall a b. (a -> b) -> a -> b
$ do
          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)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz
          (() -> IO (Array r ix e))
-> Either ShapeException ()
-> IO (Either ShapeException (Array r ix e))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b)
-> Either ShapeException a -> f (Either ShapeException b)
traverse (\()
_ -> 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)
forall ix (m :: * -> *).
(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
forall ix e. Array r' ix e -> Comp
getComp Array r' ix e
arr) MArray RealWorld r ix e
MArray (PrimState IO) r ix e
marr)
            (Either ShapeException ()
 -> IO (Either ShapeException (Array r ix e)))
-> IO (Either ShapeException ())
-> IO (Either ShapeException (Array r ix e))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO () -> IO (Either ShapeException ())
forall e a. Exception e => IO a -> IO (Either e a)
try
              ( Comp -> (Scheduler RealWorld () -> IO ()) -> IO ()
withMassivScheduler_ (Array r' ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. 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 RealWorld () -> IO ()
forall a. ST RealWorld a -> IO a
stToIO (ST RealWorld () -> IO ()) -> ST RealWorld () -> IO ()
forall a b. (a -> b) -> a -> b
$ Scheduler RealWorld ()
-> Array r' ix e
-> (Ix1 -> e -> ST RealWorld ())
-> Ix1
-> Ix1
-> Sz ix
-> ST RealWorld ()
forall s.
Scheduler s ()
-> Array r' ix e
-> (Ix1 -> e -> ST s ())
-> Ix1
-> Ix1
-> Sz ix
-> ST s ()
forall r ix e s.
Ragged r ix e =>
Scheduler s ()
-> Array r ix e
-> (Ix1 -> e -> ST s ())
-> Ix1
-> Ix1
-> Sz ix
-> ST s ()
loadRaggedST 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 ()
forall ix (m :: * -> *).
(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) Ix1
0 (Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz) Sz ix
sz
              )
{-# INLINE fromRaggedArrayM #-}

-- | Same as `fromRaggedArrayM`, but will throw an impure exception if its shape is not
-- rectangular.
--
-- @since 0.1.1
fromRaggedArray'
  :: forall r ix e r'
   . (HasCallStack, Manifest r e, Ragged r' ix e)
  => Array r' ix e
  -> Array r ix e
fromRaggedArray' :: forall r ix e r'.
(HasCallStack, Manifest r e, Ragged r' ix e) =>
Array r' ix e -> Array r ix e
fromRaggedArray' = Either SomeException (Array r ix e) -> Array r ix e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException (Array r ix e) -> Array r ix e)
-> (Array r' ix e -> Either SomeException (Array r ix e))
-> Array r' ix e
-> Array r ix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r' ix e -> Either SomeException (Array r ix e)
forall r ix e r' (m :: * -> *).
(Manifest r e, Ragged r' ix e, MonadThrow m) =>
Array r' ix e -> m (Array r ix e)
fromRaggedArrayM
{-# INLINE fromRaggedArray' #-}

-- | Same as `compute`, but with `Stride`.
--
-- /O(n div k)/ - Where @n@ is number of elements in the source array and @k@ is number of
-- elements in the stride.
--
-- @since 0.3.0
computeWithStride
  :: forall r ix e r'
   . (Manifest r e, StrideLoad r' ix e)
  => Stride ix
  -> Array r' ix e
  -> Array r ix e
computeWithStride :: forall r ix e r'.
(Manifest r e, StrideLoad r' ix e) =>
Stride ix -> Array r' ix e -> Array r ix e
computeWithStride Stride ix
stride !Array r' ix e
arr =
  IO (Array r ix e) -> Array r ix e
forall a. IO a -> a
unsafePerformIO (IO (Array r ix e) -> Array r ix e)
-> IO (Array r ix e) -> Array r ix e
forall a b. (a -> b) -> a -> b
$ do
    let !sz :: Sz ix
sz = Stride ix -> Sz ix -> Sz ix
forall ix. Index ix => Stride ix -> Sz ix -> Sz ix
strideSize Stride ix
stride (Array r' ix e -> Sz ix
forall e. Array r' ix e -> Sz ix
forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array r' ix e
arr)
    Comp
-> Sz ix
-> (Scheduler RealWorld () -> MArray RealWorld r ix e -> IO ())
-> IO (Array r ix e)
forall r ix e a (m :: * -> *) b.
(Manifest r e, Index ix, MonadUnliftIO m) =>
Comp
-> Sz ix
-> (Scheduler RealWorld a -> MArray RealWorld r ix e -> m b)
-> m (Array r ix e)
unsafeCreateArray_ (Array r' ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r' ix e -> Comp
getComp Array r' ix e
arr) Sz ix
sz ((Scheduler RealWorld () -> MArray RealWorld r ix e -> IO ())
 -> IO (Array r ix e))
-> (Scheduler RealWorld () -> MArray RealWorld r ix e -> IO ())
-> IO (Array r ix e)
forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld ()
scheduler MArray RealWorld r ix e
marr ->
      ST RealWorld () -> IO ()
forall a. ST RealWorld a -> IO a
stToIO (ST RealWorld () -> IO ()) -> ST RealWorld () -> IO ()
forall a b. (a -> b) -> a -> b
$ Scheduler RealWorld ()
-> Stride ix
-> Sz ix
-> Array r' ix e
-> (Ix1 -> e -> ST RealWorld ())
-> ST RealWorld ()
forall s.
Scheduler s ()
-> Stride ix
-> Sz ix
-> Array r' ix e
-> (Ix1 -> e -> ST s ())
-> ST s ()
forall r ix e s.
StrideLoad r ix e =>
Scheduler s ()
-> Stride ix
-> Sz ix
-> Array r ix e
-> (Ix1 -> e -> ST s ())
-> ST s ()
iterArrayLinearWithStrideST_ Scheduler RealWorld ()
scheduler Stride ix
stride Sz ix
sz 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 ()
forall ix (m :: * -> *).
(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 computeWithStride #-}

-- | Same as `computeWithStride`, but with ability to specify resulting array representation.
--
-- @since 0.3.0
computeWithStrideAs
  :: (Manifest r e, StrideLoad r' ix e) => r -> Stride ix -> Array r' ix e -> Array r ix e
computeWithStrideAs :: forall r e r' ix.
(Manifest r e, StrideLoad r' ix e) =>
r -> Stride ix -> Array r' ix e -> Array r ix e
computeWithStrideAs r
_ = Stride ix -> Array r' ix e -> Array r ix e
forall r ix e r'.
(Manifest r e, StrideLoad r' ix e) =>
Stride ix -> Array r' ix e -> Array r ix e
computeWithStride
{-# INLINE computeWithStrideAs #-}

-- | Load into a supplied mutable vector sequentially. Returned array is not
-- necesserally the same vector as the one that was supplied. It will be the
-- same only if it had enough space to load all the elements in.
--
-- @since 0.5.7
unsafeLoadIntoS
  :: forall r r' ix e m s
   . (Load r ix e, Manifest r' e, MonadPrim s m)
  => MVector s r' e
  -> Array r ix e
  -> m (MArray s r' ix e)
unsafeLoadIntoS :: forall r r' ix e (m :: * -> *) s.
(Load r ix e, Manifest r' e, MonadPrim s m) =>
MVector s r' e -> Array r ix e -> m (MArray s r' ix e)
unsafeLoadIntoS MVector s r' e
marr Array r ix e
arr = ST (PrimState m) (MArray s r' ix e) -> m (MArray s r' ix e)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (MArray s r' ix e) -> m (MArray s r' ix e))
-> ST (PrimState m) (MArray s r' ix e) -> m (MArray s r' ix e)
forall a b. (a -> b) -> a -> b
$ MVector s r' e -> Array r ix e -> ST s (MArray s r' ix e)
forall r r' ix e (m :: * -> *) s.
(Load r ix e, Manifest r' e, MonadPrim s m) =>
MVector s r' e -> Array r ix e -> m (MArray s r' ix e)
unsafeLoadIntoS MVector s r' e
marr Array r ix e
arr
{-# INLINE unsafeLoadIntoS #-}

-- | Same as `unsafeLoadIntoS`, but respecting computation strategy.
--
-- @since 0.5.7
unsafeLoadIntoM
  :: forall r r' ix e m
   . (Load r ix e, Manifest r' e, MonadIO m)
  => MVector RealWorld r' e
  -> Array r ix e
  -> m (MArray RealWorld r' ix e)
unsafeLoadIntoM :: forall r r' ix e (m :: * -> *).
(Load r ix e, Manifest r' e, MonadIO m) =>
MVector RealWorld r' e
-> Array r ix e -> m (MArray RealWorld r' ix e)
unsafeLoadIntoM MVector RealWorld r' e
marr Array r ix e
arr = IO (MArray RealWorld r' ix e) -> m (MArray RealWorld r' ix e)
forall a. IO a -> m a
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
$ MVector RealWorld r' e
-> Array r ix e -> IO (MArray RealWorld r' ix e)
forall r'.
Manifest r' e =>
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 unsafeLoadIntoM #-}

-- | Efficiently iterate a function until a convergence condition is satisfied. If the
-- size of array doesn't change between iterations then no more than two new arrays will be
-- allocated, regardless of the number of iterations. If the size does change from one
-- iteration to another, an attempt will be made to grow/shrink the intermediate mutable
-- array instead of allocating a new one.
--
-- ====__Example__
--
-- >>> import Data.Massiv.Array
-- >>> let arr = computeAs P $ makeLoadArrayS (Sz2 8 8) (0 :: Int) $ \ w -> () <$ w (0 :. 0) 1
-- >>> arr
-- Array P Seq (Sz (8 :. 8))
--   [ [ 1, 0, 0, 0, 0, 0, 0, 0 ]
--   , [ 0, 0, 0, 0, 0, 0, 0, 0 ]
--   , [ 0, 0, 0, 0, 0, 0, 0, 0 ]
--   , [ 0, 0, 0, 0, 0, 0, 0, 0 ]
--   , [ 0, 0, 0, 0, 0, 0, 0, 0 ]
--   , [ 0, 0, 0, 0, 0, 0, 0, 0 ]
--   , [ 0, 0, 0, 0, 0, 0, 0, 0 ]
--   , [ 0, 0, 0, 0, 0, 0, 0, 0 ]
--   ]
-- >>> let nextPascalRow cur above = if cur == 0 then above else cur
-- >>> let pascal = makeStencil (Sz2 2 2) 1 $ \ get -> nextPascalRow (get (0 :. 0)) (get (-1 :. -1) + get (-1 :. 0))
-- >>> iterateUntil (\_ _ a -> (a ! (7 :. 7)) /= 0) (\ _ -> mapStencil (Fill 0) pascal) arr
-- Array P Seq (Sz (8 :. 8))
--   [ [ 1, 0, 0, 0, 0, 0, 0, 0 ]
--   , [ 1, 1, 0, 0, 0, 0, 0, 0 ]
--   , [ 1, 2, 1, 0, 0, 0, 0, 0 ]
--   , [ 1, 3, 3, 1, 0, 0, 0, 0 ]
--   , [ 1, 4, 6, 4, 1, 0, 0, 0 ]
--   , [ 1, 5, 10, 10, 5, 1, 0, 0 ]
--   , [ 1, 6, 15, 20, 15, 6, 1, 0 ]
--   , [ 1, 7, 21, 35, 35, 21, 7, 1 ]
--   ]
--
-- @since 0.3.6
iterateUntil
  :: (Load r' ix e, Manifest r e, NFData (Array r ix e))
  => (Int -> Array r ix e -> Array r ix e -> Bool)
  -- ^ Convergence condition. Accepts current iteration counter, array at the previous
  -- state and at the current state.
  -> (Int -> Array r ix e -> Array r' ix e)
  -- ^ A modifying function to apply at each iteration. The size of resulting array may
  -- differ if necessary
  -> Array r ix e
  -- ^ Initial source array
  -> Array r ix e
iterateUntil :: forall r' ix e r.
(Load r' ix e, Manifest r e, NFData (Array r ix e)) =>
(Ix1 -> Array r ix e -> Array r ix e -> Bool)
-> (Ix1 -> Array r ix e -> Array r' ix e)
-> Array r ix e
-> Array r ix e
iterateUntil Ix1 -> Array r ix e -> Array r ix e -> Bool
convergence Ix1 -> Array r ix e -> Array r' ix e
iteration Array r ix e
initArr0 = IO (Array r ix e) -> Array r ix e
forall a. IO a -> a
unsafePerformIO (IO (Array r ix e) -> Array r ix e)
-> IO (Array r ix e) -> Array r ix e
forall a b. (a -> b) -> a -> b
$ do
  let loadArr0 :: Array r' ix e
loadArr0 = Ix1 -> Array r ix e -> Array r' ix e
iteration Ix1
0 Array r ix e
initArr0
  MVector RealWorld r e
initMVec1 <- Sz Ix1 -> IO (MArray (PrimState IO) r Ix1 e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
forall ix (m :: * -> *).
(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 e. 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
loadArr0))
  let conv :: Ix1
-> Array r ix e
-> Comp
-> MArray RealWorld r ix e
-> IO (Bool, Array r ix e)
conv Ix1
n Array r ix e
arr Comp
comp MArray RealWorld r ix e
marr' = do
        Array r ix e
arr' <- 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)
forall ix (m :: * -> *).
(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'
        Array r ix e
arr' Array r ix e -> IO (Bool, Array r ix e) -> IO (Bool, Array r ix e)
forall a b. NFData a => a -> b -> b
`deepseq` (Bool, Array r ix e) -> IO (Bool, Array r ix e)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ix1 -> Array r ix e -> Array r ix e -> Bool
convergence Ix1
n Array r ix e
arr Array r ix e
arr', Array r ix e
arr')
  (Ix1
 -> Array r ix e
 -> Comp
 -> MArray RealWorld r ix e
 -> IO (Bool, Array r ix e))
-> (Ix1 -> Array r ix e -> IO (Array r' ix e))
-> Ix1
-> Array r ix e
-> Array r' ix e
-> MVector RealWorld r e
-> IO (Array r ix e)
forall r' ix e r (m :: * -> *).
(Load r' ix e, Manifest r e, MonadIO m) =>
(Ix1
 -> Array r ix e
 -> Comp
 -> MArray RealWorld r ix e
 -> m (Bool, Array r ix e))
-> (Ix1 -> Array r ix e -> m (Array r' ix e))
-> Ix1
-> Array r ix e
-> Array r' ix e
-> MVector RealWorld r e
-> m (Array r ix e)
iterateLoop Ix1
-> Array r ix e
-> Comp
-> MArray RealWorld r ix e
-> IO (Bool, Array r ix e)
conv (\Ix1
n -> Array r' ix e -> IO (Array r' ix e)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array r' ix e -> IO (Array r' ix e))
-> (Array r ix e -> Array r' ix e)
-> Array r ix e
-> IO (Array r' ix e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ix1 -> Array r ix e -> Array r' ix e
iteration Ix1
n) Ix1
0 Array r ix e
initArr0 Array r' ix e
loadArr0 MVector RealWorld r e
initMVec1
{-# INLINE iterateUntil #-}

-- | Monadic version of `iterateUntil` where at each iteration mutable version
-- of an array is available. However it is less efficient then the pure
-- alternative, because an intermediate array must be copied at each
-- iteration.
--
-- @since 0.3.6
iterateUntilM
  :: (Load r' ix e, Manifest r e, MonadIO m)
  => (Int -> Array r ix e -> MArray RealWorld r ix e -> m Bool)
  -- ^ Convergence condition. Accepts current iteration counter, pure array at previous
  -- state and a mutable at the current state, therefore after each iteration its contents
  -- can be modifed if necessary.
  -> (Int -> Array r ix e -> m (Array r' ix e))
  -- ^ A modifying function to apply at each iteration.  The size of resulting array may
  -- differ if necessary.
  -> Array r ix e
  -- ^ Initial source array
  -> m (Array r ix e)
iterateUntilM :: forall r' ix e r (m :: * -> *).
(Load r' ix e, Manifest r e, MonadIO m) =>
(Ix1 -> Array r ix e -> MArray RealWorld r ix e -> m Bool)
-> (Ix1 -> Array r ix e -> m (Array r' ix e))
-> Array r ix e
-> m (Array r ix e)
iterateUntilM Ix1 -> Array r ix e -> MArray RealWorld r ix e -> m Bool
convergence Ix1 -> Array r ix e -> m (Array r' ix e)
iteration Array r ix e
initArr0 = do
  Array r' ix e
loadArr0 <- Ix1 -> Array r ix e -> m (Array r' ix e)
iteration Ix1
0 Array r ix e
initArr0
  MVector RealWorld r e
initMVec1 <- IO (MVector RealWorld r e) -> m (MVector RealWorld r e)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MVector RealWorld r e) -> m (MVector RealWorld r e))
-> IO (MVector RealWorld r e) -> m (MVector RealWorld r e)
forall a b. (a -> b) -> a -> b
$ Sz Ix1 -> IO (MArray (PrimState IO) r Ix1 e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
forall ix (m :: * -> *).
(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 e. 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
loadArr0))
  let conv :: Ix1
-> Array r ix e
-> Comp
-> MArray RealWorld r ix e
-> m (Bool, Array r ix e)
conv Ix1
n Array r ix e
arr Comp
comp MArray RealWorld r ix e
marr = (,) (Bool -> Array r ix e -> (Bool, Array r ix e))
-> m Bool -> m (Array r ix e -> (Bool, Array r ix e))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ix1 -> Array r ix e -> MArray RealWorld r ix e -> m Bool
convergence Ix1
n Array r ix e
arr MArray RealWorld r ix e
marr m (Array r ix e -> (Bool, Array r ix e))
-> m (Array r ix e) -> m (Bool, Array r ix e)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Comp -> MArray RealWorld r ix e -> m (Array r ix e)
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
comp MArray RealWorld r ix e
marr
  (Ix1
 -> Array r ix e
 -> Comp
 -> MArray RealWorld r ix e
 -> m (Bool, Array r ix e))
-> (Ix1 -> Array r ix e -> m (Array r' ix e))
-> Ix1
-> Array r ix e
-> Array r' ix e
-> MVector RealWorld r e
-> m (Array r ix e)
forall r' ix e r (m :: * -> *).
(Load r' ix e, Manifest r e, MonadIO m) =>
(Ix1
 -> Array r ix e
 -> Comp
 -> MArray RealWorld r ix e
 -> m (Bool, Array r ix e))
-> (Ix1 -> Array r ix e -> m (Array r' ix e))
-> Ix1
-> Array r ix e
-> Array r' ix e
-> MVector RealWorld r e
-> m (Array r ix e)
iterateLoop Ix1
-> Array r ix e
-> Comp
-> MArray RealWorld r ix e
-> m (Bool, Array r ix e)
conv Ix1 -> Array r ix e -> m (Array r' ix e)
iteration Ix1
0 Array r ix e
initArr0 Array r' ix e
loadArr0 MVector RealWorld r e
initMVec1
{-# INLINE iterateUntilM #-}

iterateLoop
  :: (Load r' ix e, Manifest r e, MonadIO m)
  => (Int -> Array r ix e -> Comp -> MArray RealWorld r ix e -> m (Bool, Array r ix e))
  -> (Int -> Array r ix e -> m (Array r' ix e))
  -> Int
  -> Array r ix e
  -> Array r' ix e
  -> MVector RealWorld r e
  -> m (Array r ix e)
iterateLoop :: forall r' ix e r (m :: * -> *).
(Load r' ix e, Manifest r e, MonadIO m) =>
(Ix1
 -> Array r ix e
 -> Comp
 -> MArray RealWorld r ix e
 -> m (Bool, Array r ix e))
-> (Ix1 -> Array r ix e -> m (Array r' ix e))
-> Ix1
-> Array r ix e
-> Array r' ix e
-> MVector RealWorld r e
-> m (Array r ix e)
iterateLoop Ix1
-> Array r ix e
-> Comp
-> MArray RealWorld r ix e
-> m (Bool, Array r ix e)
convergence Ix1 -> Array r ix e -> m (Array r' ix e)
iteration = Ix1
-> Array r ix e
-> Array r' ix e
-> MVector RealWorld r e
-> m (Array r ix e)
go
  where
    go :: Ix1
-> Array r ix e
-> Array r' ix e
-> MVector RealWorld r e
-> m (Array r ix e)
go Ix1
n !Array r ix e
arr !Array r' ix e
loadArr !MVector RealWorld r e
mvec = do
      let !comp :: Comp
comp = Array r' ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r' ix e -> Comp
getComp Array r' ix e
loadArr
      MArray RealWorld r ix e
marr' <- MVector RealWorld r e
-> Array r' ix e -> m (MArray RealWorld r ix e)
forall r r' ix e (m :: * -> *).
(Load r ix e, Manifest r' e, MonadIO m) =>
MVector RealWorld r' e
-> Array r ix e -> m (MArray RealWorld r' ix e)
unsafeLoadIntoM MVector RealWorld r e
mvec Array r' ix e
loadArr
      (Bool
shouldStop, Array r ix e
arr') <- Ix1
-> Array r ix e
-> Comp
-> MArray RealWorld r ix e
-> m (Bool, Array r ix e)
convergence Ix1
n Array r ix e
arr Comp
comp MArray RealWorld r ix e
marr'
      if Bool
shouldStop
        then Array r ix e -> m (Array r ix e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Array r ix e
arr'
        else do
          MArray RealWorld r ix e
nextMArr <- IO (MArray RealWorld r ix e) -> m (MArray RealWorld r ix e)
forall a. IO a -> m a
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
$ Array r ix e -> IO (MArray (PrimState IO) r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Array r ix e -> m (MArray (PrimState m) r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Array r ix e -> m (MArray (PrimState m) r ix e)
unsafeThaw Array r ix e
arr
          Array r' ix e
arr'' <- Ix1 -> Array r ix e -> m (Array r' ix e)
iteration (Ix1
n Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) Array r ix e
arr'
          Ix1
-> Array r ix e
-> Array r' ix e
-> MVector RealWorld r e
-> m (Array r ix e)
go (Ix1
n Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) Array r ix e
arr' Array r' ix e
arr'' (MVector RealWorld r e -> m (Array r ix e))
-> MVector RealWorld r e -> m (Array r ix e)
forall a b. (a -> b) -> a -> b
$ MArray RealWorld r ix e -> MVector RealWorld r e
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> MVector s r e
flattenMArray MArray RealWorld r ix e
nextMArr
{-# INLINE iterateLoop #-}