{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
-- |
-- Module      : Data.Massiv.Array.Ops.Slice
-- Copyright   : (c) Alexey Kuleshevich 2018-2021
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Data.Massiv.Array.Ops.Slice
  (
  -- ** From the outside
    (!>)
  , (!?>)
  , (??>)
  -- ** From the inside
  , (<!)
  , (<!?)
  , (<??)
  -- ** From within
  , (<!>)
  , (<!?>)
  , (<??>)
  -- ** Many slices
  , outerSlices
  , innerSlices
  , withinSlices
  , withinSlicesM
  ) where

import Control.Monad (unless)
import Data.Massiv.Array.Delayed.Pull
import Data.Massiv.Core.Common


infixl 4 !>, !?>, ??>, <!, <!?, <??, <!>, <!?>, <??>


-- | /O(1)/ - Slices the array from the outside. For 2-dimensional array this will
-- be equivalent of taking a row. Throws an error when index is out of bounds.
--
-- ===__Examples__
--
-- You could say that slicing from outside is synonymous to slicing from the end or slicing at the
-- highermost dimension. For example with rank-3 arrays outer slice would be equivalent to getting a
-- page:
--
-- >>> import Data.Massiv.Array
-- >>> arr = makeArrayR U Seq (Sz (3 :> 2 :. 4)) fromIx3
-- >>> arr
-- Array U Seq (Sz (3 :> 2 :. 4))
--   [ [ [ (0,0,0), (0,0,1), (0,0,2), (0,0,3) ]
--     , [ (0,1,0), (0,1,1), (0,1,2), (0,1,3) ]
--     ]
--   , [ [ (1,0,0), (1,0,1), (1,0,2), (1,0,3) ]
--     , [ (1,1,0), (1,1,1), (1,1,2), (1,1,3) ]
--     ]
--   , [ [ (2,0,0), (2,0,1), (2,0,2), (2,0,3) ]
--     , [ (2,1,0), (2,1,1), (2,1,2), (2,1,3) ]
--     ]
--   ]
-- >>> arr !> 2
-- Array U Seq (Sz (2 :. 4))
--   [ [ (2,0,0), (2,0,1), (2,0,2), (2,0,3) ]
--   , [ (2,1,0), (2,1,1), (2,1,2), (2,1,3) ]
--   ]
--
-- There is nothing wrong with chaining, mixing and matching slicing operators:
--
-- >>> arr !> 2 !> 0 ! 3
-- (2,0,3)
-- >>> evaluateM (arr !> 2 <! 3) 0
-- (2,0,3)
-- >>> (arr !> 2 !> 0 ! 3) == (arr ! 2 :> 0 :. 3)
-- True
--
--
-- @since 0.1.0
(!>) ::
     forall r ix e. (HasCallStack, Index ix, Index (Lower ix), Source r e)
  => Array r ix e
  -> Int
  -> Array r (Lower ix) e
!> :: Array r ix e -> Int -> Array r (Lower ix) e
(!>) !Array r ix e
arr !Int
ix = Either SomeException (Array r (Lower ix) e) -> Array r (Lower ix) e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Array r ix e
arr Array r ix e -> Int -> Either SomeException (Array r (Lower ix) e)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
Array r ix e -> Int -> m (Array r (Lower ix) e)
!?> Int
ix)
{-# INLINE (!>) #-}


-- | /O(1)/ - Just like `!>` slices the array from the outside, but returns
-- `Nothing` when index is out of bounds.
--
-- @since 0.1.0
(!?>) ::
     forall r ix e m. (MonadThrow m, Index ix, Index (Lower ix), Source r e)
  => Array r ix e
  -> Int
  -> m (Array r (Lower ix) e)
!?> :: Array r ix e -> Int -> m (Array r (Lower ix) e)
(!?>) !Array r ix e
arr !Int
i = do
  let (Sz Int
k, Sz (Lower ix)
szL) = 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. Size r => Array r ix e -> Sz ix
size Array r ix e
arr)
  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
k 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
k Int
i
  Array r (Lower ix) e -> m (Array r (Lower ix) e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array r (Lower ix) e -> m (Array r (Lower ix) e))
-> Array r (Lower ix) e -> m (Array r (Lower ix) e)
forall a b. (a -> b) -> a -> b
$ Array r ix e -> Sz (Lower ix) -> Int -> Array r (Lower ix) e
forall r e ix.
(Source r e, 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)
szL Int
i
{-# INLINE (!?>) #-}


-- | /O(1)/ - Safe slicing continuation from the outside. Similarly to (`!>`) slices the array from
-- the outside, but takes `Maybe` array as input and returns `Nothing` when index is out of bounds.
--
-- ===__Examples__
--
-- >>> import Data.Massiv.Array
-- >>> arr = makeArrayR U Seq (Sz (3 :> 2 :. 4)) fromIx3
-- >>> arr !?> 2 ??> 0 ?? 3 :: Maybe Ix3T
-- Just (2,0,3)
-- >>> arr !?> 2 ??> 0 ?? -1 :: Maybe Ix3T
-- Nothing
-- >>> arr !?> 2 ??> -10 ?? 1
-- *** Exception: IndexOutOfBoundsException: -10 is not safe for (Sz1 2)
--
-- @since 0.1.0
(??>) ::
     forall r ix e m. (MonadThrow m, Index ix, Index (Lower ix), Source r e)
  => m (Array r ix e)
  -> Int
  -> m (Array r (Lower ix) e)
??> :: m (Array r ix e) -> Int -> m (Array r (Lower 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 (Array r (Lower ix) e))
-> m (Array r (Lower ix) e)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Array r ix e -> Int -> m (Array r (Lower ix) e)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
Array r ix e -> Int -> m (Array r (Lower ix) e)
!?> Int
ix)
{-# INLINE (??>) #-}


-- | /O(1)/ - Safe slice from the inside
--
-- @since 0.1.0
(<!?) ::
     forall r ix e m. (MonadThrow m, Index ix, Source r e)
  => Array r ix e
  -> Int
  -> m (Array D (Lower ix) e)
<!? :: Array r ix e -> Int -> m (Array D (Lower ix) e)
(<!?) !Array r ix e
arr !Int
i = do
  let (Sz (Lower ix)
szL, 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. Size r => Array r ix e -> Sz ix
size Array r ix e
arr)
  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
  Array D (Lower ix) e -> m (Array D (Lower ix) e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array D (Lower ix) e -> m (Array D (Lower ix) e))
-> Array D (Lower ix) e -> m (Array D (Lower ix) e)
forall a b. (a -> b) -> a -> b
$ Array r ix e -> Sz (Lower ix) -> Int -> Array D (Lower ix) e
forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> Sz (Lower ix) -> Int -> Array D (Lower ix) e
unsafeInnerSlice Array r ix e
arr Sz (Lower ix)
szL Int
i
{-# INLINE (<!?) #-}


-- | /O(1)/ - Similarly to (`!>`) slice an array from an opposite direction.
--
-- @since 0.1.0
(<!) ::
     forall r ix e. (HasCallStack, Index ix, Source r e)
  => Array r ix e
  -> Int
  -> Array D (Lower ix) e
<! :: Array r ix e -> Int -> Array D (Lower ix) e
(<!) !Array r ix e
arr !Int
ix = Either SomeException (Array D (Lower ix) e) -> Array D (Lower ix) e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Array r ix e
arr Array r ix e -> Int -> Either SomeException (Array D (Lower ix) e)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
Array r ix e -> Int -> m (Array D (Lower ix) e)
<!? Int
ix)
{-# INLINE (<!) #-}


-- | /O(1)/ - Safe slicing continuation from the inside
--
-- @since 0.1.0
(<??) ::
     forall r ix e m. (MonadThrow m, Index ix, Source r e)
  => m (Array r ix e)
  -> Int
  -> m (Array D (Lower ix) e)
<?? :: m (Array r ix e) -> Int -> m (Array D (Lower 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 (Array D (Lower ix) e))
-> m (Array D (Lower ix) e)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Array r ix e -> Int -> m (Array D (Lower ix) e)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Source r e) =>
Array r ix e -> Int -> m (Array D (Lower ix) e)
<!? Int
ix)
{-# INLINE (<??) #-}


-- | /O(1)/ - Same as (`<!>`), but fails gracefully with a `Nothing`, instead of an error
--
-- @since 0.1.0
(<!?>) ::
     forall r ix e m. (MonadThrow m, Index ix, Index (Lower ix), Source r e)
  => Array r ix e
  -> (Dim, Int)
  -> m (Array D (Lower ix) e)
<!?> :: Array r ix e -> (Dim, Int) -> m (Array D (Lower 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. Size r => 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 (Array D (Lower ix) e)
forall (m :: * -> *) ix r e.
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
Dim -> Sz ix -> Array r ix e -> Int -> m (Array D (Lower ix) e)
internalInnerSlice Dim
dim Sz ix
cutSz Array r ix e
arr Int
i
{-# INLINE (<!?>) #-}


internalInnerSlice ::
     (MonadThrow m, Index ix, Index (Lower ix), Source r e)
  => Dim
  -> Sz ix
  -> Array r ix e
  -> Ix1
  -> m (Array D (Lower ix) e)
internalInnerSlice :: Dim -> Sz ix -> Array r ix e -> Int -> m (Array D (Lower 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 (Array D (Lower ix) e)
forall r e ix (m :: * -> *).
(Source r e, Index ix, Index (Lower ix), MonadThrow m) =>
Array r ix e -> ix -> Sz ix -> Dim -> m (Array D (Lower ix) e)
unsafeSlice Array r ix e
arr ix
start Sz ix
cutSz Dim
dim
{-# INLINE internalInnerSlice #-}



-- prop> arr !> i == arr <!> (dimensions (size arr), i)
-- prop> arr <! i == arr <!> (1,i)
--
-- | /O(1)/ - Slices the array in any available dimension. Throws an error when
-- index is out of bounds or dimensions is invalid.
--
-- @since 0.1.0
(<!>) ::
     forall r ix e. (HasCallStack, Index ix, Index (Lower ix), Source r e)
  => Array r ix e
  -> (Dim, Int)
  -> Array D (Lower ix) e
<!> :: Array r ix e -> (Dim, Int) -> Array D (Lower ix) e
(<!>) !Array r ix e
arr !(Dim, Int)
dix = Either SomeException (Array D (Lower ix) e) -> Array D (Lower ix) e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Array r ix e
arr Array r ix e
-> (Dim, Int) -> Either SomeException (Array D (Lower ix) e)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
Array r ix e -> (Dim, Int) -> m (Array D (Lower ix) e)
<!?> (Dim, Int)
dix)
{-# INLINE (<!>) #-}


-- | /O(1)/ - Safe slicing continuation from within.
--
-- @since 0.1.0
(<??>) ::
     forall r ix e m. (MonadThrow m, Index ix, Index (Lower ix), Source r e)
  => m (Array r ix e)
  -> (Dim, Int)
  -> m (Array D (Lower ix) e)
<??> :: m (Array r ix e) -> (Dim, Int) -> m (Array D (Lower 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 (Array D (Lower ix) e))
-> m (Array D (Lower ix) e)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Array r ix e -> (Dim, Int) -> m (Array D (Lower ix) e)
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
Array r ix e -> (Dim, Int) -> m (Array D (Lower ix) e)
<!?> (Dim, Int)
ix)
{-# INLINE (<??>) #-}

-- | Create a delayed array of outer slices.
--
-- ====__Examples__
--
-- >>> import Data.Massiv.Array as A
-- >>> A.mapM_ print $ outerSlices (0 ..: (3 :. 2))
-- Array D Seq (Sz1 2)
--   [ 0 :. 0, 0 :. 1 ]
-- Array D Seq (Sz1 2)
--   [ 1 :. 0, 1 :. 1 ]
-- Array D Seq (Sz1 2)
--   [ 2 :. 0, 2 :. 1 ]
--
-- @since 0.5.4
outerSlices ::
     forall r ix e. (Index ix, Index (Lower ix), Source r e)
  => Array r ix e
  -> Array D Ix1 (Array r (Lower ix) e)
outerSlices :: Array r ix e -> Array D Int (Array r (Lower ix) e)
outerSlices Array r ix e
arr = Comp
-> Sz Int
-> (Int -> Array r (Lower ix) e)
-> Array D Int (Array r (Lower ix) e)
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) Sz Int
k (Array r ix e -> Sz (Lower ix) -> Int -> Array r (Lower ix) e
forall r e ix.
(Source r e, Index ix, Index (Lower ix)) =>
Array r ix e -> Sz (Lower ix) -> Int -> Array r (Lower ix) e
unsafeOuterSlice (Comp -> Array r ix e -> Array r ix e
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
setComp Comp
Seq Array r ix e
arr) Sz (Lower ix)
szL)
  where
    (Sz Int
k, Sz (Lower ix)
szL) = 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. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
{-# INLINE outerSlices #-}


-- | Create a delayed array of inner slices.
--
-- ====__Examples__
--
-- >>> import Data.Massiv.Array as A
-- >>> A.mapM_ print $ innerSlices (0 ..: (3 :. 2))
-- Array D Seq (Sz1 3)
--   [ 0 :. 0, 1 :. 0, 2 :. 0 ]
-- Array D Seq (Sz1 3)
--   [ 0 :. 1, 1 :. 1, 2 :. 1 ]
--
-- @since 0.5.4
innerSlices ::
     forall r ix e. (Index ix, Source r e)
  => Array r ix e
  -> Array D Ix1 (Array D (Lower ix) e)
innerSlices :: Array r ix e -> Array D Int (Array D (Lower ix) e)
innerSlices Array r ix e
arr = Comp
-> Sz Int
-> (Int -> Array D (Lower ix) e)
-> Array D Int (Array D (Lower ix) e)
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) Sz Int
k (Array r ix e -> Sz (Lower ix) -> Int -> Array D (Lower ix) e
forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> Sz (Lower ix) -> Int -> Array D (Lower ix) e
unsafeInnerSlice (Comp -> Array r ix e -> Array r ix e
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
setComp Comp
Seq Array r ix e
arr) Sz (Lower ix)
szL)
  where
    (Sz (Lower ix)
szL, 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. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
{-# INLINE innerSlices #-}

-- | Create a delayed array of slices from within. Checks dimension at compile time.
--
-- ====__Examples__
--
-- >>> import Data.Massiv.Array as A
-- >>> arr = fromIx3 <$> (0 ..: (4 :> 3 :. 2))
-- >>> print arr
-- Array D Seq (Sz (4 :> 3 :. 2))
--   [ [ [ (0,0,0), (0,0,1) ]
--     , [ (0,1,0), (0,1,1) ]
--     , [ (0,2,0), (0,2,1) ]
--     ]
--   , [ [ (1,0,0), (1,0,1) ]
--     , [ (1,1,0), (1,1,1) ]
--     , [ (1,2,0), (1,2,1) ]
--     ]
--   , [ [ (2,0,0), (2,0,1) ]
--     , [ (2,1,0), (2,1,1) ]
--     , [ (2,2,0), (2,2,1) ]
--     ]
--   , [ [ (3,0,0), (3,0,1) ]
--     , [ (3,1,0), (3,1,1) ]
--     , [ (3,2,0), (3,2,1) ]
--     ]
--   ]
-- >>> A.mapM_ print $ withinSlices Dim2 arr
-- Array D Seq (Sz (4 :. 2))
--   [ [ (0,0,0), (0,0,1) ]
--   , [ (1,0,0), (1,0,1) ]
--   , [ (2,0,0), (2,0,1) ]
--   , [ (3,0,0), (3,0,1) ]
--   ]
-- Array D Seq (Sz (4 :. 2))
--   [ [ (0,1,0), (0,1,1) ]
--   , [ (1,1,0), (1,1,1) ]
--   , [ (2,1,0), (2,1,1) ]
--   , [ (3,1,0), (3,1,1) ]
--   ]
-- Array D Seq (Sz (4 :. 2))
--   [ [ (0,2,0), (0,2,1) ]
--   , [ (1,2,0), (1,2,1) ]
--   , [ (2,2,0), (2,2,1) ]
--   , [ (3,2,0), (3,2,1) ]
--   ]
--
-- @since 0.5.4
withinSlices ::
     forall n r ix e. (IsIndexDimension ix n, Index (Lower ix), Source r e)
  => Dimension n
  -> Array r ix e
  -> Array D Ix1 (Array D (Lower ix) e)
withinSlices :: Dimension n -> Array r ix e -> Array D Int (Array D (Lower ix) e)
withinSlices Dimension n
dim = (SomeException -> Array D Int (Array D (Lower ix) e))
-> (Array D Int (Array D (Lower ix) e)
    -> Array D Int (Array D (Lower ix) e))
-> Either SomeException (Array D Int (Array D (Lower ix) e))
-> Array D Int (Array D (Lower ix) e)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array D Int (Array D (Lower ix) e)
forall e a. (HasCallStack, Exception e) => e -> a
throwImpossible Array D Int (Array D (Lower ix) e)
-> Array D Int (Array D (Lower ix) e)
forall a. a -> a
id (Either SomeException (Array D Int (Array D (Lower ix) e))
 -> Array D Int (Array D (Lower ix) e))
-> (Array r ix e
    -> Either SomeException (Array D Int (Array D (Lower ix) e)))
-> Array r ix e
-> Array D Int (Array D (Lower ix) e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim
-> Array r ix e
-> Either SomeException (Array D Int (Array D (Lower ix) e))
forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
Dim -> Array r ix e -> m (Array D Int (Array D (Lower ix) e))
withinSlicesM (Dimension n -> Dim
forall (n :: Nat). KnownNat n => Dimension n -> Dim
fromDimension Dimension n
dim)
{-# INLINE withinSlices #-}


-- | Create a delayed array of slices from within. Same as `withinSlices`, but throws an
-- error on invalid dimension.
--
-- /__Throws Exceptions__/: `IndexDimensionException`
--
-- @since 0.5.4
withinSlicesM ::
     forall r ix e m. (MonadThrow m, Index ix, Index (Lower ix), Source r e)
  => Dim
  -> Array r ix e
  -> m (Array D Ix1 (Array D (Lower ix) e))
withinSlicesM :: Dim -> Array r ix e -> m (Array D Int (Array D (Lower 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. Size r => 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 (Array D (Lower ix) e)
-> m (Array D Int (Array D (Lower ix) e))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array D Int (Array D (Lower ix) e)
 -> m (Array D Int (Array D (Lower ix) e)))
-> Array D Int (Array D (Lower ix) e)
-> m (Array D Int (Array D (Lower ix) e))
forall a b. (a -> b) -> a -> b
$ Comp
-> Sz Int
-> (Int -> Array D (Lower ix) e)
-> Array D Int (Array D (Lower ix) e)
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
Seq Sz Int
k ((SomeException -> Array D (Lower ix) e)
-> (Array D (Lower ix) e -> Array D (Lower ix) e)
-> Either SomeException (Array D (Lower ix) e)
-> Array D (Lower ix) e
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Array D (Lower ix) e
forall e a. (HasCallStack, Exception e) => e -> a
throwImpossible Array D (Lower ix) e -> Array D (Lower ix) e
forall a. a -> a
id (Either SomeException (Array D (Lower ix) e)
 -> Array D (Lower ix) e)
-> (Int -> Either SomeException (Array D (Lower ix) e))
-> Int
-> Array D (Lower ix) e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim
-> Sz ix
-> Array r ix e
-> Int
-> Either SomeException (Array D (Lower ix) e)
forall (m :: * -> *) ix r e.
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
Dim -> Sz ix -> Array r ix e -> Int -> m (Array D (Lower ix) e)
internalInnerSlice Dim
dim Sz ix
cutSz Array r ix e
arr)
{-# INLINE withinSlicesM #-}