{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
module Data.Massiv.Array.Ops.Slice
(
(!>)
, (!?>)
, (??>)
, (<!)
, (<!?)
, (<??)
, (<!>)
, (<!?>)
, (<??>)
, outerSlices
, innerSlices
, withinSlices
, withinSlicesM
) where
import Control.Monad (unless)
import Data.Massiv.Array.Delayed.Pull
import Data.Massiv.Core.Common
infixl 4 !>, !?>, ??>, <!, <!?, <??, <!>, <!?>, <??>
(!>) :: OuterSlice r ix e => Array r ix e -> Int -> Elt r ix e
!> :: Array r ix e -> Int -> Elt r ix e
(!>) !Array r ix e
arr !Int
ix = (SomeException -> Elt r ix e)
-> (Elt r ix e -> Elt r ix e)
-> Either SomeException (Elt r ix e)
-> Elt r ix e
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Elt r ix e
forall a e. Exception e => e -> a
throw Elt r ix e -> Elt r ix e
forall a. a -> a
id (Array r ix e
arr Array r ix e -> Int -> Either SomeException (Elt r ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, OuterSlice r ix e) =>
Array r ix e -> Int -> m (Elt r ix e)
!?> Int
ix)
{-# INLINE (!>) #-}
(!?>) :: (MonadThrow m, OuterSlice r ix e) => Array r ix e -> Int -> m (Elt r ix e)
!?> :: Array r ix e -> Int -> m (Elt r ix e)
(!?>) !Array r ix e
arr !Int
i
| Sz Int -> Int -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz Int
sz Int
i = Elt r ix e -> m (Elt r ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Elt r ix e -> m (Elt r ix e)) -> Elt r ix e -> m (Elt r ix e)
forall a b. (a -> b) -> a -> b
$ Array r ix e -> Int -> Elt r ix e
forall r ix e.
OuterSlice r ix e =>
Array r ix e -> Int -> Elt r ix e
unsafeOuterSlice Array r ix e
arr Int
i
| Bool
otherwise = IndexException -> m (Elt r ix e)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (IndexException -> m (Elt r ix e))
-> IndexException -> m (Elt r ix e)
forall a b. (a -> b) -> a -> b
$ Sz Int -> Int -> IndexException
forall ix. Index ix => Sz ix -> ix -> IndexException
IndexOutOfBoundsException Sz Int
sz Int
i
where
!sz :: Sz Int
sz = (Sz Int, Sz (Lower ix)) -> Sz Int
forall a b. (a, b) -> a
fst (Sz ix -> (Sz Int, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz (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))
{-# INLINE (!?>) #-}
(??>) :: (MonadThrow m, OuterSlice r ix e) => m (Array r ix e) -> Int -> m (Elt r ix e)
??> :: m (Array r ix e) -> Int -> m (Elt r ix e)
(??>) m (Array r ix e)
marr !Int
ix = m (Array r ix e)
marr m (Array r ix e)
-> (Array r ix e -> m (Elt r ix e)) -> m (Elt r ix e)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Array r ix e -> Int -> m (Elt r ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, OuterSlice r ix e) =>
Array r ix e -> Int -> m (Elt r ix e)
!?> Int
ix)
{-# INLINE (??>) #-}
(<!?) :: (MonadThrow m, InnerSlice r ix e) => Array r ix e -> Int -> m (Elt r ix e)
<!? :: Array r ix e -> Int -> m (Elt r ix e)
(<!?) !Array r ix e
arr !Int
i
| Sz Int -> Int -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz Int
m Int
i = Elt r ix e -> m (Elt r ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Elt r ix e -> m (Elt r ix e)) -> Elt r ix e -> m (Elt r ix e)
forall a b. (a -> b) -> a -> b
$ Array r ix e -> (Sz (Lower ix), Sz Int) -> Int -> Elt r ix e
forall r ix e.
InnerSlice r ix e =>
Array r ix e -> (Sz (Lower ix), Sz Int) -> Int -> Elt r ix e
unsafeInnerSlice Array r ix e
arr (Sz (Lower ix), Sz Int)
sz Int
i
| Bool
otherwise = IndexException -> m (Elt r ix e)
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (IndexException -> m (Elt r ix e))
-> IndexException -> m (Elt r ix e)
forall a b. (a -> b) -> a -> b
$ Sz Int -> Int -> IndexException
forall ix. Index ix => Sz ix -> ix -> IndexException
IndexOutOfBoundsException Sz Int
m Int
i
where
!sz :: (Sz (Lower ix), Sz Int)
sz@(Sz (Lower ix)
_, Sz Int
m) = Sz ix -> (Sz (Lower ix), Sz Int)
forall ix. Index ix => Sz ix -> (Sz (Lower ix), Sz Int)
unsnocSz (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)
{-# INLINE (<!?) #-}
(<!) :: InnerSlice r ix e => Array r ix e -> Int -> Elt r ix e
<! :: Array r ix e -> Int -> Elt r ix e
(<!) !Array r ix e
arr !Int
ix =
case Array r ix e
arr Array r ix e -> Int -> Either SomeException (Elt r ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, InnerSlice r ix e) =>
Array r ix e -> Int -> m (Elt r ix e)
<!? Int
ix of
Right res -> Elt r ix e
res
Left exc -> SomeException -> Elt r ix e
forall a e. Exception e => e -> a
throw SomeException
exc
{-# INLINE (<!) #-}
(<??) :: (MonadThrow m, InnerSlice r ix e) => m (Array r ix e) -> Int -> m (Elt r ix e)
<?? :: m (Array r ix e) -> Int -> m (Elt r ix e)
(<??) m (Array r ix e)
marr !Int
ix = m (Array r ix e)
marr m (Array r ix e)
-> (Array r ix e -> m (Elt r ix e)) -> m (Elt r ix e)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Array r ix e -> Int -> m (Elt r ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, InnerSlice r ix e) =>
Array r ix e -> Int -> m (Elt r ix e)
<!? Int
ix)
{-# INLINE (<??) #-}
(<!?>) :: (MonadThrow m, Slice r ix e) => Array r ix e -> (Dim, Int) -> m (Elt r ix e)
<!?> :: Array r ix e -> (Dim, Int) -> m (Elt r ix e)
(<!?>) !Array r ix e
arr (Dim
dim, Int
i) = do
(Sz Int
m, Sz (Lower ix)
szl) <- Sz ix -> Dim -> m (Sz Int, Sz (Lower ix))
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz ix -> Dim -> m (Sz Int, Sz (Lower ix))
pullOutSzM (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) Dim
dim
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Sz Int -> Int -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz Int
m Int
i) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ IndexException -> m ()
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (IndexException -> m ()) -> IndexException -> m ()
forall a b. (a -> b) -> a -> b
$ Sz Int -> Int -> IndexException
forall ix. Index ix => Sz ix -> ix -> IndexException
IndexOutOfBoundsException Sz Int
m Int
i
Sz ix
cutSz <- Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
insertSzM Sz (Lower ix)
szl Dim
dim Sz Int
forall ix. Index ix => Sz ix
oneSz
Dim -> Sz ix -> Array r ix e -> Int -> m (Elt r ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Slice r ix e) =>
Dim -> Sz ix -> Array r ix e -> Int -> m (Elt r ix e)
internalInnerSlice Dim
dim Sz ix
cutSz Array r ix e
arr Int
i
{-# INLINE (<!?>) #-}
internalInnerSlice ::
(MonadThrow m, Slice r ix e) => Dim -> Sz ix -> Array r ix e -> Int -> m (Elt r ix e)
internalInnerSlice :: Dim -> Sz ix -> Array r ix e -> Int -> m (Elt r ix e)
internalInnerSlice Dim
dim Sz ix
cutSz Array r ix e
arr Int
i = do
ix
start <- ix -> Dim -> Int -> m ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
forall ix. Index ix => ix
zeroIndex Dim
dim Int
i
Array r ix e -> ix -> Sz ix -> Dim -> m (Elt r ix e)
forall r ix e (m :: * -> *).
(Slice r ix e, MonadThrow m) =>
Array r ix e -> ix -> Sz ix -> Dim -> m (Elt r ix e)
unsafeSlice Array r ix e
arr ix
start Sz ix
cutSz Dim
dim
{-# INLINE internalInnerSlice #-}
(<!>) :: Slice r ix e => Array r ix e -> (Dim, Int) -> Elt r ix e
<!> :: Array r ix e -> (Dim, Int) -> Elt r ix e
(<!>) !Array r ix e
arr !(Dim, Int)
dix =
case Array r ix e
arr Array r ix e -> (Dim, Int) -> Either SomeException (Elt r ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Slice r ix e) =>
Array r ix e -> (Dim, Int) -> m (Elt r ix e)
<!?> (Dim, Int)
dix of
Right res -> Elt r ix e
res
Left exc -> SomeException -> Elt r ix e
forall a e. Exception e => e -> a
throw SomeException
exc
{-# INLINE (<!>) #-}
(<??>) :: (MonadThrow m, Slice r ix e) => m (Array r ix e) -> (Dim, Int) -> m (Elt r ix e)
<??> :: m (Array r ix e) -> (Dim, Int) -> m (Elt r ix e)
(<??>) !m (Array r ix e)
marr !(Dim, Int)
ix = m (Array r ix e)
marr m (Array r ix e)
-> (Array r ix e -> m (Elt r ix e)) -> m (Elt r ix e)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Array r ix e -> (Dim, Int) -> m (Elt r ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Slice r ix e) =>
Array r ix e -> (Dim, Int) -> m (Elt r ix e)
<!?> (Dim, Int)
ix)
{-# INLINE (<??>) #-}
outerSlices :: OuterSlice r ix e => Array r ix e -> Array D Ix1 (Elt r ix e)
outerSlices :: Array r ix e -> Array D Int (Elt r ix e)
outerSlices Array r ix e
arr = Comp -> Sz Int -> (Int -> Elt r ix e) -> Array D Int (Elt r ix e)
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
Seq Sz Int
k (Array r ix e -> Int -> Elt r ix e
forall r ix e.
OuterSlice r ix e =>
Array r ix e -> Int -> Elt r ix e
unsafeOuterSlice Array r ix e
arr)
where
(Sz Int
k, Sz (Lower ix)
_) = Sz ix -> (Sz Int, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz (Sz ix -> (Sz Int, Sz (Lower ix)))
-> Sz ix -> (Sz Int, Sz (Lower ix))
forall a b. (a -> b) -> a -> b
$ 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
{-# INLINE outerSlices #-}
innerSlices :: InnerSlice r ix e => Array r ix e -> Array D Ix1 (Elt r ix e)
innerSlices :: Array r ix e -> Array D Int (Elt r ix e)
innerSlices Array r ix e
arr = Comp -> Sz Int -> (Int -> Elt r ix e) -> Array D Int (Elt r ix e)
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
Seq Sz Int
k (Array r ix e -> (Sz (Lower ix), Sz Int) -> Int -> Elt r ix e
forall r ix e.
InnerSlice r ix e =>
Array r ix e -> (Sz (Lower ix), Sz Int) -> Int -> Elt r ix e
unsafeInnerSlice Array r ix e
arr (Sz (Lower ix), Sz Int)
sz)
where
sz :: (Sz (Lower ix), Sz Int)
sz@(Sz (Lower ix)
_, Sz Int
k) = Sz ix -> (Sz (Lower ix), Sz Int)
forall ix. Index ix => Sz ix -> (Sz (Lower ix), Sz Int)
unsnocSz (Sz ix -> (Sz (Lower ix), Sz Int))
-> Sz ix -> (Sz (Lower ix), Sz Int)
forall a b. (a -> b) -> a -> b
$ 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
{-# INLINE innerSlices #-}
withinSlices ::
(IsIndexDimension ix n, Slice r ix e)
=> Dimension n
-> Array r ix e
-> Array D Ix1 (Elt r ix e)
withinSlices :: Dimension n -> Array r ix e -> Array D Int (Elt r ix e)
withinSlices Dimension n
dim = (SomeException -> Array D Int (Elt r ix e))
-> (Array D Int (Elt r ix e) -> Array D Int (Elt r ix e))
-> Either SomeException (Array D Int (Elt r ix e))
-> Array D Int (Elt r ix e)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array D Int (Elt r ix e)
forall e a. Exception e => e -> a
throwImpossible Array D Int (Elt r ix e) -> Array D Int (Elt r ix e)
forall a. a -> a
id (Either SomeException (Array D Int (Elt r ix e))
-> Array D Int (Elt r ix e))
-> (Array r ix e
-> Either SomeException (Array D Int (Elt r ix e)))
-> Array r ix e
-> Array D Int (Elt r ix e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim
-> Array r ix e -> Either SomeException (Array D Int (Elt r ix e))
forall (m :: * -> *) r ix e.
(MonadThrow m, Slice r ix e) =>
Dim -> Array r ix e -> m (Array D Int (Elt r ix e))
withinSlicesM (Dimension n -> Dim
forall (n :: Nat). KnownNat n => Dimension n -> Dim
fromDimension Dimension n
dim)
{-# INLINE withinSlices #-}
withinSlicesM :: (MonadThrow m, Slice r ix e) => Dim -> Array r ix e -> m (Array D Ix1 (Elt r ix e))
withinSlicesM :: Dim -> Array r ix e -> m (Array D Int (Elt r ix e))
withinSlicesM Dim
dim Array r ix e
arr = do
(Sz Int
k, Sz (Lower ix)
szl) <- Sz ix -> Dim -> m (Sz Int, Sz (Lower ix))
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz ix -> Dim -> m (Sz Int, Sz (Lower ix))
pullOutSzM (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) Dim
dim
Sz ix
cutSz <- Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
insertSzM Sz (Lower ix)
szl Dim
dim Sz Int
forall ix. Index ix => Sz ix
oneSz
Array D Int (Elt r ix e) -> m (Array D Int (Elt r ix e))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array D Int (Elt r ix e) -> m (Array D Int (Elt r ix e)))
-> Array D Int (Elt r ix e) -> m (Array D Int (Elt r ix e))
forall a b. (a -> b) -> a -> b
$ Comp -> Sz Int -> (Int -> Elt r ix e) -> Array D Int (Elt r ix e)
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
Seq Sz Int
k ((SomeException -> Elt r ix e)
-> (Elt r ix e -> Elt r ix e)
-> Either SomeException (Elt r ix e)
-> Elt r ix e
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Elt r ix e
forall e a. Exception e => e -> a
throwImpossible Elt r ix e -> Elt r ix e
forall a. a -> a
id (Either SomeException (Elt r ix e) -> Elt r ix e)
-> (Int -> Either SomeException (Elt r ix e)) -> Int -> Elt r ix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim
-> Sz ix
-> Array r ix e
-> Int
-> Either SomeException (Elt r ix e)
forall (m :: * -> *) r ix e.
(MonadThrow m, Slice r ix e) =>
Dim -> Sz ix -> Array r ix e -> Int -> m (Elt r ix e)
internalInnerSlice Dim
dim Sz ix
cutSz Array r ix e
arr)
{-# INLINE withinSlicesM #-}