{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PatternSynonyms #-}

-- |
-- Module      : Data.Massiv.Core.Index
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- 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,
  iterA_,
  iterM_,
  iterLinearM,
  iterLinearM_,
  module Data.Massiv.Core.Loop,
  module Data.Massiv.Core.Index.Iterator,
  module Data.Massiv.Core.Index.Tuple,

  -- * Exceptions
  IndexException (..),
  SizeException (..),
  ShapeException (..),
  guardNumberOfElements,
  indexAssert,
  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.Iterator
import Data.Massiv.Core.Index.Ix
import Data.Massiv.Core.Index.Stride
import Data.Massiv.Core.Index.Tuple
import Data.Massiv.Core.Loop
import GHC.TypeLits

#include "massiv.h"

-- | 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 in a constant element.
    --
    -- @
    --            outside |  Array  | outside
    -- ('Fill' 0) : 0 0 0 0 | 1 2 3 4 | 0 0 0 0
    -- @
    Fill e
  | -- | 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
    -- @
    Wrap
  | -- | Replicate the element at the edge.
    --
    -- @
    --            outside |  Array  | outside
    -- 'Edge' :     1 1 1 1 | 1 2 3 4 | 4 4 4 4
    -- @
    Edge
  | -- | Mirror like reflection.
    --
    -- @
    --            outside |  Array  | outside
    -- 'Reflect' :  4 3 2 1 | 1 2 3 4 | 4 3 2 1
    -- @
    Reflect
  | -- | 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
    -- @
    Continue
  deriving (Border e -> Border e -> Bool
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
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 -> 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 :: forall ix e. Index ix => 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 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 (forall ix.
Index ix =>
Sz ix
-> ix -> (Sz Int -> Int -> Int) -> (Sz Int -> Int -> Int) -> ix
repairIndex Sz ix
sz ix
ix forall {a}. Integral a => Sz a -> a -> a
wrap forall {a}. Integral a => Sz a -> a -> a
wrap)
    Border e
Edge -> ix -> e
getVal (forall ix.
Index ix =>
Sz ix
-> ix -> (Sz Int -> Int -> Int) -> (Sz Int -> Int -> Int) -> ix
repairIndex Sz ix
sz ix
ix (forall a b. a -> b -> a
const (forall a b. a -> b -> a
const Int
0)) (\(SafeSz Int
k) Int
_ -> Int
k forall a. Num a => a -> a -> a
- Int
1))
    Border e
Reflect ->
      ix -> e
getVal
        ( 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 -> (forall a. Num a => a -> a
abs Int
i forall a. Num a => a -> a -> a
- Int
1) forall a. Integral a => a -> a -> a
`mod` Int
k)
            (\(SafeSz Int
k) !Int
i -> (-Int
i forall a. Num a => a -> a -> a
- Int
1) forall a. Integral a => a -> a -> a
`mod` Int
k)
        )
    Border e
Continue ->
      ix -> e
getVal
        ( 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 -> forall a. Num a => a -> a
abs Int
i forall a. Integral a => a -> a -> a
`mod` Int
k)
            (\(SafeSz Int
k) !Int
i -> (-Int
i forall a. Num a => a -> a -> a
- Int
2) forall a. Integral a => a -> a -> a
`mod` Int
k)
        )
  where
    wrap :: Sz a -> a -> a
wrap (SafeSz a
k) a
i = a
i 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 :: forall ix. Index ix => ix
zeroIndex = 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 :: forall ix. Index ix => ix
oneIndex = 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 :: forall ix. Index ix => Sz ix -> Bool
isNotZeroSz !Sz ix
sz = forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz 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 :: forall ix. Index ix => Sz ix -> Bool
isZeroSz = Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall ix. Index ix => Sz ix -> Sz Int
toLinearSz = coerce :: forall a b. Coercible a b => a -> b
coerce forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall ix. Index ix => ix -> Int
headDim = forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall ix. Index ix => ix -> Lower ix
tailDim = forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall ix. Index ix => ix -> Int
lastDim = forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall ix. Index ix => ix -> Lower ix
initDim = forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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' :: forall ix. (HasCallStack, Index ix) => ix -> Dim -> Int -> ix
setDim' ix
ix Dim
dim = forall a. HasCallStack => Either SomeException a -> a
throwEither forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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' :: forall ix. (HasCallStack, Index ix) => ix -> Dim -> Int
getDim' ix
ix = forall a. HasCallStack => Either SomeException a -> a
throwEither forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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' :: forall ix.
(HasCallStack, Index ix) =>
ix -> Dim -> (Int -> Int) -> (Int, ix)
modifyDim' ix
ix Dim
dim = forall a. HasCallStack => Either SomeException a -> a
throwEither forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
ix -> Dim -> m (Lower ix)
dropDimM ix
ix = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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' :: forall ix. (HasCallStack, Index ix) => ix -> Dim -> Lower ix
dropDim' ix
ix = forall a. HasCallStack => Either SomeException a -> a
throwEither forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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' :: forall ix. (HasCallStack, Index ix) => ix -> Dim -> (Int, Lower ix)
pullOutDim' ix
ix = forall a. HasCallStack => Either SomeException a -> a
throwEither forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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' :: forall ix. (HasCallStack, Index ix) => Lower ix -> Dim -> Int -> ix
insertDim' Lower ix
ix Dim
dim = forall a. HasCallStack => Either SomeException a -> a
throwEither forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall (n :: Nat). KnownNat n => Dimension n -> Dim
fromDimension = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall ix (n :: Nat).
IsIndexDimension ix n =>
ix -> Dimension n -> Int -> ix
setDimension ix
ix = forall ix. (HasCallStack, Index ix) => ix -> Dim -> Int -> ix
setDim' ix
ix forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall ix (n :: Nat).
IsIndexDimension ix n =>
ix -> Dimension n -> (Int -> Int) -> (Int, ix)
modifyDimension ix
ix = forall ix.
(HasCallStack, Index ix) =>
ix -> Dim -> (Int -> Int) -> (Int, ix)
modifyDim' ix
ix forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall ix (n :: Nat).
IsIndexDimension ix n =>
ix -> Dimension n -> Int
getDimension ix
ix = forall ix. (HasCallStack, Index ix) => ix -> Dim -> Int
getDim' ix
ix forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall ix (n :: Nat).
IsIndexDimension ix n =>
ix -> Dimension n -> Lower ix
dropDimension ix
ix = forall ix. (HasCallStack, Index ix) => ix -> Dim -> Lower ix
dropDim' ix
ix forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall ix (n :: Nat).
IsIndexDimension ix n =>
ix -> Dimension n -> (Int, Lower ix)
pullOutDimension ix
ix = forall ix. (HasCallStack, Index ix) => ix -> Dim -> (Int, Lower ix)
pullOutDim' ix
ix forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall ix (n :: Nat).
IsIndexDimension ix n =>
Lower ix -> Dimension n -> Int -> ix
insertDimension Lower ix
ix = forall ix. (HasCallStack, Index ix) => Lower ix -> Dim -> Int -> ix
insertDim' Lower ix
ix forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall ix a.
Index ix =>
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 =
  forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$ 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 -> forall (m :: * -> *) a. Monad m => a -> m a
return 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 -> (acc + k) <$ print (fromLinearIndex sz k == ix)
-- 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 :: forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
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 =
  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) (forall a. Num a => a -> a -> a
+ Int
inc) a
acc forall a b. (a -> b) -> a -> b
$ \ !Int
i !a
acc0 -> Int -> ix -> a -> m a
f Int
i (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_ :: forall ix (m :: * -> *).
(Index ix, Monad m) =>
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 =
  forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
k0 (Int -> Int -> Bool
`cond` Int
k1) (forall a. Num a => a -> a -> a
+ Int
inc) forall a b. (a -> b) -> a -> b
$ \ !Int
i -> Int -> ix -> m ()
f Int
i (forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i)
{-# INLINE iterLinearM_ #-}

-- | This is used by the @unsafe-checks@ cabal flag.
--
-- @since 1.1.0
#ifdef MASSIV_UNSAFE_CHECKS
indexAssert :: (HasCallStack, Index ix) => String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert funName getSize f arr ix
  | isSafeIndex sz ix = f arr ix
  | otherwise = _errorIx ("<" ++ funName ++ ">") sz ix
  where
    sz = getSize arr
#else
indexAssert :: String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert :: forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert String
_funName a -> Sz ix
_getSize a -> ix -> e
f a
arr ix
ix = a -> ix -> e
f a
arr ix
ix
#endif
{-# INLINE indexAssert #-}

-- | 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 :: forall ix arr e.
Index ix =>
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
  | forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz ix
ix = arr -> ix -> e
f arr
arr ix
ix
  | Bool
otherwise = forall ix ix' a.
(HasCallStack, Show ix, Show ix') =>
String -> ix -> ix' -> a
_errorIx (String
"<" forall a. [a] -> [a] -> [a]
++ String
fileName forall a. [a] -> [a] -> [a]
++ String
":" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
lineNo forall a. [a] -> [a] -> [a]
++ String
"> " forall a. [a] -> [a] -> [a]
++ String
funName) Sz ix
sz ix
ix
  where
    sz :: Sz ix
sz = arr -> Sz ix
getSize arr
arr
{-# DEPRECATED indexWith "In favor of `indexAssert` that uses HasCallStack" #-}

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