{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
module Data.Massiv.Array.Manifest.Internal
( M
, Manifest(..)
, Array(..)
, toManifest
, compute
, computeS
, computeP
, computeIO
, computePrimM
, computeAs
, computeProxy
, computeSource
, computeWithStride
, computeWithStrideAs
, clone
, convert
, convertAs
, convertProxy
, gcastArr
, fromRaggedArrayM
, fromRaggedArray'
, sizeofArray
, sizeofMutableArray
, iterateUntil
, iterateUntilM
) where
import Control.Exception (try)
import Control.Monad.ST
import Control.Scheduler
import qualified Data.Foldable as F (Foldable(..))
import Data.Massiv.Array.Delayed.Pull
import Data.Massiv.Array.Mutable
import Data.Massiv.Array.Ops.Fold.Internal as A
import Data.Massiv.Array.Mutable.Internal (unsafeCreateArray_)
import Data.Massiv.Vector.Stream as S (steps, isteps)
import Data.Massiv.Core.Common
import Data.Massiv.Core.List
import Data.Maybe (fromMaybe)
import Data.Typeable
import GHC.Base hiding (ord)
import System.IO.Unsafe (unsafePerformIO)
#if MIN_VERSION_primitive(0,6,2)
import Data.Primitive.Array (sizeofArray, sizeofMutableArray)
#else
import qualified Data.Primitive.Array as A (Array(..), MutableArray(..))
import GHC.Exts (sizeofArray#, sizeofMutableArray#)
sizeofArray :: A.Array a -> Int
sizeofArray (A.Array a) = I# (sizeofArray# a)
{-# INLINE sizeofArray #-}
sizeofMutableArray :: A.MutableArray s a -> Int
sizeofMutableArray (A.MutableArray ma) = I# (sizeofMutableArray# ma)
{-# INLINE sizeofMutableArray #-}
#endif
data M
data instance Array M ix e = MArray { Array M ix e -> Comp
mComp :: !Comp
, Array M ix e -> Sz ix
mSize :: !(Sz ix)
, Array M ix e -> Int -> e
mLinearIndex :: Int -> e }
instance (Ragged L ix e, Show e) => Show (Array M ix e) where
showsPrec :: Int -> Array M ix e -> ShowS
showsPrec = (Array M ix e -> Array M ix e) -> Int -> Array M ix e -> ShowS
forall r r' ix ix' e.
(Ragged L ix' e, Load r ix e, Source r' ix' e, Show e) =>
(Array r ix e -> Array r' ix' e) -> Int -> Array r ix e -> ShowS
showsArrayPrec Array M ix e -> Array M ix e
forall a. a -> a
id
showList :: [Array M ix e] -> ShowS
showList = [Array M ix e] -> ShowS
forall arr. Show arr => [arr] -> ShowS
showArrayList
instance (Eq e, Index ix) => Eq (Array M ix e) where
== :: Array M ix e -> Array M ix e -> Bool
(==) = (e -> e -> Bool) -> Array M ix e -> Array M ix e -> Bool
forall r1 ix e1 r2 e2.
(Source r1 ix e1, Source r2 ix e2) =>
(e1 -> e2 -> Bool) -> Array r1 ix e1 -> Array r2 ix e2 -> Bool
eq e -> e -> Bool
forall a. Eq a => a -> a -> Bool
(==)
{-# INLINE (==) #-}
instance (Ord e, Index ix) => Ord (Array M ix e) where
compare :: Array M ix e -> Array M ix e -> Ordering
compare = (e -> e -> Ordering) -> Array M ix e -> Array M ix e -> Ordering
forall r1 ix e1 r2 e2.
(Source r1 ix e1, Source r2 ix e2) =>
(e1 -> e2 -> Ordering)
-> Array r1 ix e1 -> Array r2 ix e2 -> Ordering
ord e -> e -> Ordering
forall a. Ord a => a -> a -> Ordering
compare
{-# INLINE compare #-}
toManifest :: Manifest r ix e => Array r ix e -> Array M ix e
toManifest :: Array r ix e -> Array M ix e
toManifest !Array r ix e
arr = Comp -> Sz ix -> (Int -> e) -> Array M ix e
forall ix e. Comp -> Sz ix -> (Int -> e) -> Array M ix e
MArray (Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr) (Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr) (Array r ix e -> Int -> e
forall r ix e. Manifest r ix e => Array r ix e -> Int -> e
unsafeLinearIndexM Array r ix e
arr)
{-# INLINE toManifest #-}
instance Index ix => Foldable (Array M ix) where
fold :: Array M ix m -> m
fold = Array M ix m -> m
forall e r ix. (Monoid e, Source r ix e) => Array r ix e -> e
fold
{-# INLINE fold #-}
foldMap :: (a -> m) -> Array M ix a -> m
foldMap = (a -> m) -> Array M ix a -> m
forall r ix e m.
(Source r ix e, Monoid m) =>
(e -> m) -> Array r ix e -> m
foldMono
{-# INLINE foldMap #-}
foldl :: (b -> a -> b) -> b -> Array M ix a -> b
foldl = (b -> a -> b) -> b -> Array M ix a -> b
forall r ix e a.
Source r ix e =>
(a -> e -> a) -> a -> Array r ix e -> a
lazyFoldlS
{-# INLINE foldl #-}
foldl' :: (b -> a -> b) -> b -> Array M ix a -> b
foldl' = (b -> a -> b) -> b -> Array M ix a -> b
forall r ix e a.
Source r ix e =>
(a -> e -> a) -> a -> Array r ix e -> a
foldlS
{-# INLINE foldl' #-}
foldr :: (a -> b -> b) -> b -> Array M ix a -> b
foldr = (a -> b -> b) -> b -> Array M ix a -> b
forall r ix e b.
Source r ix e =>
(e -> b -> b) -> b -> Array r ix e -> b
foldrFB
{-# INLINE foldr #-}
foldr' :: (a -> b -> b) -> b -> Array M ix a -> b
foldr' = (a -> b -> b) -> b -> Array M ix a -> b
forall r ix e b.
Source r ix e =>
(e -> b -> b) -> b -> Array r ix e -> b
foldrS
{-# INLINE foldr' #-}
null :: Array M ix a -> Bool
null (MArray _ sz _) = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
{-# INLINE null #-}
length :: Array M ix a -> Int
length = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (Sz ix -> Int) -> (Array M ix a -> Sz ix) -> Array M ix a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array M ix a -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size
{-# INLINE length #-}
elem :: a -> Array M ix a -> Bool
elem a
e = (a -> Bool) -> Array M ix a -> Bool
forall r ix e. Source r ix e => (e -> Bool) -> Array r ix e -> Bool
A.any (a
e a -> a -> Bool
forall a. Eq a => a -> a -> Bool
==)
{-# INLINE elem #-}
toList :: Array M ix a -> [a]
toList Array M ix a
arr = (forall b. (a -> b -> b) -> b -> b) -> [a]
forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (\ a -> b -> b
c b
n -> (a -> b -> b) -> b -> Array M ix a -> b
forall r ix e b.
Source r ix e =>
(e -> b -> b) -> b -> Array r ix e -> b
foldrFB a -> b -> b
c b
n Array M ix a
arr)
{-# INLINE toList #-}
instance Index ix => Source M ix e where
unsafeLinearIndex :: Array M ix e -> Int -> e
unsafeLinearIndex = Array M ix e -> Int -> e
forall ix e. Array M ix e -> Int -> e
mLinearIndex
{-# INLINE unsafeLinearIndex #-}
unsafeLinearSlice :: Int -> Sz1 -> Array M ix e -> Array M Int e
unsafeLinearSlice Int
ix Sz1
sz Array M ix e
arr = Int -> Sz1 -> Array M Int e -> Array (R M) Int e
forall r ix e.
Extract r ix e =>
ix -> Sz ix -> Array r ix e -> Array (R r) ix e
unsafeExtract Int
ix Sz1
sz (Sz1 -> Array M ix e -> Array M Int e
forall r ix ix' e.
(Resize r ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
unsafeResize Sz1
sz Array M ix e
arr)
{-# INLINE unsafeLinearSlice #-}
instance Index ix => Manifest M ix e where
unsafeLinearIndexM :: Array M ix e -> Int -> e
unsafeLinearIndexM = Array M ix e -> Int -> e
forall ix e. Array M ix e -> Int -> e
mLinearIndex
{-# INLINE unsafeLinearIndexM #-}
instance Index ix => Resize M ix where
unsafeResize :: Sz ix' -> Array M ix e -> Array M ix' e
unsafeResize !Sz ix'
sz !Array M ix e
arr = Array M ix e
R:ArrayMixe ix e
arr { mSize :: Sz ix'
mSize = Sz ix'
sz }
{-# INLINE unsafeResize #-}
instance Index ix => Extract M ix e where
unsafeExtract :: ix -> Sz ix -> Array M ix e -> Array (R M) ix e
unsafeExtract !ix
sIx !Sz ix
newSz !Array M ix e
arr =
Comp -> Sz ix -> (Int -> e) -> Array M ix e
forall ix e. Comp -> Sz ix -> (Int -> e) -> Array M ix e
MArray (Array M ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array M ix e
arr) Sz ix
newSz ((Int -> e) -> Array M ix e) -> (Int -> e) -> Array M ix e
forall a b. (a -> b) -> a -> b
$ \ Int
i ->
Array M ix e -> ix -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array M ix e
arr ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) (Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
newSz Int
i) ix
sIx)
{-# INLINE unsafeExtract #-}
instance {-# OVERLAPPING #-} Slice M Ix1 e where
unsafeSlice :: Array M Int e -> Int -> Sz1 -> Dim -> m (Elt M Int e)
unsafeSlice Array M Int e
arr Int
i Sz1
_ Dim
_ = e -> m e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array M Int e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array M Int e
arr Int
i)
{-# INLINE unsafeSlice #-}
instance ( Index ix
, Index (Lower ix)
, Elt M ix e ~ Array M (Lower ix) e
) =>
Slice M ix e where
unsafeSlice :: Array M ix e -> ix -> Sz ix -> Dim -> m (Elt M ix e)
unsafeSlice Array M ix e
arr ix
start Sz ix
cutSz Dim
dim = do
(Sz1
_, Sz (Lower ix)
newSz) <- Sz ix -> Dim -> m (Sz1, Sz (Lower ix))
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz ix -> Dim -> m (Sz1, Sz (Lower ix))
pullOutSzM Sz ix
cutSz Dim
dim
Array M (Lower ix) e -> m (Array M (Lower ix) e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array M (Lower ix) e -> m (Array M (Lower ix) e))
-> Array M (Lower ix) e -> m (Array M (Lower ix) e)
forall a b. (a -> b) -> a -> b
$ Sz (Lower ix) -> Array M ix e -> Array M (Lower ix) e
forall r ix ix' e.
(Resize r ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
unsafeResize Sz (Lower ix)
newSz (ix -> Sz ix -> Array M ix e -> Array (R M) ix e
forall r ix e.
Extract r ix e =>
ix -> Sz ix -> Array r ix e -> Array (R r) ix e
unsafeExtract ix
start Sz ix
cutSz Array M ix e
arr)
{-# INLINE unsafeSlice #-}
instance {-# OVERLAPPING #-} OuterSlice M Ix1 e where
unsafeOuterSlice :: Array M Int e -> Int -> Elt M Int e
unsafeOuterSlice !Array M Int e
arr = Array M Int e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array M Int e
arr
{-# INLINE unsafeOuterSlice #-}
instance (Elt M ix e ~ Array M (Lower ix) e, Index ix, Index (Lower ix)) => OuterSlice M ix e where
unsafeOuterSlice :: Array M ix e -> Int -> Elt M ix e
unsafeOuterSlice !Array M ix e
arr !Int
i =
Comp -> Sz (Lower ix) -> (Int -> e) -> Array M (Lower ix) e
forall ix e. Comp -> Sz ix -> (Int -> e) -> Array M ix e
MArray (Array M ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array M ix e
arr) ((Sz1, Sz (Lower ix)) -> Sz (Lower ix)
forall a b. (a, b) -> b
snd (Sz ix -> (Sz1, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz1, Sz (Lower ix))
unconsSz (Array M ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array M ix e
arr))) (Array M ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array M ix e
arr (Int -> e) -> (Int -> Int) -> Int -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
kStart))
where
!kStart :: Int
kStart = Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex (Array M ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array M ix e
arr) (Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
i (Lower ix
forall ix. Index ix => ix
zeroIndex :: Lower ix))
{-# INLINE unsafeOuterSlice #-}
instance {-# OVERLAPPING #-} InnerSlice M Ix1 e where
unsafeInnerSlice :: Array M Int e -> (Sz (Lower Int), Sz1) -> Int -> Elt M Int e
unsafeInnerSlice !Array M Int e
arr (Sz (Lower Int), Sz1)
_ = Array M Int e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array M Int e
arr
{-# INLINE unsafeInnerSlice #-}
instance (Elt M ix e ~ Array M (Lower ix) e, Index ix, Index (Lower ix)) => InnerSlice M ix e where
unsafeInnerSlice :: Array M ix e -> (Sz (Lower ix), Sz1) -> Int -> Elt M ix e
unsafeInnerSlice !Array M ix e
arr (Sz (Lower ix)
szL, Sz1
m) !Int
i =
Comp -> Sz (Lower ix) -> (Int -> e) -> Array M (Lower ix) e
forall ix e. Comp -> Sz ix -> (Int -> e) -> Array M ix e
MArray (Array M ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array M ix e
arr) Sz (Lower ix)
szL (\Int
k -> Array M ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array M ix e
arr (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* Sz1 -> Int
forall ix. Sz ix -> ix
unSz Sz1
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
kStart))
where
!kStart :: Int
kStart = Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex (Array M ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array M ix e
arr) (Lower ix -> Int -> ix
forall ix. Index ix => Lower ix -> Int -> ix
snocDim (Lower ix
forall ix. Index ix => ix
zeroIndex :: Lower ix) Int
i)
{-# INLINE unsafeInnerSlice #-}
instance Index ix => Load M ix e where
size :: Array M ix e -> Sz ix
size = Array M ix e -> Sz ix
forall ix e. Array M ix e -> Sz ix
mSize
{-# INLINE size #-}
getComp :: Array M ix e -> Comp
getComp = Array M ix e -> Comp
forall ix e. Array M ix e -> Comp
mComp
{-# INLINE getComp #-}
loadArrayM :: Scheduler m () -> Array M ix e -> (Int -> e -> m ()) -> m ()
loadArrayM Scheduler m ()
scheduler (MArray _ sz f) = Scheduler m () -> Int -> (Int -> e) -> (Int -> e -> m ()) -> m ()
forall (m :: * -> *) b.
Monad m =>
Scheduler m () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_ Scheduler m ()
scheduler (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) Int -> e
f
{-# INLINE loadArrayM #-}
instance Index ix => StrideLoad M ix e
instance Index ix => Stream M ix e where
toStream :: Array M ix e -> Steps Id e
toStream = Array M ix e -> Steps Id e
forall r ix e (m :: * -> *).
(Monad m, Source r ix e) =>
Array r ix e -> Steps m e
S.steps
{-# INLINE toStream #-}
toStreamIx :: Array M ix e -> Steps Id (ix, e)
toStreamIx = Array M ix e -> Steps Id (ix, e)
forall r ix e (m :: * -> *).
(Monad m, Source r ix e) =>
Array r ix e -> Steps m (ix, e)
S.isteps
{-# INLINE toStreamIx #-}
compute :: forall r ix e r' . (Mutable r ix e, Load r' ix e) => Array r' ix e -> Array r ix e
compute :: 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 :: * -> *).
(Mutable r ix 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' . (Mutable r ix e, Load r' ix e) => Array r' ix e -> Array r ix e
computeS :: 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 :: * -> *).
(Mutable r ix 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'. (Mutable r ix e, Construct r' ix e, Load r' ix e)
=> Array r' ix e
-> Array r ix e
computeP :: 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.
Construct r ix e =>
Comp -> Array r ix e -> Array r ix e
setComp (Array r' ix e -> Comp
forall r ix e. Load r 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'.
(Mutable r ix 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.
Construct r ix e =>
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. (Mutable r ix e, Load r' ix e, MonadIO m)
=> Array r' ix e
-> m (Array r ix e)
computeIO :: 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 (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, Mutable r ix 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 (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 ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (Array r' ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix e
arr))
{-# INLINE computeIO #-}
computePrimM ::
forall r ix e r' m. (Mutable r ix e, Load r' ix e, PrimMonad m)
=> Array r' ix e
-> m (Array r ix e)
computePrimM :: 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, Mutable r ix 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 (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 ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (Array r' ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix e
arr)
{-# INLINE computePrimM #-}
computeAs :: (Mutable r ix e, Load r' ix e) => r -> Array r' ix e -> Array r ix e
computeAs :: r -> Array r' ix e -> Array r ix e
computeAs r
_ = Array r' ix e -> Array r ix e
forall r ix e r'.
(Mutable r ix e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute
{-# INLINE computeAs #-}
computeProxy :: (Mutable r ix 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
computeProxy proxy r
_ = Array r' ix e -> Array r ix e
forall r ix e r'.
(Mutable r ix e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute
{-# INLINE computeProxy #-}
computeSource :: forall r ix e r' . (Mutable r ix e, Source r' ix e)
=> Array r' ix e -> Array r ix e
computeSource :: 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 r' ix e -> Array r ix e
forall r ix e r'.
(Mutable r ix e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute 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 #-}
clone :: Mutable r ix e => Array r ix e -> Array r ix e
clone :: 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 :: * -> *).
(Mutable r ix e, 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 (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 ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (Array r ix e -> Comp
forall r ix e. Load r ix e => 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 :: 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 (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'))
convert :: forall r ix e r' . (Mutable r ix e, Load r' ix e)
=> Array r' ix e -> Array r ix e
convert :: 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'.
(Mutable r ix 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 #-}
convertAs :: (Mutable r ix e, Load r' ix e)
=> r -> Array r' ix e -> Array r ix e
convertAs :: r -> Array r' ix e -> Array r ix e
convertAs r
_ = Array r' ix e -> Array r ix e
forall r ix e r'.
(Mutable r ix e, Load r' ix e) =>
Array r' ix e -> Array r ix e
convert
{-# INLINE convertAs #-}
convertProxy :: (Mutable r ix 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
convertProxy proxy r
_ = Array r' ix e -> Array r ix e
forall r ix e r'.
(Mutable r ix e, Load r' ix e) =>
Array r' ix e -> Array r ix e
convert
{-# INLINE convertProxy #-}
fromRaggedArrayM ::
forall r ix e r' m . (Mutable r ix e, Ragged r' ix e, Load r' ix e, MonadThrow m)
=> Array r' ix e
-> m (Array r ix e)
fromRaggedArrayM :: 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 r ix e. Ragged r ix e => Array r ix e -> Sz ix
edgeSize 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 (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM ShapeException
e) Array r ix e -> m (Array r ix e)
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 ix e (m :: * -> *).
(Mutable r ix e, 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)
traverse (\()
_ -> Comp -> MArray (PrimState IO) r ix e -> IO (Array r ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (Array r' ix e -> Comp
forall r ix e. Load r 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 IO () -> IO ()) -> IO ()
withMassivScheduler_ (Array r' ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix e
arr) ((Scheduler IO () -> IO ()) -> IO ())
-> (Scheduler IO () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Scheduler IO ()
scheduler ->
(IO () -> IO ())
-> (Int -> e -> IO ())
-> Int
-> Int
-> Sz ix
-> Array r' ix e
-> IO ()
forall r ix e (m :: * -> *) a.
(Ragged r ix e, Monad m) =>
(m () -> m ())
-> (Int -> e -> m a) -> Int -> Int -> Sz ix -> Array r ix e -> m ()
loadRagged (Scheduler IO () -> IO () -> IO ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler IO ()
scheduler) (MArray (PrimState IO) r ix e -> Int -> e -> IO ()
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r ix e
MArray (PrimState IO) r ix e
marr) Int
0 (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) Sz ix
sz Array r' ix e
arr)
{-# INLINE fromRaggedArrayM #-}
fromRaggedArray' ::
forall r ix e r'. (Mutable r ix e, Load r' ix e, Ragged r' ix e)
=> Array r' ix e
-> Array r ix e
fromRaggedArray' :: Array r' ix e -> Array r ix e
fromRaggedArray' Array r' ix e
arr = (SomeException -> Array r ix e)
-> (Array r ix e -> Array r ix e)
-> Either SomeException (Array r ix e)
-> Array r ix e
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array r ix e
forall a e. Exception e => e -> a
throw Array r ix e -> Array r ix e
forall a. a -> a
id (Either SomeException (Array r ix e) -> Array r ix e)
-> Either SomeException (Array r ix e) -> Array r ix e
forall a b. (a -> b) -> a -> b
$ Array r' ix e -> Either SomeException (Array r ix e)
forall r ix e r' (m :: * -> *).
(Mutable r ix e, Ragged r' ix e, Load r' ix e, MonadThrow m) =>
Array r' ix e -> m (Array r ix e)
fromRaggedArrayM Array r' ix e
arr
{-# INLINE fromRaggedArray' #-}
computeWithStride ::
forall r ix e r'. (Mutable r ix e, StrideLoad r' ix e)
=> Stride ix
-> Array r' ix e
-> Array r ix e
computeWithStride :: 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 r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r' ix e
arr)
Comp
-> Sz ix
-> (Scheduler IO () -> MArray (PrimState IO) r ix e -> IO ())
-> IO (Array r ix e)
forall r ix e a (m :: * -> *) b.
(Mutable r ix e, PrimMonad m, MonadUnliftIO m) =>
Comp
-> Sz ix
-> (Scheduler m a -> MArray (PrimState m) r ix e -> m b)
-> m (Array r ix e)
unsafeCreateArray_ (Array r' ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix e
arr) Sz ix
sz ((Scheduler IO () -> MArray (PrimState IO) r ix e -> IO ())
-> IO (Array r ix e))
-> (Scheduler IO () -> MArray (PrimState IO) r ix e -> IO ())
-> IO (Array r ix e)
forall a b. (a -> b) -> a -> b
$ \Scheduler IO ()
scheduler MArray (PrimState IO) r ix e
marr ->
Scheduler IO ()
-> Stride ix
-> Sz ix
-> Array r' ix e
-> (Int -> e -> IO ())
-> IO ()
forall r ix e (m :: * -> *).
(StrideLoad r ix e, Monad m) =>
Scheduler m ()
-> Stride ix -> Sz ix -> Array r ix e -> (Int -> e -> m ()) -> m ()
loadArrayWithStrideM Scheduler IO ()
scheduler Stride ix
stride Sz ix
sz Array r' ix e
arr (MArray (PrimState IO) r ix e -> Int -> e -> IO ()
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray (PrimState IO) r ix e
marr)
{-# INLINE computeWithStride #-}
computeWithStrideAs ::
(Mutable r ix 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
computeWithStrideAs r
_ = Stride ix -> Array r' ix e -> Array r ix e
forall r ix e r'.
(Mutable r ix e, StrideLoad r' ix e) =>
Stride ix -> Array r' ix e -> Array r ix e
computeWithStride
{-# INLINE computeWithStrideAs #-}
iterateUntil ::
(Load r' ix e, Mutable 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 :: (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 Int -> Array r ix e -> Array r ix e -> Bool
convergence Int -> Array r ix e -> Array r' ix e
iteration Array r ix e
initArr0
| Int -> Array r ix e -> Array r ix e -> Bool
convergence Int
0 Array r ix e
initArr0 Array r ix e
initArr1 = Array r ix e
initArr1
| Bool
otherwise =
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 loadArr :: Array r' ix e
loadArr = Int -> Array r ix e -> Array r' ix e
iteration Int
1 Array r ix e
initArr1
MArray RealWorld r ix e
marr <- Sz ix -> IO (MArray (PrimState IO) r ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew (Array r' ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r' ix e
loadArr)
(Int
-> Array r ix e
-> Array r ix e
-> MArray (PrimState IO) r ix e
-> IO Bool)
-> (Int -> Array r ix e -> Array r' ix e)
-> Int
-> Array r ix e
-> Array r' ix e
-> MArray (PrimState IO) r ix e
-> IO (Array r ix e)
forall r' ix e r (m :: * -> *).
(Load r' ix e, Mutable r ix e, PrimMonad m, MonadIO m,
PrimState m ~ RealWorld) =>
(Int
-> Array r ix e
-> Array r ix e
-> MArray (PrimState m) r ix e
-> m Bool)
-> (Int -> Array r ix e -> Array r' ix e)
-> Int
-> Array r ix e
-> Array r' ix e
-> MArray (PrimState m) r ix e
-> m (Array r ix e)
iterateLoop
(\Int
n Array r ix e
a Array r ix e
a' MArray (PrimState IO) r ix e
_ -> Bool -> IO Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> IO Bool) -> Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ Int -> Array r ix e -> Array r ix e -> Bool
convergence Int
n Array r ix e
a Array r ix e
a')
Int -> Array r ix e -> Array r' ix e
iteration
Int
1
Array r ix e
initArr1
Array r' ix e
loadArr
(Array r ix e -> MArray RealWorld r ix e -> MArray RealWorld r ix e
forall r ix e s. Array r ix e -> MArray s r ix e -> MArray s r ix e
asArr Array r ix e
initArr0 MArray RealWorld r ix e
marr)
where
!initArr1 :: Array r ix e
initArr1 = Array r' ix e -> Array r ix e
forall r ix e r'.
(Mutable r ix e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute (Array r' ix e -> Array r ix e) -> Array r' ix e -> Array r ix e
forall a b. (a -> b) -> a -> b
$ Int -> Array r ix e -> Array r' ix e
iteration Int
0 Array r ix e
initArr0
asArr :: Array r ix e -> MArray s r ix e -> MArray s r ix e
asArr :: Array r ix e -> MArray s r ix e -> MArray s r ix e
asArr Array r ix e
_ = MArray s r ix e -> MArray s r ix e
forall a. a -> a
id
{-# INLINE iterateUntil #-}
iterateUntilM ::
(Load r' ix e, Mutable r ix e, PrimMonad m, MonadIO m, PrimState m ~ RealWorld)
=> (Int -> Array r ix e -> MArray (PrimState m) r ix e -> m Bool)
-> (Int -> Array r ix e -> Array r' ix e)
-> Array r ix e
-> m (Array r ix e)
iterateUntilM :: (Int -> Array r ix e -> MArray (PrimState m) r ix e -> m Bool)
-> (Int -> Array r ix e -> Array r' ix e)
-> Array r ix e
-> m (Array r ix e)
iterateUntilM Int -> Array r ix e -> MArray (PrimState m) r ix e -> m Bool
convergence Int -> Array r ix e -> Array r' ix e
iteration Array r ix e
initArr0 = do
let loadArr0 :: Array r' ix e
loadArr0 = Int -> Array r ix e -> Array r' ix e
iteration Int
0 Array r ix e
initArr0
MArray RealWorld r ix e
initMArr1 <- Sz ix -> m (MArray (PrimState m) r ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew (Array r' ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r' ix e
loadArr0)
MArray RealWorld r ix e -> Array r' ix e -> m ()
forall r' ix' e r ix (m :: * -> *).
(Load r' ix' e, Mutable r ix e, MonadIO m) =>
MArray RealWorld r ix e -> Array r' ix' e -> m ()
computeInto MArray RealWorld r ix e
initMArr1 Array r' ix e
loadArr0
Bool
shouldStop <- Int -> Array r ix e -> MArray (PrimState m) r ix e -> m Bool
convergence Int
0 Array r ix e
initArr0 MArray RealWorld r ix e
MArray (PrimState m) r ix e
initMArr1
Array r ix e
initArr1 <- Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (Array r' ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix e
loadArr0) MArray RealWorld r ix e
MArray (PrimState m) r ix e
initMArr1
if Bool
shouldStop
then Array r ix e -> m (Array r ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Array r ix e
initArr1
else do
let loadArr1 :: Array r' ix e
loadArr1 = Int -> Array r ix e -> Array r' ix e
iteration Int
1 Array r ix e
initArr1
MArray RealWorld r ix e
marr <- Sz ix -> m (MArray (PrimState m) r ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew (Array r' ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r' ix e
loadArr1)
(Int
-> Array r ix e
-> Array r ix e
-> MArray (PrimState m) r ix e
-> m Bool)
-> (Int -> Array r ix e -> Array r' ix e)
-> Int
-> Array r ix e
-> Array r' ix e
-> MArray (PrimState m) r ix e
-> m (Array r ix e)
forall r' ix e r (m :: * -> *).
(Load r' ix e, Mutable r ix e, PrimMonad m, MonadIO m,
PrimState m ~ RealWorld) =>
(Int
-> Array r ix e
-> Array r ix e
-> MArray (PrimState m) r ix e
-> m Bool)
-> (Int -> Array r ix e -> Array r' ix e)
-> Int
-> Array r ix e
-> Array r' ix e
-> MArray (PrimState m) r ix e
-> m (Array r ix e)
iterateLoop (\Int
n Array r ix e
a Array r ix e
_ -> Int -> Array r ix e -> MArray (PrimState m) r ix e -> m Bool
convergence Int
n Array r ix e
a) Int -> Array r ix e -> Array r' ix e
iteration Int
1 Array r ix e
initArr1 Array r' ix e
loadArr1 MArray RealWorld r ix e
MArray (PrimState m) r ix e
marr
{-# INLINE iterateUntilM #-}
iterateLoop ::
(Load r' ix e, Mutable r ix e, PrimMonad m, MonadIO m, PrimState m ~ RealWorld)
=> (Int -> Array r ix e -> Array r ix e -> MArray (PrimState m) r ix e -> m Bool)
-> (Int -> Array r ix e -> Array r' ix e)
-> Int
-> Array r ix e
-> Array r' ix e
-> MArray (PrimState m) r ix e
-> m (Array r ix e)
iterateLoop :: (Int
-> Array r ix e
-> Array r ix e
-> MArray (PrimState m) r ix e
-> m Bool)
-> (Int -> Array r ix e -> Array r' ix e)
-> Int
-> Array r ix e
-> Array r' ix e
-> MArray (PrimState m) r ix e
-> m (Array r ix e)
iterateLoop Int
-> Array r ix e
-> Array r ix e
-> MArray (PrimState m) r ix e
-> m Bool
convergence Int -> Array r ix e -> Array r' ix e
iteration = Int
-> Array r ix e
-> Array r' ix e
-> MArray RealWorld r ix e
-> m (Array r ix e)
Int
-> Array r ix e
-> Array r' ix e
-> MArray (PrimState m) r ix e
-> m (Array r ix e)
go
where
go :: Int
-> Array r ix e
-> Array r' ix e
-> MArray RealWorld r ix e
-> m (Array r ix e)
go !Int
n !Array r ix e
arr !Array r' ix e
loadArr !MArray RealWorld r ix e
marr = do
let !sz :: Sz ix
sz = Array r' ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r' ix e
loadArr
!k :: Int
k = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
!mk :: Int
mk = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (MArray RealWorld r ix e -> Sz ix
forall r ix e s. Mutable r ix e => MArray s r ix e -> Sz ix
msize MArray RealWorld r ix e
marr)
MArray RealWorld r ix e
marr' <-
if Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
mk
then MArray RealWorld r ix e -> m (MArray RealWorld r ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure MArray RealWorld r ix e
marr
else if Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
mk
then MArray (PrimState m) r ix e
-> Sz ix -> m (MArray (PrimState m) r ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
MArray (PrimState m) r ix e
-> Sz ix -> m (MArray (PrimState m) r ix e)
unsafeLinearShrink MArray RealWorld r ix e
MArray (PrimState m) r ix e
marr Sz ix
sz
else MArray (PrimState m) r ix e
-> Sz ix -> m (MArray (PrimState m) r ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
MArray (PrimState m) r ix e
-> Sz ix -> m (MArray (PrimState m) r ix e)
unsafeLinearGrow MArray RealWorld r ix e
MArray (PrimState m) r ix e
marr Sz ix
sz
MArray RealWorld r ix e -> Array r' ix e -> m ()
forall r' ix' e r ix (m :: * -> *).
(Load r' ix' e, Mutable r ix e, MonadIO m) =>
MArray RealWorld r ix e -> Array r' ix' e -> m ()
computeInto MArray RealWorld r ix e
marr' Array r' ix e
loadArr
Array r ix e
arr' <- Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (Array r' ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r' ix e
loadArr) MArray RealWorld r ix e
MArray (PrimState m) r ix e
marr'
Bool
shouldStop <- Int
-> Array r ix e
-> Array r ix e
-> MArray (PrimState m) r ix e
-> m Bool
convergence Int
n Array r ix e
arr Array r ix e
arr' MArray RealWorld r ix e
MArray (PrimState m) r ix e
marr'
if Bool
shouldStop
then Array r ix e -> m (Array r ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Array r ix e
arr'
else do
MArray RealWorld r ix e
nextMArr <- Array r ix e -> m (MArray (PrimState m) r ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Array r ix e -> m (MArray (PrimState m) r ix e)
unsafeThaw Array r ix e
arr
Int
-> Array r ix e
-> Array r' ix e
-> MArray RealWorld r ix e
-> m (Array r ix e)
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Array r ix e
arr' (Int -> Array r ix e -> Array r' ix e
iteration (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Array r ix e
arr') MArray RealWorld r ix e
nextMArr
{-# INLINE iterateLoop #-}