{-# 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 (
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)
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 = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ 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 #-}
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 a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ 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 #-}
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 = forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
setComp (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix e
arr) forall a b. (a -> b) -> a -> b
$ 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. Strategy r => Comp -> Array r ix e -> Array r ix e
setComp Comp
Par Array r' ix e
arr)
{-# INLINE computeP #-}
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 = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (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 forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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 (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix e
arr))
{-# INLINE computeIO #-}
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 = 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 forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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 (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix e
arr)
{-# INLINE computePrimM #-}
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
_ = forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute
{-# INLINE computeAs #-}
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
_ = forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute
{-# INLINE computeProxy #-}
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 = forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute forall a b. (a -> b) -> a -> b
$ 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
arr) (forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT :: Maybe (r' :~: r))
{-# INLINE computeSource #-}
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 = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ 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 forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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 (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr)
{-# INLINE clone #-}
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 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\r :~: r'
Refl -> Array r' ix e
arr) (forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT :: Maybe (r :~: r'))
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 = forall a. a -> Maybe a -> a
fromMaybe (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) (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 #-}
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
_ = forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
convert
{-# INLINE convertAs #-}
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
_ = forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
convert
{-# INLINE convertProxy #-}
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 = forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array r' ix e
arr
in forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (\(ShapeException
e :: ShapeException) -> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM ShapeException
e) forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ do
MArray RealWorld r ix e
marr <- 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
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\()
_ -> 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 (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix e
arr) MArray RealWorld r ix e
marr)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall e a. Exception e => IO a -> IO (Either e a)
try
( Comp -> (Scheduler RealWorld () -> IO ()) -> IO ()
withMassivScheduler_ (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix e
arr) forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld ()
scheduler ->
forall a. ST RealWorld a -> IO a
stToIO forall a b. (a -> b) -> a -> b
$ 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 (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
marr) Ix1
0 (forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz) Sz ix
sz
)
{-# INLINE fromRaggedArrayM #-}
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' = forall a. HasCallStack => Either SomeException a -> a
throwEither forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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' #-}
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 =
forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ do
let !sz :: Sz ix
sz = forall ix. Index ix => Stride ix -> Sz ix -> Sz ix
strideSize Stride ix
stride (forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array r' ix e
arr)
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_ (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix e
arr) Sz ix
sz forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld ()
scheduler MArray RealWorld r ix e
marr ->
forall a. ST RealWorld a -> IO a
stToIO forall a b. (a -> b) -> a -> b
$ 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 (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
marr)
{-# INLINE computeWithStride #-}
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
_ = 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 #-}
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 = forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim forall a b. (a -> b) -> a -> b
$ 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 #-}
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 = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ 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 #-}
iterateUntil
:: (Load r' ix e, Manifest r e, NFData (Array r ix e))
=> (Int -> Array r ix e -> Array r ix e -> Bool)
-> (Int -> Array r ix e -> Array r' ix e)
-> Array r ix e
-> 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 = forall a. IO a -> a
unsafePerformIO 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 <- forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew (forall a. a -> Maybe a -> a
fromMaybe forall ix. Index ix => Sz ix
zeroSz (forall r ix e. Shape r ix => Array r ix e -> Maybe Sz1
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' <- 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
marr'
Array r ix e
arr' forall a b. NFData a => a -> b -> b
`deepseq` 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')
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 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure 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 #-}
iterateUntilM
:: (Load r' ix e, Manifest r e, MonadIO m)
=> (Int -> Array r ix e -> MArray RealWorld r ix e -> m Bool)
-> (Int -> Array r ix e -> m (Array r' ix e))
-> Array r ix e
-> 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 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew (forall a. a -> Maybe a -> a
fromMaybe forall ix. Index ix => Sz ix
zeroSz (forall r ix e. Shape r ix => Array r ix e -> Maybe Sz1
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 = (,) 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 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> 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
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 = forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r' ix e
loadArr
MArray RealWorld r ix e
marr' <- 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 forall (f :: * -> *) a. Applicative f => a -> f a
pure Array r ix e
arr'
else do
MArray RealWorld r ix e
nextMArr <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall r e ix (m :: * -> *).
(Manifest r e, 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 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 forall a. Num a => a -> a -> a
+ Ix1
1) Array r ix e
arr' Array r' ix e
arr'' forall a b. (a -> b) -> a -> b
$ 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 #-}