{-# 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
-- Copyright   : (c) Alexey Kuleshevich 2018-2021
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
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.Massiv.Core.Operations
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


-- | General Manifest representation
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
eqArrays 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
compareArrays e -> e -> Ordering
forall a. Ord a => a -> a -> Ordering
compare
  {-# INLINE compare #-}


-- | /O(1)/ - Conversion of `Manifest` arrays to `M` representation.
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 #-}


-- | Row-major sequentia folding over a Manifest array.
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 #-}


instance Num e => FoldNumeric M e where
  unsafeDotProduct :: Array M ix e -> Array M ix e -> e
unsafeDotProduct = Array M ix e -> Array M ix e -> e
forall e r ix.
(Num e, Source r ix e) =>
Array r ix e -> Array r ix e -> e
defaultUnsafeDotProduct
  {-# INLINE unsafeDotProduct #-}
  powerSumArray :: Array M ix e -> Int -> e
powerSumArray = Array M ix e -> Int -> e
forall r ix e. (Source r ix e, Num e) => Array r ix e -> Int -> e
defaultPowerSumArray
  {-# INLINE powerSumArray #-}
  foldArray :: (e -> e -> e) -> e -> Array M ix e -> e
foldArray = (e -> e -> e) -> e -> Array M ix e -> e
forall r ix e.
Source r ix e =>
(e -> e -> e) -> e -> Array r ix e -> e
defaultFoldArray
  {-# INLINE foldArray #-}

-- | 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' . (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 #-}

-- | 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' . (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 #-}


-- | 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'. (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 #-}

-- | 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. (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 #-}

-- | Compute an array in `PrimMonad` sequentially disregarding predefined computation
-- strategy.
--
-- @since 0.4.5
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 #-}


-- | 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 :: (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 #-}


-- | 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 :: (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 #-}


-- | 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' . (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 #-}


-- | /O(n)/ - Make an exact immutable copy of an Array.
--
-- @since 0.1.0
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 #-}


-- | /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 :: 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'))


-- | /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' . (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 #-}

-- | Same as `convert`, but let's you supply resulting representation type as an argument.
--
-- @since 0.1.0
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 #-}


-- | Same as `convert` and `convertAs`, but let's you supply resulting representation type as a
-- proxy argument.
--
-- @since 0.1.1
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 #-}


-- | 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 . (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 #-}


-- | Same as `fromRaggedArrayM`, but will throw a pure exception if its shape is not
-- rectangular.
--
-- @since 0.1.1
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' #-}


-- | Same as `compute`, but with `Stride`.
--
-- /O(n div k)/ - Where @n@ is numer 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'. (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 #-}


-- | Same as `computeWithStride`, but with ability to specify resulting array representation.
--
-- @since 0.3.0
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 #-}



-- | 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
-- >>> a = computeAs P $ makeLoadArrayS (Sz2 8 8) (0 :: Int) $ \ w -> () <$ w (0 :. 0) 1
-- >>> a
-- 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 ]
--   ]
-- >>> nextPascalRow cur above = if cur == 0 then above else cur
-- >>> 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) a
-- 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, Mutable 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 :: (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 -> Comp -> 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 -> Comp -> 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 Comp
comp MArray (PrimState IO) r ix e
marr' -> Int -> Array r ix e -> Array r ix e -> Bool
convergence Int
n Array r ix e
a (Array r ix e -> Bool) -> IO (Array r ix e) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Comp -> MArray (PrimState IO) r ix e -> IO (Array r ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
comp MArray (PrimState IO) r ix e
marr')
        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 #-}

-- | Monadic version of `iterateUntil` where at each iteration mutable version of an array
-- is available.
--
-- @since 0.3.6
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)
  -- ^ 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 -> 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 :: (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 -> Comp -> 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 -> Comp -> 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 Comp
_ -> 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 -> Comp -> 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 -> Comp -> 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 -> Comp -> 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)
          !comp :: Comp
comp = 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
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
      Bool
shouldStop <- Int
-> Array r ix e -> Comp -> MArray (PrimState m) r ix e -> m Bool
convergence Int
n Array r ix e
arr Comp
comp MArray RealWorld r ix e
MArray (PrimState m) r ix e
marr'
      Array r ix e
arr' <- Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
comp 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 #-}