{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ExplicitNamespaces #-}
-- |
-- Module      : Data.Massiv.Core.Index
-- Copyright   : (c) Alexey Kuleshevich 2018-2021
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <alexey@kuleshevi.ch>
-- Stability   : experimental
-- Portability : non-portable
--
module Data.Massiv.Core.Index
  ( Ix0(..)
  , type Ix1
  , pattern Ix1
  , type Ix2(Ix2, (:.))
  , IxN((:>), Ix3, Ix4, Ix5)
  , HighIxN
  , type Ix3
  , type Ix4
  , type Ix5
  , Ix
  -- ** Size
  , type Sz1
  , type Sz2
  , type Sz3
  , type Sz4
  , type Sz5
  , Sz(Sz, Sz1, Sz2, Sz3, Sz4, Sz5)
  , unSz
  , zeroSz
  , oneSz
  , liftSz
  , liftSz2
  , consSz
  , unconsSz
  , snocSz
  , unsnocSz
  , setSzM
  , insertSzM
  , pullOutSzM
  , toLinearSz
  , mkSzM
  -- ** Dimension
  , Dim(..)
  , Dimension(Dim1, Dim2, Dim3, Dim4, Dim5, DimN)
  , IsIndexDimension
  , IsDimValid
  , ReportInvalidDim
  -- ** Stride
  , Stride(Stride)
  , unStride
  , toLinearIndexStride
  , strideStart
  , strideSize
  , oneStride
  -- ** Border
  , Border(..)
  , handleBorderIndex
  -- ** Index functions
  , Lower
  , Index(..)
  , zeroIndex
  , oneIndex
  , isZeroSz
  , isNotZeroSz
  , headDim
  , tailDim
  , lastDim
  , initDim
  , getDim'
  , setDim'
  , modifyDim'
  , dropDimM
  , dropDim'
  , pullOutDim'
  , insertDim'
  , fromDimension
  , getDimension
  , setDimension
  , modifyDimension
  , dropDimension
  , pullOutDimension
  , insertDimension
  -- * Iterators
  , iter
  , iterLinearM
  , iterLinearM_
  , module Data.Massiv.Core.Iterator
  , module Data.Massiv.Core.Index.Tuple
  -- * Exceptions
  , IndexException(..)
  , SizeException(..)
  , ShapeException(..)
  , guardNumberOfElements
  , indexWith
  ) where

import Control.DeepSeq
import Control.Monad.Catch (MonadThrow(..))
import Data.Coerce
import Data.Functor.Identity (runIdentity)
import Data.Massiv.Core.Exception
import Data.Massiv.Core.Index.Internal
import Data.Massiv.Core.Index.Ix
import Data.Massiv.Core.Index.Stride
import Data.Massiv.Core.Index.Tuple
import Data.Massiv.Core.Iterator
import GHC.TypeLits


-- | 1-dimensional type synonym for size.
--
-- @since 0.3.0
type Sz1 = Sz Ix1

-- | 2-dimensional size type synonym.
--
-- @since 0.3.0
type Sz2 = Sz Ix2

-- | 3-dimensional size type synonym.
--
-- @since 0.3.0
type Sz3 = Sz Ix3

-- | 4-dimensional size type synonym.
--
-- @since 0.3.0
type Sz4 = Sz Ix4

-- | 5-dimensional size type synonym.
--
-- @since 0.3.0
type Sz5 = Sz Ix5


-- | Approach to be used near the borders during various transformations.
-- Whenever a function needs information not only about an element of interest, but
-- also about it's neighbors, it will go out of bounds near the array edges,
-- hence is this set of approaches that specify how to handle such situation.
data Border e =
  Fill e    -- ^ Fill in a constant element.
              --
              -- @
              --            outside |  Array  | outside
              -- ('Fill' 0) : 0 0 0 0 | 1 2 3 4 | 0 0 0 0
              -- @
              --
  | Wrap      -- ^ Wrap around from the opposite border of the array.
              --
              -- @
              --            outside |  Array  | outside
              -- 'Wrap' :     1 2 3 4 | 1 2 3 4 | 1 2 3 4
              -- @
              --
  | Edge      -- ^ Replicate the element at the edge.
              --
              -- @
              --            outside |  Array  | outside
              -- 'Edge' :     1 1 1 1 | 1 2 3 4 | 4 4 4 4
              -- @
              --
  | Reflect   -- ^ Mirror like reflection.
              --
              -- @
              --            outside |  Array  | outside
              -- 'Reflect' :  4 3 2 1 | 1 2 3 4 | 4 3 2 1
              -- @
              --
  | Continue  -- ^ Also mirror like reflection, but without repeating the edge element.
              --
              -- @
              --            outside |  Array  | outside
              -- 'Continue' : 1 4 3 2 | 1 2 3 4 | 3 2 1 4
              -- @
              --
  deriving (Border e -> Border e -> Bool
(Border e -> Border e -> Bool)
-> (Border e -> Border e -> Bool) -> Eq (Border e)
forall e. Eq e => Border e -> Border e -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Border e -> Border e -> Bool
$c/= :: forall e. Eq e => Border e -> Border e -> Bool
== :: Border e -> Border e -> Bool
$c== :: forall e. Eq e => Border e -> Border e -> Bool
Eq, Int -> Border e -> ShowS
[Border e] -> ShowS
Border e -> String
(Int -> Border e -> ShowS)
-> (Border e -> String) -> ([Border e] -> ShowS) -> Show (Border e)
forall e. Show e => Int -> Border e -> ShowS
forall e. Show e => [Border e] -> ShowS
forall e. Show e => Border e -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Border e] -> ShowS
$cshowList :: forall e. Show e => [Border e] -> ShowS
show :: Border e -> String
$cshow :: forall e. Show e => Border e -> String
showsPrec :: Int -> Border e -> ShowS
$cshowsPrec :: forall e. Show e => Int -> Border e -> ShowS
Show)

instance NFData e => NFData (Border e) where
  rnf :: Border e -> ()
rnf Border e
b = case Border e
b of
            Fill e
e   -> e -> ()
forall a. NFData a => a -> ()
rnf e
e
            Border e
Wrap     -> ()
            Border e
Edge     -> ()
            Border e
Reflect  -> ()
            Border e
Continue -> ()


-- | Apply a border resolution technique to an index
--
-- ==== __Examples__
--
-- >>> handleBorderIndex (Fill 100) (Sz (2 :. 3)) id (2 :. 3)
-- 100 :. 100
-- >>> handleBorderIndex Wrap (Sz (2 :. 3)) id (2 :. 3)
-- 0 :. 0
-- >>> handleBorderIndex Edge (Sz (2 :. 3)) id (2 :. 3)
-- 1 :. 2
--
-- @since 0.1.0
handleBorderIndex ::
     Index ix
  => Border e -- ^ Broder resolution technique
  -> Sz ix -- ^ Size
  -> (ix -> e) -- ^ Index function that produces an element
  -> ix -- ^ Index
  -> e
handleBorderIndex :: Border e -> Sz ix -> (ix -> e) -> ix -> e
handleBorderIndex Border e
border !Sz ix
sz ix -> e
getVal !ix
ix =
  case Border e
border of
    Fill e
val -> if Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz ix
ix then ix -> e
getVal ix
ix else e
val
    Border e
Wrap     -> ix -> e
getVal (Sz ix
-> ix -> (Sz Int -> Int -> Int) -> (Sz Int -> Int -> Int) -> ix
forall ix.
Index ix =>
Sz ix
-> ix -> (Sz Int -> Int -> Int) -> (Sz Int -> Int -> Int) -> ix
repairIndex Sz ix
sz ix
ix Sz Int -> Int -> Int
forall a. Integral a => Sz a -> a -> a
wrap Sz Int -> Int -> Int
forall a. Integral a => Sz a -> a -> a
wrap)
    Border e
Edge     -> ix -> e
getVal (Sz ix
-> ix -> (Sz Int -> Int -> Int) -> (Sz Int -> Int -> Int) -> ix
forall ix.
Index ix =>
Sz ix
-> ix -> (Sz Int -> Int -> Int) -> (Sz Int -> Int -> Int) -> ix
repairIndex Sz ix
sz ix
ix ((Int -> Int) -> Sz Int -> Int -> Int
forall a b. a -> b -> a
const (Int -> Int -> Int
forall a b. a -> b -> a
const Int
0)) (\ (SafeSz Int
k) Int
_ -> Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
    Border e
Reflect  -> ix -> e
getVal (Sz ix
-> ix -> (Sz Int -> Int -> Int) -> (Sz Int -> Int -> Int) -> ix
forall ix.
Index ix =>
Sz ix
-> ix -> (Sz Int -> Int -> Int) -> (Sz Int -> Int -> Int) -> ix
repairIndex Sz ix
sz ix
ix (\ (SafeSz Int
k) !Int
i -> (Int -> Int
forall a. Num a => a -> a
abs Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
k)
                        (\ (SafeSz Int
k) !Int
i -> (-Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
k))
    Border e
Continue -> ix -> e
getVal (Sz ix
-> ix -> (Sz Int -> Int -> Int) -> (Sz Int -> Int -> Int) -> ix
forall ix.
Index ix =>
Sz ix
-> ix -> (Sz Int -> Int -> Int) -> (Sz Int -> Int -> Int) -> ix
repairIndex Sz ix
sz ix
ix (\ (SafeSz Int
k) !Int
i -> Int -> Int
forall a. Num a => a -> a
abs Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
k)
                        (\ (SafeSz Int
k) !Int
i -> (-Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
k))

  where wrap :: Sz a -> a -> a
wrap (SafeSz a
k) a
i = a
i a -> a -> a
forall a. Integral a => a -> a -> a
`mod` a
k
        {-# INLINE [1] wrap #-}
{-# INLINE [1] handleBorderIndex #-}

-- | Index with all zeros
--
-- ==== __Examples__
--
-- >>> zeroIndex :: Ix4
-- 0 :> 0 :> 0 :. 0
--
-- @since 0.1.0
zeroIndex :: Index ix => ix
zeroIndex :: ix
zeroIndex = Int -> ix
forall ix. Index ix => Int -> ix
pureIndex Int
0
{-# INLINE [1] zeroIndex #-}

-- | Index with all ones
--
-- @since 0.3.0
oneIndex :: Index ix => ix
oneIndex :: ix
oneIndex = Int -> ix
forall ix. Index ix => Int -> ix
pureIndex Int
1
{-# INLINE [1] oneIndex #-}

-- | Checks whether size can hold at least one element.
--
-- ==== __Examples__
--
-- >>> isNotZeroSz (Sz3 1 0 2)
-- False
--
-- @since 1.0.0
isNotZeroSz :: Index ix => Sz ix -> Bool
isNotZeroSz :: Sz ix -> Bool
isNotZeroSz !Sz ix
sz = Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz ix
forall ix. Index ix => ix
zeroIndex
{-# INLINE [1] isNotZeroSz #-}
-- TODO: benchmark against (also adjust `isEmpty` with fastest):
-- - foldlIndex (*) 1 (unSz sz) /= 0
-- - foldlIndex (\a x -> a && x /= 0) True (unSz sz)
-- - totalElem sz == 0

-- | Checks whether size can hold at least one element.
--
-- ==== __Examples__
--
-- >>> isZeroSz (Sz3 1 0 2)
-- True
--
-- @since 1.0.0
isZeroSz :: Index ix => Sz ix -> Bool
isZeroSz :: Sz ix -> Bool
isZeroSz = Bool -> Bool
not (Bool -> Bool) -> (Sz ix -> Bool) -> Sz ix -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Bool
forall ix. Index ix => Sz ix -> Bool
isNotZeroSz
{-# INLINE [1] isZeroSz #-}


-- | Convert a size to a linear size.
--
-- @since 0.5.8
toLinearSz :: Index ix => Sz ix -> Sz1
toLinearSz :: Sz ix -> Sz Int
toLinearSz = Int -> Sz Int
coerce (Int -> Sz Int) -> (Sz ix -> Int) -> Sz ix -> Sz Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem
{-# INLINE [1] toLinearSz #-}

-- | Get the outmost dimension of the index.
--
-- ==== __Examples__
--
-- >>> headDim (2 :> 3 :> 4 :. 5)
-- 2
--
-- @since 0.1.0
headDim :: Index ix => ix -> Int
headDim :: ix -> Int
headDim = (Int, Lower ix) -> Int
forall a b. (a, b) -> a
fst ((Int, Lower ix) -> Int) -> (ix -> (Int, Lower ix)) -> ix -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim
{-# INLINE [1] headDim #-}

-- | Drop the outmost dimension from the index
--
-- ==== __Examples__
--
-- >>> tailDim (2 :> 3 :> 4 :. 5)
-- 3 :> 4 :. 5
--
-- @since 0.1.0
tailDim :: Index ix => ix -> Lower ix
tailDim :: ix -> Lower ix
tailDim = (Int, Lower ix) -> Lower ix
forall a b. (a, b) -> b
snd ((Int, Lower ix) -> Lower ix)
-> (ix -> (Int, Lower ix)) -> ix -> Lower ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim
{-# INLINE [1] tailDim #-}

-- | Get the innermost dimension from the index
--
-- ==== __Examples__
--
-- >>> lastDim (2 :> 3 :> 4 :. 5)
-- 5
--
-- @since 0.1.0
lastDim :: Index ix => ix -> Int
lastDim :: ix -> Int
lastDim = (Lower ix, Int) -> Int
forall a b. (a, b) -> b
snd ((Lower ix, Int) -> Int) -> (ix -> (Lower ix, Int)) -> ix -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> (Lower ix, Int)
forall ix. Index ix => ix -> (Lower ix, Int)
unsnocDim
{-# INLINE [1] lastDim #-}

-- | Drop the innermost dimension from the index
--
-- ==== __Examples__
--
-- >>> initDim (2 :> 3 :> 4 :. 5)
-- 2 :> 3 :. 4
--
-- @since 0.1.0
initDim :: Index ix => ix -> Lower ix
initDim :: ix -> Lower ix
initDim = (Lower ix, Int) -> Lower ix
forall a b. (a, b) -> a
fst ((Lower ix, Int) -> Lower ix)
-> (ix -> (Lower ix, Int)) -> ix -> Lower ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> (Lower ix, Int)
forall ix. Index ix => ix -> (Lower ix, Int)
unsnocDim
{-# INLINE [1] initDim #-}

-- | Change the value of a specific dimension within the index. See `setDimM` for a safer
-- version and `setDimension` for a type safe version.
--
-- ==== __Examples__
--
-- >>> setDim' (2 :> 3 :> 4 :. 5) 3 10
-- 2 :> 10 :> 4 :. 5
--
-- @since 0.2.4
setDim' :: (HasCallStack, Index ix) => ix -> Dim -> Int -> ix
setDim' :: ix -> Dim -> Int -> ix
setDim' ix
ix Dim
dim = Either SomeException ix -> ix
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException ix -> ix)
-> (Int -> Either SomeException ix) -> Int -> ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> Dim -> Int -> Either SomeException ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
ix Dim
dim
{-# INLINE [1] setDim' #-}

-- | Change the value from a specific dimension within the index. See
-- `getDimM` for a safer version and `getDimension` for a type safe version.
--
-- ==== __Examples__
--
-- >>> getDim' (2 :> 3 :> 4 :. 5) 3
-- 3
--
-- @since 0.2.4
getDim' :: (HasCallStack, Index ix) => ix -> Dim -> Int
getDim' :: ix -> Dim -> Int
getDim' ix
ix = Either SomeException Int -> Int
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException Int -> Int)
-> (Dim -> Either SomeException Int) -> Dim -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> Dim -> Either SomeException Int
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m Int
getDimM ix
ix
{-# INLINE [1] getDim' #-}

-- | Update the value of a specific dimension within the index. See
-- `modifyDimM` for a safer version and `modifyDimension` for a type safe version.
--
-- ==== __Examples__
--
-- >>> modifyDim' (2 :> 3 :> 4 :. 5) 2 (+ 10)
-- (4,2 :> 3 :> 14 :. 5)
--
-- @since 0.4.1
modifyDim' :: (HasCallStack, Index ix) => ix -> Dim -> (Int -> Int) -> (Int, ix)
modifyDim' :: ix -> Dim -> (Int -> Int) -> (Int, ix)
modifyDim' ix
ix Dim
dim = Either SomeException (Int, ix) -> (Int, ix)
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException (Int, ix) -> (Int, ix))
-> ((Int -> Int) -> Either SomeException (Int, ix))
-> (Int -> Int)
-> (Int, ix)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> Dim -> (Int -> Int) -> Either SomeException (Int, ix)
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> (Int -> Int) -> m (Int, ix)
modifyDimM ix
ix Dim
dim
{-# INLINE [1] modifyDim' #-}

-- | Remove a dimension from the index.
--
-- ==== __Examples__
--
-- >>> dropDimM (2 :> 3 :> 4 :. 5) 3 :: Maybe Ix3
-- Just (2 :> 4 :. 5)
-- >>> dropDimM (2 :> 3 :> 4 :. 5) 6 :: Maybe Ix3
-- Nothing
--
-- @since 0.3.0
dropDimM :: (MonadThrow m, Index ix) => ix -> Dim -> m (Lower ix)
dropDimM :: ix -> Dim -> m (Lower ix)
dropDimM ix
ix = ((Int, Lower ix) -> Lower ix) -> m (Int, Lower ix) -> m (Lower ix)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int, Lower ix) -> Lower ix
forall a b. (a, b) -> b
snd (m (Int, Lower ix) -> m (Lower ix))
-> (Dim -> m (Int, Lower ix)) -> Dim -> m (Lower ix)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> Dim -> m (Int, Lower ix)
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m (Int, Lower ix)
pullOutDimM ix
ix
{-# INLINE [1] dropDimM #-}

-- | Remove a dimension from the index.
--
-- ==== __Examples__
--
-- >>> dropDim' (2 :> 3 :> 4 :. 5) 3
-- 2 :> 4 :. 5
--
-- @since 0.2.4
dropDim' :: (HasCallStack, Index ix) => ix -> Dim -> Lower ix
dropDim' :: ix -> Dim -> Lower ix
dropDim' ix
ix = Either SomeException (Lower ix) -> Lower ix
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException (Lower ix) -> Lower ix)
-> (Dim -> Either SomeException (Lower ix)) -> Dim -> Lower ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> Dim -> Either SomeException (Lower ix)
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
ix -> Dim -> m (Lower ix)
dropDimM ix
ix
{-# INLINE [1] dropDim' #-}

-- | Lower the dimension of the index by pulling the specified dimension. See
-- `pullOutDimM` for a safer version and `pullOutDimension` for a type safe version.
--
-- ==== __Examples__
--
-- >>> pullOutDim' (2 :> 3 :> 4 :. 5) 3
-- (3,2 :> 4 :. 5)
--
-- @since 0.2.4
pullOutDim' :: (HasCallStack, Index ix) => ix -> Dim -> (Int, Lower ix)
pullOutDim' :: ix -> Dim -> (Int, Lower ix)
pullOutDim' ix
ix = Either SomeException (Int, Lower ix) -> (Int, Lower ix)
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException (Int, Lower ix) -> (Int, Lower ix))
-> (Dim -> Either SomeException (Int, Lower ix))
-> Dim
-> (Int, Lower ix)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> Dim -> Either SomeException (Int, Lower ix)
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m (Int, Lower ix)
pullOutDimM ix
ix
{-# INLINE [1] pullOutDim' #-}

-- | Raise the dimension of the index by inserting one in the specified dimension. See
-- `insertDimM` for a safer version and `insertDimension` for a type safe version.
--
-- ==== __Examples__
--
-- >>> insertDim' (2 :> 3 :> 4 :. 5) 3 10 :: Ix5
-- 2 :> 3 :> 10 :> 4 :. 5
--
-- @since 0.2.4
insertDim' :: (HasCallStack, Index ix) => Lower ix -> Dim -> Int -> ix
insertDim' :: Lower ix -> Dim -> Int -> ix
insertDim' Lower ix
ix Dim
dim = Either SomeException ix -> ix
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException ix -> ix)
-> (Int -> Either SomeException ix) -> Int -> ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lower ix -> Dim -> Int -> Either SomeException ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
Lower ix -> Dim -> Int -> m ix
insertDimM Lower ix
ix Dim
dim
{-# INLINE [1] insertDim' #-}

-- | Get the value level `Dim` from the type level equivalent.
--
-- ==== __Examples__
--
-- >>> fromDimension Dim4
-- (Dim 4)
-- >>> :set -XDataKinds
-- >>> fromDimension (DimN :: Dimension 10)
-- (Dim 10)
--
-- @since 0.2.4
fromDimension :: KnownNat n => Dimension n -> Dim
fromDimension :: Dimension n -> Dim
fromDimension = Integer -> Dim
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Dim) -> (Dimension n -> Integer) -> Dimension n -> Dim
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dimension n -> Integer
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Integer
natVal
{-# INLINE [1] fromDimension #-}

-- | Type safe way to set value of index at a particular dimension.
--
-- ==== __Examples__
--
-- >>> setDimension (2 :> 3 :> 4 :. 5) Dim4 10
-- 10 :> 3 :> 4 :. 5
--
-- @since 0.2.4
setDimension :: IsIndexDimension ix n => ix -> Dimension n -> Int -> ix
setDimension :: ix -> Dimension n -> Int -> ix
setDimension ix
ix = ix -> Dim -> Int -> ix
forall ix. (HasCallStack, Index ix) => ix -> Dim -> Int -> ix
setDim' ix
ix (Dim -> Int -> ix)
-> (Dimension n -> Dim) -> Dimension n -> Int -> ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dimension n -> Dim
forall (n :: Nat). KnownNat n => Dimension n -> Dim
fromDimension
{-# INLINE [1] setDimension #-}

-- | Type safe way to set value of index at a particular dimension.
--
-- ==== __Examples__
--
-- >>> modifyDimension (2 :> 3 :> 4 :. 5) Dim3 (+ 2)
-- (3,2 :> 5 :> 4 :. 5)
--
-- @since 0.4.1
modifyDimension :: IsIndexDimension ix n => ix -> Dimension n -> (Int -> Int) -> (Int, ix)
modifyDimension :: ix -> Dimension n -> (Int -> Int) -> (Int, ix)
modifyDimension ix
ix = ix -> Dim -> (Int -> Int) -> (Int, ix)
forall ix.
(HasCallStack, Index ix) =>
ix -> Dim -> (Int -> Int) -> (Int, ix)
modifyDim' ix
ix (Dim -> (Int -> Int) -> (Int, ix))
-> (Dimension n -> Dim) -> Dimension n -> (Int -> Int) -> (Int, ix)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dimension n -> Dim
forall (n :: Nat). KnownNat n => Dimension n -> Dim
fromDimension
{-# INLINE [1] modifyDimension #-}

-- | Type safe way to extract value of index at a particular dimension.
--
-- ==== __Examples__
--
-- >>> getDimension (2 :> 3 :> 4 :. 5) Dim2
-- 4
--
-- @since 0.2.4
getDimension :: IsIndexDimension ix n => ix -> Dimension n -> Int
getDimension :: ix -> Dimension n -> Int
getDimension ix
ix = ix -> Dim -> Int
forall ix. (HasCallStack, Index ix) => ix -> Dim -> Int
getDim' ix
ix (Dim -> Int) -> (Dimension n -> Dim) -> Dimension n -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dimension n -> Dim
forall (n :: Nat). KnownNat n => Dimension n -> Dim
fromDimension
{-# INLINE [1] getDimension #-}


-- | Type safe way of dropping a particular dimension, thus lowering index
-- dimensionality.
--
-- ==== __Examples__
--
-- >>> dropDimension (2 :> 3 :> 4 :. 5) Dim2
-- 2 :> 3 :. 5
--
-- @since 0.2.4
dropDimension :: IsIndexDimension ix n => ix -> Dimension n -> Lower ix
dropDimension :: ix -> Dimension n -> Lower ix
dropDimension ix
ix = ix -> Dim -> Lower ix
forall ix. (HasCallStack, Index ix) => ix -> Dim -> Lower ix
dropDim' ix
ix (Dim -> Lower ix)
-> (Dimension n -> Dim) -> Dimension n -> Lower ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dimension n -> Dim
forall (n :: Nat). KnownNat n => Dimension n -> Dim
fromDimension
{-# INLINE [1] dropDimension #-}

-- | Type safe way of pulling out a particular dimension, thus lowering index
-- dimensionality and returning the value at specified dimension.
--
-- ==== __Examples__
--
-- >>> pullOutDimension (2 :> 3 :> 4 :. 5) Dim2
-- (4,2 :> 3 :. 5)
--
-- @since 0.2.4
pullOutDimension :: IsIndexDimension ix n => ix -> Dimension n -> (Int, Lower ix)
pullOutDimension :: ix -> Dimension n -> (Int, Lower ix)
pullOutDimension ix
ix = ix -> Dim -> (Int, Lower ix)
forall ix. (HasCallStack, Index ix) => ix -> Dim -> (Int, Lower ix)
pullOutDim' ix
ix (Dim -> (Int, Lower ix))
-> (Dimension n -> Dim) -> Dimension n -> (Int, Lower ix)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dimension n -> Dim
forall (n :: Nat). KnownNat n => Dimension n -> Dim
fromDimension
{-# INLINE [1] pullOutDimension #-}

-- | Type safe way of inserting a particular dimension, thus raising index dimensionality.
--
-- ==== __Examples__
--
-- >>> insertDimension (2 :> 3 :> 4 :. 5) Dim5 10 :: Ix5
-- 10 :> 2 :> 3 :> 4 :. 5
-- >>> insertDimension (2 :> 3 :> 4 :. 5) Dim4 10 :: Ix5
-- 2 :> 10 :> 3 :> 4 :. 5
-- >>> insertDimension (2 :> 3 :> 4 :. 5) Dim3 10 :: Ix5
-- 2 :> 3 :> 10 :> 4 :. 5
-- >>> insertDimension (2 :> 3 :> 4 :. 5) Dim2 10 :: Ix5
-- 2 :> 3 :> 4 :> 10 :. 5
-- >>> insertDimension (2 :> 3 :> 4 :. 5) Dim1 10 :: Ix5
-- 2 :> 3 :> 4 :> 5 :. 10
--
-- @since 0.2.5
insertDimension :: IsIndexDimension ix n => Lower ix -> Dimension n -> Int -> ix
insertDimension :: Lower ix -> Dimension n -> Int -> ix
insertDimension Lower ix
ix = Lower ix -> Dim -> Int -> ix
forall ix. (HasCallStack, Index ix) => Lower ix -> Dim -> Int -> ix
insertDim' Lower ix
ix (Dim -> Int -> ix)
-> (Dimension n -> Dim) -> Dimension n -> Int -> ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dimension n -> Dim
forall (n :: Nat). KnownNat n => Dimension n -> Dim
fromDimension
{-# INLINE [1] insertDimension #-}

-- | Row-major iterator for the index. Same as `iterM`, but pure.
--
-- ==== __Examples__
--
-- >>> iter (Ix1 0) 1000 1 (<) 0 (+)
-- 499500
-- >>> iter (0 :. 0) (2 :. 3) oneIndex (<) 100 $ \ (i :. j) acc -> (acc + i) * (j + 1)
-- 3615
--
-- @since 0.1.0
iter :: Index ix
  => ix -- ^ Start index
  -> ix -- ^ End index
  -> ix -- ^ Increment
  -> (Int -> Int -> Bool) -- ^ Continuation condition
  -> a -- ^ Accumulator
  -> (ix -> a -> a) -- ^ Iterating function
  -> a
iter :: ix -> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> a) -> a
iter ix
sIx ix
eIx ix
incIx Int -> Int -> Bool
cond a
acc ix -> a -> a
f =
  Identity a -> a
forall a. Identity a -> a
runIdentity (Identity a -> a) -> Identity a -> a
forall a b. (a -> b) -> a -> b
$ ix
-> ix
-> ix
-> (Int -> Int -> Bool)
-> a
-> (ix -> a -> Identity a)
-> Identity a
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix
-> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a
iterM ix
sIx ix
eIx ix
incIx Int -> Int -> Bool
cond a
acc (\ix
ix -> a -> Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Identity a) -> (a -> a) -> a -> Identity a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> a -> a
f ix
ix)
{-# INLINE iter #-}


-- | Iterate over N-dimensional space linearly from start to end in row-major fashion with an
-- accumulator
--
-- ==== __Examples__
--
-- >>> sz = Sz2 3 4
-- >>> iterLinearM sz 0 3 1 (<) 100 $ \ k ix acc -> print (fromLinearIndex sz k == ix) >> pure (acc + k)
-- True
-- True
-- True
-- 103
--
-- @since 0.1.0
iterLinearM :: (Index ix, Monad m)
            => Sz ix -- ^ Size
            -> Int -- ^ Linear start (must be non-negative)
            -> Int -- ^ Linear end (must be less than or equal to @`totalElem` sz@)
            -> Int -- ^ Increment (must not be zero)
            -> (Int -> Int -> Bool) -- ^ Continuation condition (continue if @True@)
            -> a -- ^ Accumulator
            -> (Int -> ix -> a -> m a)
            -> m a
iterLinearM :: Sz ix
-> Int
-> Int
-> Int
-> (Int -> Int -> Bool)
-> a
-> (Int -> ix -> a -> m a)
-> m a
iterLinearM !Sz ix
sz !Int
k0 !Int
k1 !Int
inc Int -> Int -> Bool
cond !a
acc Int -> ix -> a -> m a
f =
  Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
k0 (Int -> Int -> Bool
`cond` Int
k1) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
inc) a
acc ((Int -> a -> m a) -> m a) -> (Int -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \ !Int
i !a
acc0 -> Int -> ix -> a -> m a
f Int
i (Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i) a
acc0
{-# INLINE iterLinearM #-}

-- | Same as `iterLinearM`, except without an accumulator.
--
-- ==== __Examples__
--
-- >>> sz = Sz2 3 4
-- >>> iterLinearM_ sz 0 3 1 (<) $ \ k ix -> print (toLinearIndex sz ix == k)
-- True
-- True
-- True
--
-- @since 0.1.0
iterLinearM_ :: (Index ix, Monad m) =>
                Sz ix -- ^ Size
             -> Int -- ^ Start (must be non-negative)
             -> Int -- ^ End
             -> Int -- ^ Increment (must not be zero)
             -> (Int -> Int -> Bool) -- ^ Continuation condition (continue if @True@)
             -> (Int -> ix -> m ()) -- ^ Monadic action that takes index in both forms
             -> m ()
iterLinearM_ :: Sz ix
-> Int
-> Int
-> Int
-> (Int -> Int -> Bool)
-> (Int -> ix -> m ())
-> m ()
iterLinearM_ Sz ix
sz !Int
k0 !Int
k1 !Int
inc Int -> Int -> Bool
cond Int -> ix -> m ()
f =
  Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
k0 (Int -> Int -> Bool
`cond` Int
k1) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
inc) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
i -> Int -> ix -> m ()
f Int
i (Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i)
{-# INLINE iterLinearM_ #-}


-- | This is used by @INDEX_CHECK@ macro and thus used whenever the @unsafe-checks@ cabal
-- flag is on.
--
-- @since 0.4.0
indexWith ::
     Index ix
  => String -- ^ Source file name, eg. __FILE__
  -> Int -- ^ Line number in th source file, eg. __LINE__
  -> String
  -> (arr -> Sz ix) -- ^ Get size of the array
  -> (arr -> ix -> e) -- ^ Indexing function
  -> arr -- ^ Array
  -> ix -- ^ Index
  -> e
indexWith :: String
-> Int
-> String
-> (arr -> Sz ix)
-> (arr -> ix -> e)
-> arr
-> ix
-> e
indexWith String
fileName Int
lineNo String
funName arr -> Sz ix
getSize arr -> ix -> e
f arr
arr ix
ix
  | Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz ix
ix = arr -> ix -> e
f arr
arr ix
ix
  | Bool
otherwise = String -> Sz ix -> ix -> e
forall ix ix' a. (Show ix, Show ix') => String -> ix -> ix' -> a
errorIx (String
"<" String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
fileName String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
":" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
lineNo String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"> " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
funName) Sz ix
sz ix
ix
  where
    sz :: Sz ix
sz = arr -> Sz ix
getSize arr
arr

-- | Helper function for throwing out of bounds error. Used by `indexWith`
errorIx :: (Show ix, Show ix') => String -> ix -> ix' -> a
errorIx :: String -> ix -> ix' -> a
errorIx String
fName ix
sz ix'
ix =
  String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$
  String
fName String -> ShowS
forall a. [a] -> [a] -> [a]
++
  String
": Index out of bounds: (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ ix' -> String
forall a. Show a => a -> String
show ix'
ix String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") for Array of size: (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ ix -> String
forall a. Show a => a -> String
show ix
sz String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
{-# NOINLINE errorIx #-}