{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module      : Data.Massiv.Core.Common
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
module Data.Massiv.Core.Common (
  Array,
  Vector,
  Matrix,
  MArray,
  MVector,
  MMatrix,
  Steps (..),
  Stream (..),
  Strategy (..),
  Source (..),
  PrefIndex (..),
  Load (..),
  StrideLoad (..),
  Size (..),
  Shape (..),
  Manifest (..),
  Mutable,
  Comp (..),
  Scheduler,
  numWorkers,
  scheduleWork,
  scheduleWork_,
  WorkerStates,
  unsafeRead,
  unsafeWrite,
  unsafeModify,
  unsafeLinearModify,
  unsafeSwap,
  unsafeLinearSwap,
  unsafeDefaultLinearShrink,
  Ragged (..),
  empty,
  singleton,

  -- * Size
  elemsCount,
  isNotNull,
  isEmpty,
  isNotEmpty,
  Sz (SafeSz),
  LengthHint (..),

  -- * Indexing
  (!?),
  index,
  indexM,
  (!),
  index',
  (??),
  defaultIndex,
  borderIndex,
  evaluateM,
  evaluate',
  inline0,
  inline1,
  inline2,
  module Data.Massiv.Core.Index,

  -- * Common Operations
  Semigroup ((<>)),

  -- * Exceptions
  MonadThrow (..),
  IndexException (..),
  SizeException (..),
  ShapeException (..),
  module Data.Massiv.Core.Exception,
  Proxy (..),
  Id (..),

  -- * Stateful Monads
  runST,
  ST,
  MonadUnliftIO (..),
  MonadIO (liftIO),
  PrimMonad (PrimState),
  RealWorld,
) where

#if !MIN_VERSION_base(4,11,0)
import Data.Semigroup (Semigroup((<>)))
#endif
import Control.Monad.Catch (MonadThrow (..))
import Control.Monad.IO.Unlift (MonadIO (liftIO), MonadUnliftIO (..))
import Control.Monad.Primitive
import Control.Monad.ST
import Control.Scheduler (
  Comp (..),
  Scheduler,
  WorkerStates,
  numWorkers,
  scheduleWork,
  scheduleWork_,
  trivialScheduler_,
 )
import Data.Kind
import Data.Massiv.Core.Exception
import Data.Massiv.Core.Index
import Data.Massiv.Core.Index.Internal (Sz (SafeSz))
import Data.Typeable
import qualified Data.Stream.Monadic as S (Stream)
import Data.Vector.Fusion.Util
import GHC.Exts (IsList)

#include "massiv.h"

-- | The array family. Representations @r@ describe how data is arranged or computed. All
-- arrays have a common property that each index @ix@ always maps to the same unique
-- element @e@, even if that element does not yet exist in memory and the array has to be
-- computed in order to get the value of that element. Data is always arranged in a nested
-- row-major fashion. Rank of an array is specified by @`Dimensions` ix@.
--
-- @since 0.1.0
data family Array r ix e :: Type

-- | Type synonym for a single dimension array, or simply a flat vector.
--
-- @since 0.5.0
type Vector r e = Array r Ix1 e

-- | Type synonym for a two-dimentsional array, or simply a matrix.
--
-- @since 0.5.0
type Matrix r e = Array r Ix2 e

-- | Mutable version of a `Manifest` `Array`. The extra type argument @s@ is for
-- the state token used by `IO` and `ST`.
--
-- @since 0.1.0
data family MArray s r ix e :: Type

-- | Type synonym for a single dimension mutable array, or simply a flat mutable vector.
--
-- @since 0.5.0
type MVector s r e = MArray s r Ix1 e

-- | Type synonym for a two-dimentsional mutable array, or simply a mutable matrix.
--
-- @since 0.5.0
type MMatrix s r e = MArray s r Ix2 e

class Load r ix e => Stream r ix e where
  toStream :: Array r ix e -> Steps Id e

  toStreamIx :: Array r ix e -> Steps Id (ix, e)

data Steps m e = Steps
  { forall (m :: * -> *) e. Steps m e -> Stream m e
stepsStream :: S.Stream m e
  , forall (m :: * -> *) e. Steps m e -> LengthHint
stepsSize :: LengthHint
  }

class Typeable r => Strategy r where
  -- | Set computation strategy for this array
  --
  -- ==== __Example__
  --
  -- >>> :set -XTypeApplications
  -- >>> import Data.Massiv.Array
  -- >>> a = singleton @DL @Ix1 @Int 0
  -- >>> a
  -- Array DL Seq (Sz1 1)
  --   [ 0 ]
  -- >>> setComp (ParN 6) a -- use 6 capabilities
  -- Array DL (ParN 6) (Sz1 1)
  --   [ 0 ]
  setComp :: Comp -> Array r ix e -> Array r ix e

  -- | Get computation strategy of this array
  --
  -- @since 0.1.0
  getComp :: Array r ix e -> Comp

  -- | Array representation. Representation is never evaluated in @massiv@,
  -- therefore default implementation is bottom. However, it is recommended to
  -- supply a constructor that doesn't result in an error when evaluated.
  --
  -- @since 1.0.2
  repr :: r
  repr =
    [Char] -> r
forall a. HasCallStack => [Char] -> a
error ([Char] -> r) -> [Char] -> r
forall a b. (a -> b) -> a -> b
$
      [Char]
"Array representation should never be evaluated: "
        [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ TypeRep -> [Char]
forall a. Show a => a -> [Char]
show (Proxy r -> TypeRep
forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (Proxy r
forall {k} (t :: k). Proxy t
Proxy :: Proxy r))

-- | Size hint
--
-- @since 1.0.0
data LengthHint
  = -- | Exact known size
    LengthExact Sz1
  | -- | Upper bound on the size
    LengthMax Sz1
  | -- | Unknown size
    LengthUnknown
  deriving (LengthHint -> LengthHint -> Bool
(LengthHint -> LengthHint -> Bool)
-> (LengthHint -> LengthHint -> Bool) -> Eq LengthHint
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: LengthHint -> LengthHint -> Bool
== :: LengthHint -> LengthHint -> Bool
$c/= :: LengthHint -> LengthHint -> Bool
/= :: LengthHint -> LengthHint -> Bool
Eq, Int -> LengthHint -> [Char] -> [Char]
[LengthHint] -> [Char] -> [Char]
LengthHint -> [Char]
(Int -> LengthHint -> [Char] -> [Char])
-> (LengthHint -> [Char])
-> ([LengthHint] -> [Char] -> [Char])
-> Show LengthHint
forall a.
(Int -> a -> [Char] -> [Char])
-> (a -> [Char]) -> ([a] -> [Char] -> [Char]) -> Show a
$cshowsPrec :: Int -> LengthHint -> [Char] -> [Char]
showsPrec :: Int -> LengthHint -> [Char] -> [Char]
$cshow :: LengthHint -> [Char]
show :: LengthHint -> [Char]
$cshowList :: [LengthHint] -> [Char] -> [Char]
showList :: [LengthHint] -> [Char] -> [Char]
Show)

-- | The shape of an array. It is different from `Size` in that it can be applicable to
-- non-square matrices and might not be available in constant time.
--
-- @since 1.0.0
class Index ix => Shape r ix where
  -- | /O(1)/ - Check what do we know about the number of elements without doing any work
  --
  -- @since 1.0.0
  linearSizeHint :: Array r ix e -> LengthHint
  linearSizeHint = Sz1 -> LengthHint
LengthExact (Sz1 -> LengthHint)
-> (Array r ix e -> Sz1) -> Array r ix e -> LengthHint
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> Sz1
forall e. Array r ix e -> Sz1
forall r ix e. Shape r ix => Array r ix e -> Sz1
linearSize
  {-# INLINE linearSizeHint #-}

  -- | /O(n)/ - possibly iterate over the whole array before producing the answer
  --
  -- @since 0.5.8
  linearSize :: Array r ix e -> Sz1
  default linearSize :: Size r => Array r ix e -> Sz1
  linearSize = Int -> Sz1
forall ix. ix -> Sz ix
SafeSz (Int -> Sz1) -> (Array r ix e -> Int) -> Array r ix e -> Sz1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> Int
forall ix r e. (Index ix, Size r) => Array r ix e -> Int
elemsCount
  {-# INLINE linearSize #-}

  -- | /O(n)/ - Rectangular size of an array that is inferred from looking at the first row in
  -- each dimensions. For rectangular arrays this is the same as `size`
  --
  -- @since 1.0.0
  outerSize :: Array r ix e -> Sz ix
  default outerSize :: Size r => Array r ix e -> Sz ix
  outerSize = Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size
  {-# INLINE outerSize #-}

  -- | /O(1)/ - Get the possible maximum linear size of an immutabe array. If the lookup
  -- of size in constant time is not possible, `Nothing` will be returned. This value
  -- will be used as the initial size of the mutable array into which the loading will
  -- happen.
  --
  -- @since 1.0.0
  maxLinearSize :: Array r ix e -> Maybe Sz1
  maxLinearSize = LengthHint -> Maybe Sz1
lengthHintUpperBound (LengthHint -> Maybe Sz1)
-> (Array r ix e -> LengthHint) -> Array r ix e -> Maybe Sz1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> LengthHint
forall e. Array r ix e -> LengthHint
forall r ix e. Shape r ix => Array r ix e -> LengthHint
linearSizeHint
  {-# INLINE maxLinearSize #-}

  -- | /O(1)/ - Check whether an array is empty or not.
  --
  -- ==== __Examples__
  --
  -- >>> import Data.Massiv.Array
  -- >>> isNull $ range Seq (Ix2 10 20) (11 :. 21)
  -- False
  -- >>> isNull $ range Seq (Ix2 10 20) (10 :. 21)
  -- True
  -- >>> isNull (empty :: Array D Ix5 Int)
  -- True
  -- >>> isNull $ sfromList []
  -- True
  --
  -- @since 1.0.0
  isNull :: Array r ix e -> Bool
  isNull = (Sz1
forall ix. Index ix => Sz ix
zeroSz Sz1 -> Sz1 -> Bool
forall a. Eq a => a -> a -> Bool
==) (Sz1 -> Bool) -> (Array r ix e -> Sz1) -> Array r ix e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> Sz1
forall e. Array r ix e -> Sz1
forall r ix e. Shape r ix => Array r ix e -> Sz1
linearSize
  {-# INLINE isNull #-}

lengthHintUpperBound :: LengthHint -> Maybe Sz1
lengthHintUpperBound :: LengthHint -> Maybe Sz1
lengthHintUpperBound = \case
  LengthExact Sz1
sz -> Sz1 -> Maybe Sz1
forall a. a -> Maybe a
Just Sz1
sz
  LengthMax Sz1
sz -> Sz1 -> Maybe Sz1
forall a. a -> Maybe a
Just Sz1
sz
  LengthHint
LengthUnknown -> Maybe Sz1
forall a. Maybe a
Nothing
{-# INLINE lengthHintUpperBound #-}

-- | Arrays that have information about their size availible in constant
-- time.
class Size r where
  -- | /O(1)/ - Get the exact size of an immutabe array. Most of the time will
  -- produce the size in constant time, except for `Data.Massiv.Array.DS`
  -- representation, which could result in evaluation of the whole stream. See
  -- `maxLinearSize` and `Data.Massiv.Vector.slength` for more info.
  --
  -- @since 0.1.0
  size :: Array r ix e -> Sz ix

  -- | /O(1)/ - Change the size of an array. Total number of elements should be the same, but it is
  -- not validated.
  --
  -- @since 0.1.0
  unsafeResize :: (Index ix, Index ix') => Sz ix' -> Array r ix e -> Array r ix' e

-- | Prefered indexing function.
data PrefIndex ix e
  = PrefIndex (ix -> e)
  | PrefIndexLinear (Int -> e)

instance Functor (PrefIndex ix) where
  fmap :: forall a b. (a -> b) -> PrefIndex ix a -> PrefIndex ix b
fmap a -> b
f = \case
    PrefIndex ix -> a
ig -> (ix -> b) -> PrefIndex ix b
forall ix e. (ix -> e) -> PrefIndex ix e
PrefIndex (a -> b
f (a -> b) -> (ix -> a) -> ix -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> a
ig)
    PrefIndexLinear Int -> a
ig -> (Int -> b) -> PrefIndex ix b
forall ix e. (Int -> e) -> PrefIndex ix e
PrefIndexLinear (a -> b
f (a -> b) -> (Int -> a) -> Int -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> a
ig)
  {-# INLINE fmap #-}
  <$ :: forall a b. a -> PrefIndex ix b -> PrefIndex ix a
(<$) a
e PrefIndex ix b
_ = (Int -> a) -> PrefIndex ix a
forall ix e. (Int -> e) -> PrefIndex ix e
PrefIndexLinear (a -> Int -> a
forall a b. a -> b -> a
const a
e)
  {-# INLINE (<$) #-}

-- | Arrays that can be used as source to practically any manipulation function.
class (Strategy r, Size r) => Source r e where
  {-# MINIMAL (unsafeIndex | unsafeLinearIndex), unsafeLinearSlice #-}

  -- | Lookup element in the array. No bounds check is performed and access of
  -- arbitrary memory is possible when invalid index is supplied.
  --
  -- @since 0.1.0
  unsafeIndex :: HAS_CALL_STACK => Index ix => Array r ix e -> ix -> e
  unsafeIndex !Array r ix e
arr = Array r ix e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr (Int -> e) -> (ix -> Int) -> ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex (Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
arr)
  {-# INLINE unsafeIndex #-}

  -- | Lookup element in the array using flat index in a row-major fashion. No
  -- bounds check is performed
  --
  -- @since 0.1.0
  unsafeLinearIndex :: HAS_CALL_STACK => Index ix => Array r ix e -> Int -> e
  unsafeLinearIndex !Array r ix e
arr = Array r ix e -> ix -> e
forall ix. Index ix => Array r ix e -> ix -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr (ix -> e) -> (Int -> ix) -> Int -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex (Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
arr)
  {-# INLINE unsafeLinearIndex #-}

  -- | Alternative indexing function that can choose an index that is most
  -- efficient for underlying representation
  --
  -- @since 1.0.2
  unsafePrefIndex :: HAS_CALL_STACK => Index ix => Array r ix e -> PrefIndex ix e
  unsafePrefIndex !Array r ix e
arr = (Int -> e) -> PrefIndex ix e
forall ix e. (Int -> e) -> PrefIndex ix e
PrefIndexLinear (Array r ix e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr)
  {-# INLINE unsafePrefIndex #-}

  -- | /O(1)/ - Take a slice out of an array from the outside
  --
  -- @since 0.1.0
  unsafeOuterSlice
    :: HAS_CALL_STACK
    => (Index ix, Index (Lower ix))
    => Array r ix e
    -> Sz (Lower ix)
    -> Int
    -> Array r (Lower ix) e
  unsafeOuterSlice Array r ix e
arr Sz (Lower ix)
sz Int
i = Sz (Lower ix) -> Array r Int e -> Array r (Lower ix) e
forall r ix ix' e.
(Size r, Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
forall ix ix' e.
(Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
unsafeResize Sz (Lower ix)
sz (Array r Int e -> Array r (Lower ix) e)
-> Array r Int e -> Array r (Lower ix) e
forall a b. (a -> b) -> a -> b
$ Int -> Sz1 -> Array r ix e -> Array r Int e
forall ix. Index ix => Int -> Sz1 -> Array r ix e -> Array r Int e
forall r e ix.
(Source r e, Index ix) =>
Int -> Sz1 -> Array r ix e -> Array r Int e
unsafeLinearSlice Int
i (Sz (Lower ix) -> Sz1
forall ix. Index ix => Sz ix -> Sz1
toLinearSz Sz (Lower ix)
sz) Array r ix e
arr
  {-# INLINE unsafeOuterSlice #-}

  -- | /O(1)/ - Source arrays also give us ability to look at their linear slices in
  -- constant time
  --
  -- @since 0.5.0
  unsafeLinearSlice :: Index ix => Ix1 -> Sz1 -> Array r ix e -> Array r Ix1 e

-- | Any array that can be computed and loaded into memory
class (Strategy r, Shape r ix) => Load r ix e where
  {-# MINIMAL (makeArray | makeArrayLinear), (iterArrayLinearST_ | iterArrayLinearWithSetST_) #-}

  -- | Construct an Array. Resulting type either has to be unambiguously inferred or restricted
  -- manually, like in the example below. Use "Data.Massiv.Array.makeArrayR" if you'd like to
  -- specify representation as an argument.
  --
  -- >>> import Data.Massiv.Array
  -- >>> makeArray Seq (Sz (3 :. 4)) (\ (i :. j) -> if i == j then i else 0) :: Array D Ix2 Int
  -- Array D Seq (Sz (3 :. 4))
  --   [ [ 0, 0, 0, 0 ]
  --   , [ 0, 1, 0, 0 ]
  --   , [ 0, 0, 2, 0 ]
  --   ]
  --
  -- Instead of restricting the full type manually we can use @TypeApplications@ as convenience:
  --
  -- >>> :set -XTypeApplications
  -- >>> makeArray @P @_ @Double Seq (Sz2 3 4) $ \(i :. j) -> logBase (fromIntegral i) (fromIntegral j)
  -- Array P Seq (Sz (3 :. 4))
  --   [ [ NaN, -0.0, -0.0, -0.0 ]
  --   , [ -Infinity, NaN, Infinity, Infinity ]
  --   , [ -Infinity, 0.0, 1.0, 1.5849625007211563 ]
  --   ]
  --
  -- @since 0.1.0
  makeArray
    :: Comp
    -- ^ Computation strategy. Useful constructors are `Seq` and `Par`
    -> Sz ix
    -- ^ Size of the result array.
    -> (ix -> e)
    -- ^ Function to generate elements at a particular index
    -> Array r ix e
  makeArray Comp
comp Sz ix
sz ix -> e
f = Comp -> Sz ix -> (Int -> e) -> Array r ix e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear Comp
comp Sz ix
sz (ix -> e
f (ix -> e) -> (Int -> ix) -> Int -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz)
  {-# INLINE makeArray #-}

  -- | Same as `makeArray`, but produce elements using linear row-major index.
  --
  -- >>> import Data.Massiv.Array
  -- >>> makeArrayLinear Seq (Sz (2 :. 4)) id :: Array D Ix2 Int
  -- Array D Seq (Sz (2 :. 4))
  --   [ [ 0, 1, 2, 3 ]
  --   , [ 4, 5, 6, 7 ]
  --   ]
  --
  -- @since 0.3.0
  makeArrayLinear :: Comp -> Sz ix -> (Int -> e) -> Array r ix e
  makeArrayLinear Comp
comp Sz ix
sz Int -> e
f = Comp -> Sz ix -> (ix -> e) -> Array r ix e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
comp Sz ix
sz (Int -> e
f (Int -> e) -> (ix -> Int) -> ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
sz)
  {-# INLINE makeArrayLinear #-}

  -- | Construct an array of the specified size that contains the same element in all of
  -- the cells.
  --
  -- @since 0.3.0
  replicate :: Comp -> Sz ix -> e -> Array r ix e
  replicate Comp
comp Sz ix
sz !e
e = Comp -> Sz ix -> (Int -> e) -> Array r ix e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear Comp
comp Sz ix
sz (e -> Int -> e
forall a b. a -> b -> a
const e
e)
  {-# INLINE replicate #-}

  -- | Iterate over an array with a ST action that is applied to each element and its index.
  --
  -- @since 1.0.0
  iterArrayLinearST_
    :: Scheduler s ()
    -> Array r ix e
    -- ^ Array that is being loaded
    -> (Int -> e -> ST s ())
    -- ^ Function that writes an element into target array
    -> ST s ()
  iterArrayLinearST_ Scheduler s ()
scheduler Array r ix e
arr Int -> e -> ST s ()
uWrite =
    Scheduler s ()
-> Array r ix e
-> (Int -> e -> ST s ())
-> (Int -> Sz1 -> e -> ST s ())
-> ST s ()
forall s.
Scheduler s ()
-> Array r ix e
-> (Int -> e -> ST s ())
-> (Int -> Sz1 -> e -> ST s ())
-> ST s ()
forall r ix e s.
Load r ix e =>
Scheduler s ()
-> Array r ix e
-> (Int -> e -> ST s ())
-> (Int -> Sz1 -> e -> ST s ())
-> ST s ()
iterArrayLinearWithSetST_ Scheduler s ()
scheduler Array r ix e
arr Int -> e -> ST s ()
uWrite ((Int -> Sz1 -> e -> ST s ()) -> ST s ())
-> (Int -> Sz1 -> e -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
offset Sz1
sz e
e ->
      Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> ST s ()) -> ST s ()
forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
offset (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz1 -> Int
forall ix. Sz ix -> ix
unSz Sz1
sz)) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> e -> ST s ()
`uWrite` e
e)
  {-# INLINE iterArrayLinearST_ #-}

  -- | Similar to `iterArrayLinearST_`. Except it also accepts a function that is
  -- potentially optimized for setting many cells in a region to the same
  -- value.
  --
  -- @since 1.0.0
  iterArrayLinearWithSetST_
    :: Scheduler s ()
    -> Array r ix e
    -- ^ Array that is being loaded
    -> (Ix1 -> e -> ST s ())
    -- ^ Function that writes an element into target array
    -> (Ix1 -> Sz1 -> e -> ST s ())
    -- ^ Function that efficiently sets a region of an array
    -- to the supplied value target array
    -> ST s ()
  iterArrayLinearWithSetST_ Scheduler s ()
scheduler Array r ix e
arr Int -> e -> ST s ()
uWrite Int -> Sz1 -> e -> ST s ()
_ = Scheduler s () -> Array r ix e -> (Int -> e -> ST s ()) -> ST s ()
forall s.
Scheduler s () -> Array r ix e -> (Int -> e -> ST s ()) -> ST s ()
forall r ix e s.
Load r ix e =>
Scheduler s () -> Array r ix e -> (Int -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ Scheduler s ()
scheduler Array r ix e
arr Int -> e -> ST s ()
uWrite
  {-# INLINE iterArrayLinearWithSetST_ #-}

  -- | Load into a supplied mutable array sequentially. Returned array does not have to be
  -- the same.
  --
  -- @since 1.0.0
  unsafeLoadIntoST
    :: Manifest r' e
    => MVector s r' e
    -> Array r ix e
    -> ST s (MArray s r' ix e)
  unsafeLoadIntoST MVector s r' e
mvec Array r ix e
arr = do
    let sz :: Sz ix
sz = Array r ix e -> Sz ix
forall e. Array r ix e -> Sz ix
forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array r ix e
arr
    MVector s r' e
mvec' <- MArray (PrimState (ST s)) r' Int e
-> Sz1 -> ST s (MArray (PrimState (ST s)) r' Int e)
forall r e (f :: * -> *).
(Manifest r e, PrimMonad f) =>
MVector (PrimState f) r e -> Sz1 -> f (MVector (PrimState f) r e)
resizeMVector MVector s r' e
MArray (PrimState (ST s)) r' Int e
mvec (Sz1 -> ST s (MArray (PrimState (ST s)) r' Int e))
-> Sz1 -> ST s (MArray (PrimState (ST s)) r' Int e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> Sz1
forall ix. Index ix => Sz ix -> Sz1
toLinearSz Sz ix
sz
    Scheduler s ()
-> Array r ix e
-> (Int -> e -> ST s ())
-> (Int -> Sz1 -> e -> ST s ())
-> ST s ()
forall s.
Scheduler s ()
-> Array r ix e
-> (Int -> e -> ST s ())
-> (Int -> Sz1 -> e -> ST s ())
-> ST s ()
forall r ix e s.
Load r ix e =>
Scheduler s ()
-> Array r ix e
-> (Int -> e -> ST s ())
-> (Int -> Sz1 -> e -> ST s ())
-> ST s ()
iterArrayLinearWithSetST_ Scheduler s ()
forall s. Scheduler s ()
trivialScheduler_ Array r ix e
arr (MArray (PrimState (ST s)) r' Int e -> Int -> e -> ST s ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r' ix e -> Int -> e -> m ()
unsafeLinearWrite MVector s r' e
MArray (PrimState (ST s)) r' Int e
mvec') (MArray (PrimState (ST s)) r' Int e -> Int -> Sz1 -> e -> ST s ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> Sz1 -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r' ix e -> Int -> Sz1 -> e -> m ()
unsafeLinearSet MVector s r' e
MArray (PrimState (ST s)) r' Int e
mvec')
    MArray s r' ix e -> ST s (MArray s r' ix e)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MArray s r' ix e -> ST s (MArray s r' ix e))
-> MArray s r' ix e -> ST s (MArray s r' ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> MVector s r' e -> MArray s r' ix e
forall r e ix' ix s.
(Manifest r e, Index ix', Index ix) =>
Sz ix' -> MArray s r ix e -> MArray s r ix' e
forall ix' ix s.
(Index ix', Index ix) =>
Sz ix' -> MArray s r' ix e -> MArray s r' ix' e
unsafeResizeMArray Sz ix
sz MVector s r' e
mvec'
  {-# INLINE unsafeLoadIntoST #-}

  -- | Same as `unsafeLoadIntoST`, but respecting computation strategy.
  --
  -- @since 1.0.0
  unsafeLoadIntoIO
    :: Manifest r' e
    => MVector RealWorld r' e
    -> Array r ix e
    -> IO (MArray RealWorld r' ix e)
  unsafeLoadIntoIO MVector RealWorld r' e
mvec Array r ix e
arr = do
    let sz :: Sz ix
sz = Array r ix e -> Sz ix
forall e. Array r ix e -> Sz ix
forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array r ix e
arr
    MVector RealWorld r' e
mvec' <- MArray (PrimState IO) r' Int e
-> Sz1 -> IO (MArray (PrimState IO) r' Int e)
forall r e (f :: * -> *).
(Manifest r e, PrimMonad f) =>
MVector (PrimState f) r e -> Sz1 -> f (MVector (PrimState f) r e)
resizeMVector MVector RealWorld r' e
MArray (PrimState IO) r' Int e
mvec (Sz1 -> IO (MArray (PrimState IO) r' Int e))
-> Sz1 -> IO (MArray (PrimState IO) r' Int e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> Sz1
forall ix. Index ix => Sz ix -> Sz1
toLinearSz Sz ix
sz
    Comp -> (Scheduler RealWorld () -> IO ()) -> IO ()
withMassivScheduler_ (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Array r ix e
arr) ((Scheduler RealWorld () -> IO ()) -> IO ())
-> (Scheduler RealWorld () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld ()
scheduler ->
      ST RealWorld () -> IO ()
forall a. ST RealWorld a -> IO a
stToIO (ST RealWorld () -> IO ()) -> ST RealWorld () -> IO ()
forall a b. (a -> b) -> a -> b
$
        Scheduler RealWorld ()
-> Array r ix e
-> (Int -> e -> ST RealWorld ())
-> (Int -> Sz1 -> e -> ST RealWorld ())
-> ST RealWorld ()
forall s.
Scheduler s ()
-> Array r ix e
-> (Int -> e -> ST s ())
-> (Int -> Sz1 -> e -> ST s ())
-> ST s ()
forall r ix e s.
Load r ix e =>
Scheduler s ()
-> Array r ix e
-> (Int -> e -> ST s ())
-> (Int -> Sz1 -> e -> ST s ())
-> ST s ()
iterArrayLinearWithSetST_ Scheduler RealWorld ()
scheduler Array r ix e
arr (MArray (PrimState (ST RealWorld)) r' Int e
-> Int -> e -> ST RealWorld ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r' ix e -> Int -> e -> m ()
unsafeLinearWrite MVector RealWorld r' e
MArray (PrimState (ST RealWorld)) r' Int e
mvec') (MArray (PrimState (ST RealWorld)) r' Int e
-> Int -> Sz1 -> e -> ST RealWorld ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> Sz1 -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r' ix e -> Int -> Sz1 -> e -> m ()
unsafeLinearSet MVector RealWorld r' e
MArray (PrimState (ST RealWorld)) r' Int e
mvec')
    MArray RealWorld r' ix e -> IO (MArray RealWorld r' ix e)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MArray RealWorld r' ix e -> IO (MArray RealWorld r' ix e))
-> MArray RealWorld r' ix e -> IO (MArray RealWorld r' ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> MVector RealWorld r' e -> MArray RealWorld r' ix e
forall r e ix' ix s.
(Manifest r e, Index ix', Index ix) =>
Sz ix' -> MArray s r ix e -> MArray s r ix' e
forall ix' ix s.
(Index ix', Index ix) =>
Sz ix' -> MArray s r' ix e -> MArray s r' ix' e
unsafeResizeMArray Sz ix
sz MVector RealWorld r' e
mvec'
  {-# INLINE unsafeLoadIntoIO #-}

resizeMVector
  :: (Manifest r e, PrimMonad f)
  => MVector (PrimState f) r e
  -> Sz1
  -> f (MVector (PrimState f) r e)
resizeMVector :: forall r e (f :: * -> *).
(Manifest r e, PrimMonad f) =>
MVector (PrimState f) r e -> Sz1 -> f (MVector (PrimState f) r e)
resizeMVector MVector (PrimState f) r e
mvec Sz1
k =
  let mk :: Sz1
mk = MVector (PrimState f) r e -> Sz1
forall ix s. Index ix => MArray s r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MVector (PrimState f) r e
mvec
   in if Sz1
k Sz1 -> Sz1 -> Bool
forall a. Eq a => a -> a -> Bool
== Sz1
mk
        then MVector (PrimState f) r e -> f (MVector (PrimState f) r e)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector (PrimState f) r e
mvec
        else
          if Sz1
k Sz1 -> Sz1 -> Bool
forall a. Ord a => a -> a -> Bool
< Sz1
mk
            then MVector (PrimState f) r e -> Sz1 -> f (MVector (PrimState f) r e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e
-> Sz ix -> m (MArray (PrimState m) r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e
-> Sz ix -> m (MArray (PrimState m) r ix e)
unsafeLinearShrink MVector (PrimState f) r e
mvec Sz1
k
            else MVector (PrimState f) r e -> Sz1 -> f (MVector (PrimState f) r e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e
-> Sz ix -> m (MArray (PrimState m) r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e
-> Sz ix -> m (MArray (PrimState m) r ix e)
unsafeLinearGrow MVector (PrimState f) r e
mvec Sz1
k
{-# INLINE resizeMVector #-}

class Load r ix e => StrideLoad r ix e where
  -- | Load an array into memory with stride. Default implementation requires an instance of
  -- `Source`.
  iterArrayLinearWithStrideST_
    :: Scheduler s ()
    -> Stride ix
    -- ^ Stride to use
    -> Sz ix
    -- ^ Size of the target array affected by the stride.
    -> Array r ix e
    -- ^ Array that is being loaded
    -> (Int -> e -> ST s ())
    -- ^ Function that writes an element into target array
    -> ST s ()
  default iterArrayLinearWithStrideST_
    :: Source r e
    => Scheduler s ()
    -> Stride ix
    -> Sz ix
    -> Array r ix e
    -> (Int -> e -> ST s ())
    -> ST s ()
  iterArrayLinearWithStrideST_ Scheduler s ()
scheduler Stride ix
stride Sz ix
resultSize Array r ix e
arr =
    Scheduler s ()
-> Int -> (Int -> e) -> (Int -> e -> ST s ()) -> ST s ()
forall s (m :: * -> *) b.
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_ Scheduler s ()
scheduler (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
resultSize) Int -> e
unsafeLinearIndexWithStride
    where
      !strideIx :: ix
strideIx = Stride ix -> ix
forall ix. Stride ix -> ix
unStride Stride ix
stride
      unsafeLinearIndexWithStride :: Int -> e
unsafeLinearIndexWithStride =
        Array r ix e -> ix -> e
forall ix. Index ix => Array r ix e -> ix -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr (ix -> e) -> (Int -> ix) -> Int -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (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
(*) ix
strideIx (ix -> ix) -> (Int -> ix) -> Int -> ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
resultSize
      {-# INLINE unsafeLinearIndexWithStride #-}
  {-# INLINE iterArrayLinearWithStrideST_ #-}

-- class (Load r ix e) => StrideLoad r ix e where
-- class (Size r, StrideLoad r ix e) => StrideLoadP r ix e where
--
-- unsafeLoadIntoWithStrideST :: -- TODO: this would remove Size constraint and allow DS and LN instances for vectors.
--      Manifest r' ix e
--   => Array r ix e
--   -> Stride ix -- ^ Stride to use
--   -> MArray RealWorld r' ix e
--   -> m (MArray RealWorld r' ix e)

-- | Starting with massiv-1.0 `Mutable` and `Manifest` are synonymous. Since massiv-1.1
-- it is deprecated and will be removed in massiv-1.2
type Mutable r e = Manifest r e

{-# DEPRECATED Mutable "In favor of `Manifest`" #-}

-- | Manifest arrays are backed by actual memory and values are looked up versus
-- computed as it is with delayed arrays. Because manifest arrays are located in
-- memory their contents can be mutated once thawed into `MArray`. The process
-- of changed a mutable `MArray` back into an immutable `Array` is called
-- freezing.
class Source r e => Manifest r e where
  unsafeLinearIndexM :: Index ix => Array r ix e -> Int -> e

  -- | /O(1)/ - Get the size of a mutable array.
  --
  -- @since 1.0.0
  sizeOfMArray :: Index ix => MArray s r ix e -> Sz ix

  -- | /O(1)/ - Change the size of a mutable array. The actual number of
  -- elements should stay the same.
  --
  -- @since 1.0.0
  unsafeResizeMArray :: (Index ix', Index ix) => Sz ix' -> MArray s r ix e -> MArray s r ix' e

  -- | /O(1)/ - Take a linear slice out of a mutable array.
  --
  -- @since 1.0.0
  unsafeLinearSliceMArray :: Index ix => Ix1 -> Sz1 -> MArray s r ix e -> MVector s r e

  -- | Convert immutable array into a mutable array without copy.
  --
  -- @since 0.1.0
  unsafeThaw :: (Index ix, PrimMonad m) => Array r ix e -> m (MArray (PrimState m) r ix e)

  -- | Convert mutable array into an immutable array without copy.
  --
  -- @since 0.1.0
  unsafeFreeze :: (Index ix, PrimMonad m) => Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)

  -- | Create new mutable array, leaving it's elements uninitialized. Size isn't validated either.
  --
  -- @since 0.1.0
  unsafeNew :: (Index ix, PrimMonad m) => Sz ix -> m (MArray (PrimState m) r ix e)

  -- | Read an element at linear row-major index
  --
  -- @since 0.1.0
  unsafeLinearRead
    :: HAS_CALL_STACK
    => (Index ix, PrimMonad m)
    => MArray (PrimState m) r ix e
    -> Int
    -> m e

  -- | Write an element into mutable array with linear row-major index
  --
  -- @since 0.1.0
  unsafeLinearWrite
    :: HAS_CALL_STACK
    => (Index ix, PrimMonad m)
    => MArray (PrimState m) r ix e
    -> Int
    -> e
    -> m ()

  -- | Initialize mutable array to some default value.
  --
  -- @since 0.3.0
  initialize :: (Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> m ()

  -- | Create new mutable array while initializing all elements to some default value.
  --
  -- @since 0.3.0
  initializeNew :: (Index ix, PrimMonad m) => Maybe e -> Sz ix -> m (MArray (PrimState m) r ix e)
  initializeNew Maybe e
Nothing Sz ix
sz = Sz ix -> m (MArray (PrimState m) r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz m (MArray (PrimState m) r ix e)
-> (MArray (PrimState m) r ix e -> m (MArray (PrimState m) r ix e))
-> m (MArray (PrimState m) r ix e)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \MArray (PrimState m) r ix e
ma -> MArray (PrimState m) r ix e
ma MArray (PrimState m) r ix e
-> m () -> m (MArray (PrimState m) r ix e)
forall a b. a -> m b -> m a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ MArray (PrimState m) r ix e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> m ()
initialize MArray (PrimState m) r ix e
ma
  initializeNew (Just e
e) Sz ix
sz = Sz ix -> e -> m (MArray (PrimState m) r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> e -> m (MArray (PrimState m) r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Sz ix -> e -> m (MArray (PrimState m) r ix e)
newMArray Sz ix
sz e
e
  {-# INLINE initializeNew #-}

  -- | Create new mutable array while initializing all elements to the specified value.
  --
  -- @since 0.6.0
  newMArray :: (Index ix, PrimMonad m) => Sz ix -> e -> m (MArray (PrimState m) r ix e)
  newMArray Sz ix
sz e
e = do
    MArray (PrimState m) r ix e
marr <- Sz ix -> m (MArray (PrimState m) r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz
    MArray (PrimState m) r ix e
marr MArray (PrimState m) r ix e
-> m () -> m (MArray (PrimState m) r ix e)
forall a b. a -> m b -> m a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ MArray (PrimState m) r ix e -> Int -> Sz1 -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> Sz1 -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> Sz1 -> e -> m ()
unsafeLinearSet MArray (PrimState m) r ix e
marr Int
0 (Int -> Sz1
forall ix. ix -> Sz ix
SafeSz (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz)) e
e
  {-# INLINE newMArray #-}

  -- | Set all cells in the mutable array within the range to a specified value.
  --
  -- @since 0.3.0
  unsafeLinearSet
    :: HAS_CALL_STACK
    => (Index ix, PrimMonad m)
    => MArray (PrimState m) r ix e
    -> Ix1
    -> Sz1
    -> e
    -> m ()
  unsafeLinearSet MArray (PrimState m) r ix e
marr Int
offset Sz1
len e
e =
    Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
offset (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz1 -> Int
forall ix. Sz ix -> ix
unSz Sz1
len)) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (\Int
i -> MArray (PrimState m) r ix e -> Int -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix e
marr Int
i e
e)
  {-# INLINE unsafeLinearSet #-}

  -- | Copy part of one mutable array into another
  --
  -- @since 0.3.6
  unsafeLinearCopy
    :: HAS_CALL_STACK
    => (Index ix', Index ix, PrimMonad m)
    => MArray (PrimState m) r ix' e
    -- ^ Source mutable array
    -> Ix1
    -- ^ Starting index at source array
    -> MArray (PrimState m) r ix e
    -- ^ Target mutable array
    -> Ix1
    -- ^ Starting index at target array
    -> Sz1
    -- ^ Number of elements to copy
    -> m ()
  unsafeLinearCopy MArray (PrimState m) r ix' e
marrFrom Int
iFrom MArray (PrimState m) r ix e
marrTo Int
iTo (SafeSz Int
k) = do
    let delta :: Int
delta = Int
iTo Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
iFrom
    Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
iFrom (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
iFrom) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
      MArray (PrimState m) r ix' e -> Int -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MArray (PrimState m) r ix' e
marrFrom Int
i m e -> (e -> m ()) -> m ()
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MArray (PrimState m) r ix e -> Int -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix e
marrTo (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
delta)
  {-# INLINE unsafeLinearCopy #-}

  -- | Copy a part of a pure array into a mutable array
  --
  -- @since 0.3.6
  unsafeArrayLinearCopy
    :: HAS_CALL_STACK
    => (Index ix', Index ix, PrimMonad m)
    => Array r ix' e
    -- ^ Source pure array
    -> Ix1
    -- ^ Starting index at source array
    -> MArray (PrimState m) r ix e
    -- ^ Target mutable array
    -> Ix1
    -- ^ Starting index at target array
    -> Sz1
    -- ^ Number of elements to copy
    -> m ()
  unsafeArrayLinearCopy Array r ix' e
arrFrom Int
iFrom MArray (PrimState m) r ix e
marrTo Int
iTo (SafeSz Int
k) = do
    let delta :: Int
delta = Int
iTo Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
iFrom
    Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
iFrom (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
iFrom) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
      MArray (PrimState m) r ix e -> Int -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix e
marrTo (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
delta) (Array r ix' e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix' e
arrFrom Int
i)
  {-# INLINE unsafeArrayLinearCopy #-}

  -- | Linearly reduce the size of an array. Total number of elements should be smaller or
  -- equal. There is no guarantee that the original array is left unchanged, so it should
  -- no longer be used.
  --
  -- @since 0.3.6
  unsafeLinearShrink
    :: HAS_CALL_STACK
    => (Index ix, PrimMonad m)
    => MArray (PrimState m) r ix e
    -> Sz ix
    -> m (MArray (PrimState m) r ix e)
  unsafeLinearShrink = MArray (PrimState m) r ix e
-> Sz ix -> m (MArray (PrimState m) r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e
-> Sz ix -> m (MArray (PrimState m) r ix e)
unsafeDefaultLinearShrink
  {-# INLINE unsafeLinearShrink #-}

  -- | Linearly increase the size of an array. Total number of elements should be larger
  -- or equal. There is no guarantee that the original array is left unchanged, so it
  -- should no longer be used.
  --
  -- @since 0.3.6
  unsafeLinearGrow
    :: HAS_CALL_STACK
    => (Index ix, PrimMonad m)
    => MArray (PrimState m) r ix e
    -> Sz ix
    -> m (MArray (PrimState m) r ix e)
  unsafeLinearGrow MArray (PrimState m) r ix e
marr Sz ix
sz = do
    MArray (PrimState m) r ix e
marr' <- Sz ix -> m (MArray (PrimState m) r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz
    MArray (PrimState m) r ix e
-> Int -> MArray (PrimState m) r ix e -> Int -> Sz1 -> m ()
forall r e ix' ix (m :: * -> *).
(Manifest r e, Index ix', Index ix, PrimMonad m) =>
MArray (PrimState m) r ix' e
-> Int -> MArray (PrimState m) r ix e -> Int -> Sz1 -> m ()
forall ix' ix (m :: * -> *).
(Index ix', Index ix, PrimMonad m) =>
MArray (PrimState m) r ix' e
-> Int -> MArray (PrimState m) r ix e -> Int -> Sz1 -> m ()
unsafeLinearCopy MArray (PrimState m) r ix e
marr Int
0 MArray (PrimState m) r ix e
marr' Int
0 (Sz1 -> m ()) -> Sz1 -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Sz1
forall ix. ix -> Sz ix
SafeSz (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (MArray (PrimState m) r ix e -> Sz ix
forall ix s. Index ix => MArray s r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr))
    MArray (PrimState m) r ix e -> m (MArray (PrimState m) r ix e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MArray (PrimState m) r ix e
marr'
  {-# INLINE unsafeLinearGrow #-}

unsafeDefaultLinearShrink
  :: HAS_CALL_STACK
  => (Manifest r e, Index ix, PrimMonad m)
  => MArray (PrimState m) r ix e
  -> Sz ix
  -> m (MArray (PrimState m) r ix e)
unsafeDefaultLinearShrink :: forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e
-> Sz ix -> m (MArray (PrimState m) r ix e)
unsafeDefaultLinearShrink MArray (PrimState m) r ix e
marr Sz ix
sz = do
  MArray (PrimState m) r ix e
marr' <- Sz ix -> m (MArray (PrimState m) r ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz ix
sz
  MArray (PrimState m) r ix e
-> Int -> MArray (PrimState m) r ix e -> Int -> Sz1 -> m ()
forall r e ix' ix (m :: * -> *).
(Manifest r e, Index ix', Index ix, PrimMonad m) =>
MArray (PrimState m) r ix' e
-> Int -> MArray (PrimState m) r ix e -> Int -> Sz1 -> m ()
forall ix' ix (m :: * -> *).
(Index ix', Index ix, PrimMonad m) =>
MArray (PrimState m) r ix' e
-> Int -> MArray (PrimState m) r ix e -> Int -> Sz1 -> m ()
unsafeLinearCopy MArray (PrimState m) r ix e
marr Int
0 MArray (PrimState m) r ix e
marr' Int
0 (Sz1 -> m ()) -> Sz1 -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Sz1
forall ix. ix -> Sz ix
SafeSz (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz)
  MArray (PrimState m) r ix e -> m (MArray (PrimState m) r ix e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MArray (PrimState m) r ix e
marr'
{-# INLINE unsafeDefaultLinearShrink #-}

-- | Read an array element
--
-- @since 0.1.0
unsafeRead
  :: HAS_CALL_STACK
  => (Manifest r e, Index ix, PrimMonad m)
  => MArray (PrimState m) r ix e
  -> ix
  -> m e
unsafeRead :: forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> m e
unsafeRead MArray (PrimState m) r ix e
marr = MArray (PrimState m) r ix e -> Int -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MArray (PrimState m) r ix e
marr (Int -> m e) -> (ix -> Int) -> ix -> m e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex (MArray (PrimState m) r ix e -> Sz ix
forall ix s. Index ix => MArray s r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr)
{-# INLINE unsafeRead #-}

-- | Write an element into array
--
-- @since 0.1.0
unsafeWrite
  :: HAS_CALL_STACK
  => (Manifest r e, Index ix, PrimMonad m)
  => MArray (PrimState m) r ix e
  -> ix
  -> e
  -> m ()
unsafeWrite :: forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> e -> m ()
unsafeWrite MArray (PrimState m) r ix e
marr = MArray (PrimState m) r ix e -> Int -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix e
marr (Int -> e -> m ()) -> (ix -> Int) -> ix -> e -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex (MArray (PrimState m) r ix e -> Sz ix
forall ix s. Index ix => MArray s r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr)
{-# INLINE unsafeWrite #-}

-- | Modify an element in the array with a monadic action. Returns the previous value.
--
-- @since 0.4.0
unsafeLinearModify
  :: HAS_CALL_STACK
  => (Manifest r e, Index ix, PrimMonad m)
  => MArray (PrimState m) r ix e
  -> (e -> m e)
  -> Int
  -> m e
unsafeLinearModify :: forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> (e -> m e) -> Int -> m e
unsafeLinearModify !MArray (PrimState m) r ix e
marr e -> m e
f !Int
i = do
  e
v <- MArray (PrimState m) r ix e -> Int -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MArray (PrimState m) r ix e
marr Int
i
  e
v' <- e -> m e
f e
v
  MArray (PrimState m) r ix e -> Int -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix e
marr Int
i e
v'
  e -> m e
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure e
v
{-# INLINE unsafeLinearModify #-}

-- | Modify an element in the array with a monadic action. Returns the previous value.
--
-- @since 0.4.0
unsafeModify
  :: HAS_CALL_STACK
  => (Manifest r e, Index ix, PrimMonad m)
  => MArray (PrimState m) r ix e
  -> (e -> m e)
  -> ix
  -> m e
unsafeModify :: forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> (e -> m e) -> ix -> m e
unsafeModify MArray (PrimState m) r ix e
marr e -> m e
f ix
ix = MArray (PrimState m) r ix e -> (e -> m e) -> Int -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> (e -> m e) -> Int -> m e
unsafeLinearModify MArray (PrimState m) r ix e
marr e -> m e
f (Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex (MArray (PrimState m) r ix e -> Sz ix
forall ix s. Index ix => MArray s r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr) ix
ix)
{-# INLINE unsafeModify #-}

-- | Swap two elements in a mutable array under the supplied indices. Returns the previous
-- values.
--
-- @since 0.4.0
unsafeSwap
  :: HAS_CALL_STACK
  => (Manifest r e, Index ix, PrimMonad m)
  => MArray (PrimState m) r ix e
  -> ix
  -> ix
  -> m (e, e)
unsafeSwap :: forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> ix -> ix -> m (e, e)
unsafeSwap !MArray (PrimState m) r ix e
marr !ix
ix1 !ix
ix2 = MArray (PrimState m) r ix e -> Int -> Int -> m (e, e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> Int -> m (e, e)
unsafeLinearSwap MArray (PrimState m) r ix e
marr (Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
sz ix
ix1) (Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
sz ix
ix2)
  where
    sz :: Sz ix
sz = MArray (PrimState m) r ix e -> Sz ix
forall ix s. Index ix => MArray s r ix e -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) r ix e
marr
{-# INLINE unsafeSwap #-}

-- | Swap two elements in a mutable array under the supplied linear indices. Returns the
-- previous values.
--
-- @since 0.4.0
unsafeLinearSwap
  :: HAS_CALL_STACK
  => (Manifest r e, Index ix, PrimMonad m)
  => MArray (PrimState m) r ix e
  -> Int
  -> Int
  -> m (e, e)
unsafeLinearSwap :: forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> Int -> m (e, e)
unsafeLinearSwap !MArray (PrimState m) r ix e
marr !Int
i1 !Int
i2 = do
  e
val1 <- MArray (PrimState m) r ix e -> Int -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MArray (PrimState m) r ix e
marr Int
i1
  e
val2 <- MArray (PrimState m) r ix e -> Int -> m e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MArray (PrimState m) r ix e
marr Int
i2
  MArray (PrimState m) r ix e -> Int -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix e
marr Int
i1 e
val2
  MArray (PrimState m) r ix e -> Int -> e -> m ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray (PrimState m) r ix e
marr Int
i2 e
val1
  (e, e) -> m (e, e)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (e
val1, e
val2)
{-# INLINE unsafeLinearSwap #-}

class (IsList (Array r ix e), Load r ix e) => Ragged r ix e where
  generateRaggedM :: Monad m => Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)

  flattenRagged :: Array r ix e -> Vector r e

  loadRaggedST
    :: Scheduler s () -> Array r ix e -> (Ix1 -> e -> ST s ()) -> Ix1 -> Ix1 -> Sz ix -> ST s ()

  raggedFormat :: (e -> String) -> String -> Array r ix e -> String

-- | Create an Array with no elements. By itself it is not particularly useful, but it serves as a
-- nice base for constructing larger arrays.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array as A
-- >>> :set -XTypeApplications
-- >>> xs = empty @DL @Ix1 @Double
-- >>> snoc (cons 4 (cons 5 xs)) 22
-- Array DL Seq (Sz1 3)
--   [ 4.0, 5.0, 22.0 ]
--
-- @since 0.3.0
empty
  :: forall r ix e
   . Load r ix e
  => Array r ix e
empty :: forall r ix e. Load r ix e => Array r ix e
empty = Comp -> Sz ix -> (ix -> e) -> Array r ix e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
Seq Sz ix
forall ix. Index ix => Sz ix
zeroSz (e -> ix -> e
forall a b. a -> b -> a
const (Uninitialized -> e
forall e a. (HasCallStack, Exception e) => e -> a
throwImpossible Uninitialized
Uninitialized))
{-# INLINE empty #-}

-- | Create an Array with a single element.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array as A
-- >>> singleton 7 :: Array D Ix4 Double
-- Array D Seq (Sz (1 :> 1 :> 1 :. 1))
--   [ [ [ [ 7.0 ]
--       ]
--     ]
--   ]
--
-- Instead of specifying type signature we could use @TypeApplications@
--
-- >>> :set -XTypeApplications
-- >>> singleton @U @Ix4 @Double 7
-- Array U Seq (Sz (1 :> 1 :> 1 :. 1))
--   [ [ [ [ 7.0 ]
--       ]
--     ]
--   ]
--
-- @since 0.1.0
singleton
  :: forall r ix e
   . Load r ix e
  => e
  -- ^ The only element
  -> Array r ix e
singleton :: forall r ix e. Load r ix e => e -> Array r ix e
singleton = Comp -> Sz ix -> (ix -> e) -> Array r ix e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
Seq Sz ix
forall ix. Index ix => Sz ix
oneSz ((ix -> e) -> Array r ix e) -> (e -> ix -> e) -> e -> Array r ix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> ix -> e
forall a b. a -> b -> a
const
{-# INLINE singleton #-}

infixl 4 !, !?, ??

-- | /O(1)/ - Infix version of 'index''.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array as A
-- >>> a = computeAs U $ iterateN (Sz (2 :. 3)) succ (0 :: Int)
-- >>> a
-- Array U Seq (Sz (2 :. 3))
--   [ [ 1, 2, 3 ]
--   , [ 4, 5, 6 ]
--   ]
-- >>> a ! 0 :. 2
-- 3
--
-- @since 0.1.0
(!)
  :: forall r ix e
   . (HasCallStack, Manifest r e, Index ix)
  => Array r ix e
  -> ix
  -> e
! :: forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
(!) Array r ix e
arr = Either SomeException e -> e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException e -> e)
-> (ix -> Either SomeException e) -> ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> ix -> Either SomeException e
forall ix r e (m :: * -> *).
(Index ix, Source r e, MonadThrow m) =>
Array r ix e -> ix -> m e
evaluateM Array r ix e
arr
{-# INLINE (!) #-}

-- | /O(1)/ - Infix version of `indexM`.
--
-- /__Exceptions__/: `IndexOutOfBoundsException`
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array as A
-- >>> :set -XTypeApplications
-- >>> a <- fromListsM @U @Ix2 @Int Seq [[1,2,3],[4,5,6]]
-- >>> a
-- Array U Seq (Sz (2 :. 3))
--   [ [ 1, 2, 3 ]
--   , [ 4, 5, 6 ]
--   ]
-- >>> a !? 0 :. 2
-- 3
-- >>> a !? 0 :. 3
-- *** Exception: IndexOutOfBoundsException: (0 :. 3) is not safe for (Sz (2 :. 3))
-- >>> a !? 0 :. 3 :: Maybe Int
-- Nothing
--
-- @since 0.1.0
(!?)
  :: forall r ix e m
   . (Index ix, Manifest r e, MonadThrow m)
  => Array r ix e
  -> ix
  -> m e
!? :: forall r ix e (m :: * -> *).
(Index ix, Manifest r e, MonadThrow m) =>
Array r ix e -> ix -> m e
(!?) = Array r ix e -> ix -> m e
forall ix r e (m :: * -> *).
(Index ix, Manifest r e, MonadThrow m) =>
Array r ix e -> ix -> m e
indexM
{-# INLINE (!?) #-}

-- | /O(1)/ - Lookup an element in the array, where array itself is wrapped with
-- `MonadThrow`. This operator is useful when used together with slicing or other
-- functions that can fail.
--
-- /__Exceptions__/: `IndexOutOfBoundsException`
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array as A
-- >>> :set -XTypeApplications
-- >>> ma = fromListsM @U @Ix3 @Int @Maybe Seq [[[1,2,3]],[[4,5,6]]]
-- >>> ma
-- Just (Array U Seq (Sz (2 :> 1 :. 3))
--   [ [ [ 1, 2, 3 ]
--     ]
--   , [ [ 4, 5, 6 ]
--     ]
--   ]
-- )
-- >>> ma ??> 1
-- Just (Array U Seq (Sz (1 :. 3))
--   [ [ 4, 5, 6 ]
--   ]
-- )
-- >>> ma ??> 1 ?? 0 :. 2
-- Just 6
-- >>> ma ?? 1 :> 0 :. 2
-- Just 6
--
-- @since 0.1.0
(??) :: (Index ix, Manifest r e, MonadThrow m) => m (Array r ix e) -> ix -> m e
?? :: forall ix r e (m :: * -> *).
(Index ix, Manifest r e, MonadThrow m) =>
m (Array r ix e) -> ix -> m e
(??) m (Array r ix e)
marr ix
ix = m (Array r ix e)
marr m (Array r ix e) -> (Array r ix e -> m e) -> m e
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Array r ix e -> ix -> m e
forall r ix e (m :: * -> *).
(Index ix, Manifest r e, MonadThrow m) =>
Array r ix e -> ix -> m e
!? ix
ix)
{-# INLINE (??) #-}

-- | /O(1)/ - Lookup an element in the array. Returns `Nothing`, when index is out of bounds and
-- returns the element at the supplied index otherwise. Use `indexM` instead, since it is more
-- general and it can just as well be used with `Maybe`.
--
-- @since 0.1.0
index :: (Index ix, Manifest r e) => Array r ix e -> ix -> Maybe e
index :: forall ix r e.
(Index ix, Manifest r e) =>
Array r ix e -> ix -> Maybe e
index = Array r ix e -> ix -> Maybe e
forall ix r e (m :: * -> *).
(Index ix, Manifest r e, MonadThrow m) =>
Array r ix e -> ix -> m e
indexM
{-# INLINE index #-}

-- | /O(1)/ - Lookup an element in the array.
--
-- /__Exceptions__/: `IndexOutOfBoundsException`
--
-- @since 0.3.0
indexM :: (Index ix, Manifest r e, MonadThrow m) => Array r ix e -> ix -> m e
indexM :: forall ix r e (m :: * -> *).
(Index ix, Manifest r e, MonadThrow m) =>
Array r ix e -> ix -> m e
indexM = Array r ix e -> ix -> m e
forall ix r e (m :: * -> *).
(Index ix, Source r e, MonadThrow m) =>
Array r ix e -> ix -> m e
evaluateM
{-# INLINE indexM #-}

-- | /O(1)/ - Lookup an element in the array, while using default element when index is out of
-- bounds.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> :set -XOverloadedLists
-- >>> xs = [0..100] :: Array P Ix1 Int
-- >>> defaultIndex 999 xs 100
-- 100
-- >>> defaultIndex 999 xs 101
-- 999
--
-- @since 0.1.0
defaultIndex :: (Index ix, Manifest r e) => e -> Array r ix e -> ix -> e
defaultIndex :: forall ix r e.
(Index ix, Manifest r e) =>
e -> Array r ix e -> ix -> e
defaultIndex e
defVal = Border e -> Array r ix e -> ix -> e
forall ix r e.
(Index ix, Manifest r e) =>
Border e -> Array r ix e -> ix -> e
borderIndex (e -> Border e
forall e. e -> Border e
Fill e
defVal)
{-# INLINE defaultIndex #-}

-- | /O(1)/ - Lookup an element in the array. Use a border resolution technique
-- when index is out of bounds.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array as A
-- >>> :set -XOverloadedLists
-- >>> xs = [0..100] :: Array U Ix1 Int
-- >>> borderIndex Wrap xs <$> range Seq 99 104
-- Array D Seq (Sz1 5)
--   [ 99, 100, 0, 1, 2 ]
--
-- @since 0.1.0
borderIndex :: (Index ix, Manifest r e) => Border e -> Array r ix e -> ix -> e
borderIndex :: forall ix r e.
(Index ix, Manifest r e) =>
Border e -> Array r ix e -> ix -> e
borderIndex Border e
border Array r ix e
arr = Border e -> Sz ix -> (ix -> e) -> ix -> e
forall ix e. Index ix => Border e -> Sz ix -> (ix -> e) -> ix -> e
handleBorderIndex Border e
border (Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
arr) (Array r ix e -> ix -> e
forall ix. Index ix => Array r ix e -> ix -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr)
{-# INLINE borderIndex #-}

-- | /O(1)/ - Lookup an element in the array. This is a partial function and it will throw
-- an error when index is out of bounds. It is safer to use `indexM` instead.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> :set -XOverloadedLists
-- >>> xs = [0..100] :: Array U Ix1 Int
-- >>> index' xs 50
-- 50
--
-- @since 0.1.0
index' :: (HasCallStack, Index ix, Manifest r e) => Array r ix e -> ix -> e
index' :: forall ix r e.
(HasCallStack, Index ix, Manifest r e) =>
Array r ix e -> ix -> e
index' Array r ix e
arr ix
ix = Either SomeException e -> e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Array r ix e -> ix -> Either SomeException e
forall ix r e (m :: * -> *).
(Index ix, Source r e, MonadThrow m) =>
Array r ix e -> ix -> m e
evaluateM Array r ix e
arr ix
ix)
{-# INLINE index' #-}

-- | This is just like `indexM` function, but it allows getting values from
-- delayed arrays as well as `Manifest`. As the name suggests, indexing into a
-- delayed array at the same index multiple times will cause evaluation of the
-- value each time and can destroy the performace if used without care.
--
-- ==== __Examples__
--
-- >>> import Control.Exception
-- >>> import Data.Massiv.Array
-- >>> evaluateM (range Seq (Ix2 10 20) (100 :. 210)) 50 :: Either SomeException Ix2
-- Right (60 :. 70)
-- >>> evaluateM (range Seq (Ix2 10 20) (100 :. 210)) 150 :: Either SomeException Ix2
-- Left (IndexOutOfBoundsException: (150 :. 150) is not safe for (Sz (90 :. 190)))
--
-- @since 0.3.0
evaluateM :: (Index ix, Source r e, MonadThrow m) => Array r ix e -> ix -> m e
evaluateM :: forall ix r e (m :: * -> *).
(Index ix, Source r e, MonadThrow m) =>
Array r ix e -> ix -> m e
evaluateM Array r ix e
arr ix
ix
  | Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
arr) ix
ix = e -> m e
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array r ix e -> ix -> e
forall ix. Index ix => Array r ix e -> ix -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr ix
ix)
  | Bool
otherwise = IndexException -> m e
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (Sz ix -> ix -> IndexException
forall ix. Index ix => Sz ix -> ix -> IndexException
IndexOutOfBoundsException (Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
arr) ix
ix)
{-# INLINE evaluateM #-}

-- | Similar to `evaluateM`, but will throw an error on out of bounds indices.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> evaluate' (range Seq (Ix2 10 20) (100 :. 210)) 50
-- 60 :. 70
--
-- @since 0.3.0
evaluate' :: (HasCallStack, Index ix, Source r e) => Array r ix e -> ix -> e
evaluate' :: forall ix r e.
(HasCallStack, Index ix, Source r e) =>
Array r ix e -> ix -> e
evaluate' Array r ix e
arr ix
ix = Either SomeException e -> e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Array r ix e -> ix -> Either SomeException e
forall ix r e (m :: * -> *).
(Index ix, Source r e, MonadThrow m) =>
Array r ix e -> ix -> m e
evaluateM Array r ix e
arr ix
ix)
{-# INLINE evaluate' #-}

-- | /O(1)/ - Check if array has elements.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> isNotNull (singleton 1 :: Array D Ix2 Int)
-- True
-- >>> isNotNull (empty :: Array D Ix2 Int)
-- False
--
-- @since 0.5.1
isNotNull :: Shape r ix => Array r ix e -> Bool
isNotNull :: forall r ix e. Shape r ix => Array r ix e -> Bool
isNotNull = Bool -> Bool
not (Bool -> Bool) -> (Array r ix e -> Bool) -> Array r ix e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> Bool
forall e. Array r ix e -> Bool
forall r ix e. Shape r ix => Array r ix e -> Bool
isNull
{-# INLINE isNotNull #-}

-- | /O(1)/ - Check if array has elements.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> isEmpty (singleton 1 :: Array D Ix2 Int)
-- False
-- >>> isEmpty (empty :: Array D Ix2 Int)
-- True
--
-- @since 1.0.0
isEmpty :: (Index ix, Size r) => Array r ix e -> Bool
isEmpty :: forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
isEmpty = (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (Int -> Bool) -> (Array r ix e -> Int) -> Array r ix e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> Int
forall ix r e. (Index ix, Size r) => Array r ix e -> Int
elemsCount
{-# INLINE isEmpty #-}

-- | /O(1)/ - Check if array has elements.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> isNotEmpty (singleton 1 :: Array D Ix2 Int)
-- True
-- >>> isNotEmpty (empty :: Array D Ix2 Int)
-- False
--
-- @since 1.0.0
isNotEmpty :: (Index ix, Size r) => Array r ix e -> Bool
isNotEmpty :: forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
isNotEmpty = Bool -> Bool
not (Bool -> Bool) -> (Array r ix e -> Bool) -> Array r ix e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> Bool
forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
isEmpty
{-# INLINE isNotEmpty #-}

-- | /O(1)/ - Get the number of elements in the array.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> elemsCount $ range Seq (Ix1 10) 15
-- 5
--
-- @since 0.1.0
elemsCount :: (Index ix, Size r) => Array r ix e -> Int
elemsCount :: forall ix r e. (Index ix, Size r) => Array r ix e -> Int
elemsCount = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (Sz ix -> Int) -> (Array r ix e -> Sz ix) -> Array r ix e -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size
{-# INLINE elemsCount #-}

inline0 :: (a -> b) -> a -> b
inline0 :: forall a b. (a -> b) -> a -> b
inline0 a -> b
f = a -> b
f
{-# INLINE [0] inline0 #-}

inline1 :: (a -> b) -> a -> b
inline1 :: forall a b. (a -> b) -> a -> b
inline1 a -> b
f = a -> b
f
{-# INLINE [1] inline1 #-}

inline2 :: (a -> b) -> a -> b
inline2 :: forall a b. (a -> b) -> a -> b
inline2 a -> b
f = a -> b
f
{-# INLINE [2] inline2 #-}