{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-unticked-promoted-constructors #-}

-- |
-- Module      : Data.Massiv.Core.Index.Internal
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <alexey@kuleshevi.ch>
-- Stability   : experimental
-- Portability : non-portable
module Data.Massiv.Core.Index.Internal (
  Sz (SafeSz),
  pattern Sz,
  pattern Sz1,
  unSz,
  zeroSz,
  oneSz,
  liftSz,
  liftSz2,
  consSz,
  unconsSz,
  snocSz,
  unsnocSz,
  setSzM,
  insertSzM,
  pullOutSzM,
  mkSzM,
  Dim (..),
  Dimension (DimN),
  pattern Dim1,
  pattern Dim2,
  pattern Dim3,
  pattern Dim4,
  pattern Dim5,
  IsIndexDimension,
  IsDimValid,
  ReportInvalidDim,
  Lower,
  Index (..),
  iterA_,
  iterM_,
  Ix0 (..),
  type Ix1,
  pattern Ix1,
  IndexException (..),
  SizeException (..),
  ShapeException (..),
  showsPrecWrapped,
) where

import Control.DeepSeq
import Control.Exception (Exception (..), throw)
import Control.Monad (void, when)
import Control.Monad.Catch (MonadThrow (..))
import Control.Monad.ST
import Control.Scheduler
import Data.Coerce
import Data.Kind
import Data.Massiv.Core.Loop
import Data.Typeable
import GHC.TypeLits
import System.Random.Stateful

-- | `Sz` is the size of the array. It describes total number of elements along
-- each dimension in the array. It is a wrapper around an index of the same
-- dimension, however it provides type safety preventing mixup with
-- index. Moreover the @Sz@ constructor and others such as
-- `Data.Massiv.Core.Index.Sz1`, `Data.Massiv.Core.Index.Sz2`, ... that
-- are specialized to specific dimensions, prevent creation of invalid sizes with
-- negative values by clamping them to zero.
--
-- ====__Examples__
--
-- >>> import Data.Massiv.Array
-- >>> Sz (1 :> 2 :. 3)
-- Sz (1 :> 2 :. 3)
--
-- `Sz` has a `Num` instance, which is very convenient:
--
-- >>> Sz (1 :> 2 :. 3) + 5
-- Sz (6 :> 7 :. 8)
--
-- However subtraction can sometimes lead to surprising behavior, because size is not
-- allowed to take negative values it will be clamped at 0.
--
-- >>> Sz (1 :> 2 :. 3) - 2
-- Sz (0 :> 0 :. 1)
--
-- __Warning__: It is always wrong to `negate` a size, thus it will result in an
-- error. For that reason also watch out for partially applied @(`Prelude.-` sz)@, which is
-- deugared into @`negate` sz@. See more info about it in
-- [#114](https://github.com/lehins/massiv/issues/114).
--
-- @since 0.3.0
newtype Sz ix
  = -- | Safe size constructor. It is unsafe to use it without making sure that it does not contain
    -- negative components. Use `Data.Massiv.Core.Index.Sz` pattern instead.
    --
    -- @since 0.3.0
    SafeSz ix
  deriving (Sz ix -> Sz ix -> Bool
(Sz ix -> Sz ix -> Bool) -> (Sz ix -> Sz ix -> Bool) -> Eq (Sz ix)
forall ix. Eq ix => Sz ix -> Sz ix -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall ix. Eq ix => Sz ix -> Sz ix -> Bool
== :: Sz ix -> Sz ix -> Bool
$c/= :: forall ix. Eq ix => Sz ix -> Sz ix -> Bool
/= :: Sz ix -> Sz ix -> Bool
Eq, Eq (Sz ix)
Eq (Sz ix) =>
(Sz ix -> Sz ix -> Ordering)
-> (Sz ix -> Sz ix -> Bool)
-> (Sz ix -> Sz ix -> Bool)
-> (Sz ix -> Sz ix -> Bool)
-> (Sz ix -> Sz ix -> Bool)
-> (Sz ix -> Sz ix -> Sz ix)
-> (Sz ix -> Sz ix -> Sz ix)
-> Ord (Sz ix)
Sz ix -> Sz ix -> Bool
Sz ix -> Sz ix -> Ordering
Sz ix -> Sz ix -> Sz ix
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall ix. Ord ix => Eq (Sz ix)
forall ix. Ord ix => Sz ix -> Sz ix -> Bool
forall ix. Ord ix => Sz ix -> Sz ix -> Ordering
forall ix. Ord ix => Sz ix -> Sz ix -> Sz ix
$ccompare :: forall ix. Ord ix => Sz ix -> Sz ix -> Ordering
compare :: Sz ix -> Sz ix -> Ordering
$c< :: forall ix. Ord ix => Sz ix -> Sz ix -> Bool
< :: Sz ix -> Sz ix -> Bool
$c<= :: forall ix. Ord ix => Sz ix -> Sz ix -> Bool
<= :: Sz ix -> Sz ix -> Bool
$c> :: forall ix. Ord ix => Sz ix -> Sz ix -> Bool
> :: Sz ix -> Sz ix -> Bool
$c>= :: forall ix. Ord ix => Sz ix -> Sz ix -> Bool
>= :: Sz ix -> Sz ix -> Bool
$cmax :: forall ix. Ord ix => Sz ix -> Sz ix -> Sz ix
max :: Sz ix -> Sz ix -> Sz ix
$cmin :: forall ix. Ord ix => Sz ix -> Sz ix -> Sz ix
min :: Sz ix -> Sz ix -> Sz ix
Ord, Sz ix -> ()
(Sz ix -> ()) -> NFData (Sz ix)
forall ix. NFData ix => Sz ix -> ()
forall a. (a -> ()) -> NFData a
$crnf :: forall ix. NFData ix => Sz ix -> ()
rnf :: Sz ix -> ()
NFData)

-- | A safe bidirectional pattern synonym for `Sz` construction that will make sure that none of
-- the size elements are negative.
--
-- @since 0.3.0
pattern Sz :: Index ix => ix -> Sz ix
pattern $mSz :: forall {r} {ix}.
Index ix =>
Sz ix -> (ix -> r) -> ((# #) -> r) -> r
$bSz :: forall ix. Index ix => ix -> Sz ix
Sz ix <- SafeSz ix
  where
    Sz ix
ix = ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz ((Int -> Int) -> ix -> ix
forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0) ix
ix)

{-# COMPLETE Sz #-}

-- | 1-dimensional size constructor. Especially useful with literals: @(Sz1 5) == Sz (5 :: Int)@.
--
-- @since 0.3.0
pattern Sz1 :: Ix1 -> Sz Ix1
pattern $mSz1 :: forall {r}. Sz Int -> (Int -> r) -> ((# #) -> r) -> r
$bSz1 :: Int -> Sz Int
Sz1 ix <- SafeSz ix
  where
    Sz1 Int
ix = Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 Int
ix)

{-# COMPLETE Sz1 #-}

instance (UniformRange ix, Index ix) => Uniform (Sz ix) where
  uniformM :: forall g (m :: * -> *). StatefulGen g m => g -> m (Sz ix)
uniformM g
g = ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz (ix -> Sz ix) -> m ix -> m (Sz ix)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ix, ix) -> g -> m ix
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *). StatefulGen g m => (ix, ix) -> g -> m ix
uniformRM (Int -> ix
forall ix. Index ix => Int -> ix
pureIndex Int
0, Int -> ix
forall ix. Index ix => Int -> ix
pureIndex Int
forall a. Bounded a => a
maxBound) g
g
  {-# INLINE uniformM #-}

instance UniformRange ix => UniformRange (Sz ix) where
  uniformRM :: forall g (m :: * -> *).
StatefulGen g m =>
(Sz ix, Sz ix) -> g -> m (Sz ix)
uniformRM (SafeSz ix
l, SafeSz ix
u) g
g = ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz (ix -> Sz ix) -> m ix -> m (Sz ix)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ix, ix) -> g -> m ix
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *). StatefulGen g m => (ix, ix) -> g -> m ix
uniformRM (ix
l, ix
u) g
g
  {-# INLINE uniformRM #-}

instance (UniformRange ix, Index ix) => Random (Sz ix)

instance Index ix => Show (Sz ix) where
  showsPrec :: Int -> Sz ix -> ShowS
showsPrec Int
n sz :: Sz ix
sz@(SafeSz ix
usz) = Int -> ShowS -> ShowS
showsPrecWrapped Int
n (String
str String -> ShowS
forall a. [a] -> [a] -> [a]
++)
    where
      str :: String
str =
        String
"Sz"
          String -> ShowS
forall a. [a] -> [a] -> [a]
++ case Dim -> Int
unDim (Sz ix -> Dim
forall ix (proxy :: * -> *). Index ix => proxy ix -> Dim
forall (proxy :: * -> *). proxy ix -> Dim
dimensions Sz ix
sz) of
            Int
1 -> String
"1 " String -> ShowS
forall a. [a] -> [a] -> [a]
++ ix -> String
forall a. Show a => a -> String
show ix
usz
            Int
_ -> String
" (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ ix -> ShowS
forall a. Show a => a -> ShowS
shows ix
usz String
")"

-- | Calling `negate` is an error.
instance (Num ix, Index ix) => Num (Sz ix) where
  + :: Sz ix -> Sz ix -> Sz ix
(+) Sz ix
x Sz ix
y = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz (Sz ix -> ix
forall a b. Coercible a b => a -> b
coerce Sz ix
x ix -> ix -> ix
forall a. Num a => a -> a -> a
+ Sz ix -> ix
forall a b. Coercible a b => a -> b
coerce Sz ix
y)
  {-# INLINE (+) #-}
  (-) Sz ix
x Sz ix
y = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz (Sz ix -> ix
forall a b. Coercible a b => a -> b
coerce Sz ix
x ix -> ix -> ix
forall a. Num a => a -> a -> a
- Sz ix -> ix
forall a b. Coercible a b => a -> b
coerce Sz ix
y)
  {-# INLINE (-) #-}
  * :: Sz ix -> Sz ix -> Sz ix
(*) Sz ix
x Sz ix
y = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz (Sz ix -> ix
forall a b. Coercible a b => a -> b
coerce Sz ix
x ix -> ix -> ix
forall a. Num a => a -> a -> a
* Sz ix -> ix
forall a b. Coercible a b => a -> b
coerce Sz ix
y)
  {-# INLINE (*) #-}
  abs :: Sz ix -> Sz ix
abs !Sz ix
x = Sz ix
x
  {-# INLINE abs #-}
  negate :: Sz ix -> Sz ix
negate Sz ix
x
    | Sz ix
x Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix
forall ix. Index ix => Sz ix
zeroSz = Sz ix
x
    | Bool
otherwise =
        String -> Sz ix
forall a. HasCallStack => String -> a
error (String -> Sz ix) -> String -> Sz ix
forall a b. (a -> b) -> a -> b
$
          String
"Attempted to negate: "
            String -> ShowS
forall a. [a] -> [a] -> [a]
++ Sz ix -> String
forall a. Show a => a -> String
show Sz ix
x
            String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
", this can lead to unexpected behavior. See https://github.com/lehins/massiv/issues/114"
  {-# INLINE negate #-}
  signum :: Sz ix -> Sz ix
signum Sz ix
x = ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz (ix -> ix
forall a. Num a => a -> a
signum (Sz ix -> ix
forall a b. Coercible a b => a -> b
coerce Sz ix
x))
  {-# INLINE signum #-}
  fromInteger :: Integer -> Sz ix
fromInteger = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz (ix -> Sz ix) -> (Integer -> ix) -> Integer -> Sz ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> ix
forall a. Num a => Integer -> a
fromInteger
  {-# INLINE fromInteger #-}

-- | Construct size from index while checking its correctness. Throws
-- `SizeNegativeException` and `SizeOverflowException`.
--
-- @since 0.6.0
mkSzM :: (Index ix, MonadThrow m) => ix -> m (Sz ix)
mkSzM :: forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> m (Sz ix)
mkSzM ix
ix = do
  let guardNegativeOverflow :: Int -> Int -> m Int
guardNegativeOverflow Int
i !Int
acc = do
        Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ SizeException -> m ()
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (SizeException -> m ()) -> SizeException -> m ()
forall a b. (a -> b) -> a -> b
$ Sz ix -> SizeException
forall ix. Index ix => Sz ix -> SizeException
SizeNegativeException (ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz ix
ix)
        let acc' :: Int
acc' = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
acc
        Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
acc' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0 Bool -> Bool -> Bool
&& Int
acc' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
acc) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ SizeException -> m ()
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (SizeException -> m ()) -> SizeException -> m ()
forall a b. (a -> b) -> a -> b
$ Sz ix -> SizeException
forall ix. Index ix => Sz ix -> SizeException
SizeOverflowException (ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz ix
ix)
        Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
acc'
  ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ix
ix Sz ix -> m Int -> m (Sz ix)
forall a b. a -> m b -> m a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ (m Int -> Int -> m Int) -> m Int -> ix -> m Int
forall ix a. Index ix => (a -> Int -> a) -> a -> ix -> a
forall a. (a -> Int -> a) -> a -> ix -> a
foldlIndex (\m Int
acc Int
i -> m Int
acc m Int -> (Int -> m Int) -> m Int
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> Int -> m Int
guardNegativeOverflow Int
i) (Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
1) ix
ix
{-# INLINE mkSzM #-}

-- | Function for unwrapping `Sz`.
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Core.Index
-- >>> unSz $ Sz3 1 2 3
-- 1 :> 2 :. 3
--
-- @since 0.3.0
unSz :: Sz ix -> ix
unSz :: forall ix. Sz ix -> ix
unSz (SafeSz ix
ix) = ix
ix
{-# INLINE unSz #-}

-- | An empty size with all elements in size equal to @0@.
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Core.Index
-- >>> zeroSz :: Sz5
-- Sz (0 :> 0 :> 0 :> 0 :. 0)
--
-- @since 0.3.0
zeroSz :: Index ix => Sz ix
zeroSz :: forall ix. Index ix => Sz ix
zeroSz = ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz (Int -> ix
forall ix. Index ix => Int -> ix
pureIndex Int
0)
{-# INLINE zeroSz #-}

-- | A singleton size with all elements in size equal to @1@.
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Core.Index
-- >>> oneSz :: Sz3
-- Sz (1 :> 1 :. 1)
--
-- @since 0.3.0
oneSz :: Index ix => Sz ix
oneSz :: forall ix. Index ix => Sz ix
oneSz = ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz (Int -> ix
forall ix. Index ix => Int -> ix
pureIndex Int
1)
{-# INLINE oneSz #-}

-- | Same as `liftIndex`, but for `Sz`
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Core.Index
-- >>> liftSz succ (Sz2 2 3)
-- Sz (3 :. 4)
--
-- @since 0.4.0
liftSz :: Index ix => (Int -> Int) -> Sz ix -> Sz ix
liftSz :: forall ix. Index ix => (Int -> Int) -> Sz ix -> Sz ix
liftSz Int -> Int
f (SafeSz ix
ix) = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((Int -> Int) -> ix -> ix
forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex Int -> Int
f ix
ix)
{-# INLINE liftSz #-}

-- | Same as `liftIndex2`, but for `Sz`
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Core.Index
-- >>> liftSz2 (-) (Sz2 2 3) (Sz2 3 1)
-- Sz (0 :. 2)
--
-- @since 0.4.3
liftSz2 :: Index ix => (Int -> Int -> Int) -> Sz ix -> Sz ix -> Sz ix
liftSz2 :: forall ix.
Index ix =>
(Int -> Int -> Int) -> Sz ix -> Sz ix -> Sz ix
liftSz2 Int -> Int -> Int
f Sz ix
sz1 Sz ix
sz2 = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
f (Sz ix -> ix
forall a b. Coercible a b => a -> b
coerce Sz ix
sz1) (Sz ix -> ix
forall a b. Coercible a b => a -> b
coerce Sz ix
sz2))
{-# INLINE liftSz2 #-}

-- | Same as `consDim`, but for `Sz`
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Core.Index
-- >>> consSz (Sz1 1) (Sz2 2 3) :: Sz3
-- Sz (1 :> 2 :. 3)
--
-- @since 0.3.0
consSz :: Index ix => Sz Ix1 -> Sz (Lower ix) -> Sz ix
consSz :: forall ix. Index ix => Sz Int -> Sz (Lower ix) -> Sz ix
consSz (SafeSz Int
i) (SafeSz Lower ix
ix) = ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz (Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
i Lower ix
ix)
{-# INLINE consSz #-}

-- | Same as `snocDim`, but for `Sz`
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Core.Index
-- >>> snocSz (Sz2 2 3) (Sz1 1) :: Sz3
-- Sz (2 :> 3 :. 1)
--
-- @since 0.3.0
snocSz :: Index ix => Sz (Lower ix) -> Sz Ix1 -> Sz ix
snocSz :: forall ix. Index ix => Sz (Lower ix) -> Sz Int -> Sz ix
snocSz (SafeSz Lower ix
i) (SafeSz Int
ix) = ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz (Lower ix -> Int -> ix
forall ix. Index ix => Lower ix -> Int -> ix
snocDim Lower ix
i Int
ix)
{-# INLINE snocSz #-}

-- | Same as `setDimM`, but for `Sz`
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Core.Index
-- >>> setSzM (Sz2 2 3) 2 (Sz1 1) :: IO Sz2
-- Sz (1 :. 3)
-- >>> setSzM (Sz2 2 3) 3 (Sz1 1) :: IO Sz2
-- *** Exception: IndexDimensionException: (Dim 3) for (2 :. 3)
--
-- @since 0.3.0
setSzM :: (MonadThrow m, Index ix) => Sz ix -> Dim -> Sz Int -> m (Sz ix)
setSzM :: forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz ix -> Dim -> Sz Int -> m (Sz ix)
setSzM (SafeSz ix
sz) Dim
dim (SafeSz Int
sz1) = ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz (ix -> Sz ix) -> m ix -> m (Sz ix)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ix -> Dim -> Int -> m ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
forall (m :: * -> *). MonadThrow m => ix -> Dim -> Int -> m ix
setDimM ix
sz Dim
dim Int
sz1
{-# INLINE setSzM #-}

-- | Same as `insertDimM`, but for `Sz`
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Core.Index
-- >>> insertSzM (Sz2 2 3) 3 (Sz1 1) :: IO Sz3
-- Sz (1 :> 2 :. 3)
-- >>> insertSzM (Sz2 2 3) 4 (Sz1 1) :: IO Sz3
-- *** Exception: IndexDimensionException: (Dim 4) for (2 :. 3)
--
-- @since 0.3.0
insertSzM :: (MonadThrow m, Index ix) => Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
insertSzM :: forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix)
insertSzM (SafeSz Lower ix
sz) Dim
dim (SafeSz Int
sz1) = ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz (ix -> Sz ix) -> m ix -> m (Sz ix)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lower ix -> Dim -> Int -> m ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
Lower ix -> Dim -> Int -> m ix
forall (m :: * -> *).
MonadThrow m =>
Lower ix -> Dim -> Int -> m ix
insertDimM Lower ix
sz Dim
dim Int
sz1
{-# INLINE insertSzM #-}

-- | Same as `unconsDim`, but for `Sz`
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Core.Index
-- >>> unconsSz $ Sz3 1 2 3
-- (Sz1 1,Sz (2 :. 3))
--
-- @since 0.3.0
unconsSz :: Index ix => Sz ix -> (Sz Ix1, Sz (Lower ix))
unconsSz :: forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz (SafeSz ix
sz) = (Int, Lower ix) -> (Sz Int, Sz (Lower ix))
forall a b. Coercible a b => a -> b
coerce (ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
sz)
{-# INLINE unconsSz #-}

-- | Same as `unsnocDim`, but for `Sz`
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Core.Index
-- >>> unsnocSz $ Sz3 1 2 3
-- (Sz (1 :. 2),Sz1 3)
--
-- @since 0.3.0
unsnocSz :: Index ix => Sz ix -> (Sz (Lower ix), Sz Ix1)
unsnocSz :: forall ix. Index ix => Sz ix -> (Sz (Lower ix), Sz Int)
unsnocSz (SafeSz ix
sz) = (Lower ix, Int) -> (Sz (Lower ix), Sz Int)
forall a b. Coercible a b => a -> b
coerce (ix -> (Lower ix, Int)
forall ix. Index ix => ix -> (Lower ix, Int)
unsnocDim ix
sz)
{-# INLINE unsnocSz #-}

-- | Same as `pullOutDim`, but for `Sz`
--
-- >>> import Data.Massiv.Core.Index
-- >>> pullOutSzM (Sz3 1 2 3) 3
-- (Sz1 1,Sz (2 :. 3))
-- >>> pullOutSzM (Sz3 1 2 3) 0
-- *** Exception: IndexDimensionException: (Dim 0) for (1 :> 2 :. 3)
--
-- @since 0.3.0
pullOutSzM :: (MonadThrow m, Index ix) => Sz ix -> Dim -> m (Sz Ix1, Sz (Lower ix))
pullOutSzM :: forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz ix -> Dim -> m (Sz Int, Sz (Lower ix))
pullOutSzM (SafeSz ix
sz) = ((Int, Lower ix) -> (Sz Int, Sz (Lower ix)))
-> m (Int, Lower ix) -> m (Sz Int, Sz (Lower ix))
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int, Lower ix) -> (Sz Int, Sz (Lower ix))
forall a b. Coercible a b => a -> b
coerce (m (Int, Lower ix) -> m (Sz Int, Sz (Lower ix)))
-> (Dim -> m (Int, Lower ix)) -> Dim -> m (Sz Int, Sz (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)
forall (m :: * -> *).
MonadThrow m =>
ix -> Dim -> m (Int, Lower ix)
pullOutDimM ix
sz
{-# INLINE pullOutSzM #-}

-- | A way to select Array dimension at a value level.
--
-- @since 0.1.0
newtype Dim = Dim {Dim -> Int
unDim :: Int} deriving (Dim -> Dim -> Bool
(Dim -> Dim -> Bool) -> (Dim -> Dim -> Bool) -> Eq Dim
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Dim -> Dim -> Bool
== :: Dim -> Dim -> Bool
$c/= :: Dim -> Dim -> Bool
/= :: Dim -> Dim -> Bool
Eq, Eq Dim
Eq Dim =>
(Dim -> Dim -> Ordering)
-> (Dim -> Dim -> Bool)
-> (Dim -> Dim -> Bool)
-> (Dim -> Dim -> Bool)
-> (Dim -> Dim -> Bool)
-> (Dim -> Dim -> Dim)
-> (Dim -> Dim -> Dim)
-> Ord Dim
Dim -> Dim -> Bool
Dim -> Dim -> Ordering
Dim -> Dim -> Dim
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Dim -> Dim -> Ordering
compare :: Dim -> Dim -> Ordering
$c< :: Dim -> Dim -> Bool
< :: Dim -> Dim -> Bool
$c<= :: Dim -> Dim -> Bool
<= :: Dim -> Dim -> Bool
$c> :: Dim -> Dim -> Bool
> :: Dim -> Dim -> Bool
$c>= :: Dim -> Dim -> Bool
>= :: Dim -> Dim -> Bool
$cmax :: Dim -> Dim -> Dim
max :: Dim -> Dim -> Dim
$cmin :: Dim -> Dim -> Dim
min :: Dim -> Dim -> Dim
Ord, Integer -> Dim
Dim -> Dim
Dim -> Dim -> Dim
(Dim -> Dim -> Dim)
-> (Dim -> Dim -> Dim)
-> (Dim -> Dim -> Dim)
-> (Dim -> Dim)
-> (Dim -> Dim)
-> (Dim -> Dim)
-> (Integer -> Dim)
-> Num Dim
forall a.
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> a)
-> (a -> a)
-> (Integer -> a)
-> Num a
$c+ :: Dim -> Dim -> Dim
+ :: Dim -> Dim -> Dim
$c- :: Dim -> Dim -> Dim
- :: Dim -> Dim -> Dim
$c* :: Dim -> Dim -> Dim
* :: Dim -> Dim -> Dim
$cnegate :: Dim -> Dim
negate :: Dim -> Dim
$cabs :: Dim -> Dim
abs :: Dim -> Dim
$csignum :: Dim -> Dim
signum :: Dim -> Dim
$cfromInteger :: Integer -> Dim
fromInteger :: Integer -> Dim
Num, Num Dim
Ord Dim
(Num Dim, Ord Dim) => (Dim -> Rational) -> Real Dim
Dim -> Rational
forall a. (Num a, Ord a) => (a -> Rational) -> Real a
$ctoRational :: Dim -> Rational
toRational :: Dim -> Rational
Real, Enum Dim
Real Dim
(Real Dim, Enum Dim) =>
(Dim -> Dim -> Dim)
-> (Dim -> Dim -> Dim)
-> (Dim -> Dim -> Dim)
-> (Dim -> Dim -> Dim)
-> (Dim -> Dim -> (Dim, Dim))
-> (Dim -> Dim -> (Dim, Dim))
-> (Dim -> Integer)
-> Integral Dim
Dim -> Integer
Dim -> Dim -> (Dim, Dim)
Dim -> Dim -> Dim
forall a.
(Real a, Enum a) =>
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> (a, a))
-> (a -> a -> (a, a))
-> (a -> Integer)
-> Integral a
$cquot :: Dim -> Dim -> Dim
quot :: Dim -> Dim -> Dim
$crem :: Dim -> Dim -> Dim
rem :: Dim -> Dim -> Dim
$cdiv :: Dim -> Dim -> Dim
div :: Dim -> Dim -> Dim
$cmod :: Dim -> Dim -> Dim
mod :: Dim -> Dim -> Dim
$cquotRem :: Dim -> Dim -> (Dim, Dim)
quotRem :: Dim -> Dim -> (Dim, Dim)
$cdivMod :: Dim -> Dim -> (Dim, Dim)
divMod :: Dim -> Dim -> (Dim, Dim)
$ctoInteger :: Dim -> Integer
toInteger :: Dim -> Integer
Integral, Int -> Dim
Dim -> Int
Dim -> [Dim]
Dim -> Dim
Dim -> Dim -> [Dim]
Dim -> Dim -> Dim -> [Dim]
(Dim -> Dim)
-> (Dim -> Dim)
-> (Int -> Dim)
-> (Dim -> Int)
-> (Dim -> [Dim])
-> (Dim -> Dim -> [Dim])
-> (Dim -> Dim -> [Dim])
-> (Dim -> Dim -> Dim -> [Dim])
-> Enum Dim
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
$csucc :: Dim -> Dim
succ :: Dim -> Dim
$cpred :: Dim -> Dim
pred :: Dim -> Dim
$ctoEnum :: Int -> Dim
toEnum :: Int -> Dim
$cfromEnum :: Dim -> Int
fromEnum :: Dim -> Int
$cenumFrom :: Dim -> [Dim]
enumFrom :: Dim -> [Dim]
$cenumFromThen :: Dim -> Dim -> [Dim]
enumFromThen :: Dim -> Dim -> [Dim]
$cenumFromTo :: Dim -> Dim -> [Dim]
enumFromTo :: Dim -> Dim -> [Dim]
$cenumFromThenTo :: Dim -> Dim -> Dim -> [Dim]
enumFromThenTo :: Dim -> Dim -> Dim -> [Dim]
Enum, Dim -> ()
(Dim -> ()) -> NFData Dim
forall a. (a -> ()) -> NFData a
$crnf :: Dim -> ()
rnf :: Dim -> ()
NFData)

instance Show Dim where
  show :: Dim -> String
show (Dim Int
d) = String
"(Dim " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
d String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

instance Uniform Dim where
  uniformM :: forall g (m :: * -> *). StatefulGen g m => g -> m Dim
uniformM g
g = Int -> Dim
Dim (Int -> Dim) -> m Int -> m Dim
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int, Int) -> g -> m Int
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *). StatefulGen g m => (Int, Int) -> g -> m Int
uniformRM (Int
1, Int
forall a. Bounded a => a
maxBound) g
g

instance UniformRange Dim where
  uniformRM :: forall g (m :: * -> *). StatefulGen g m => (Dim, Dim) -> g -> m Dim
uniformRM (Dim, Dim)
r g
g = Int -> Dim
Dim (Int -> Dim) -> m Int -> m Dim
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int, Int) -> g -> m Int
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *). StatefulGen g m => (Int, Int) -> g -> m Int
uniformRM ((Dim, Dim) -> (Int, Int)
forall a b. Coercible a b => a -> b
coerce (Dim, Dim)
r) g
g

instance Random Dim

-- | A way to select Array dimension at a type level.
--
-- @since 0.2.4
data Dimension (n :: Nat) where
  DimN :: (1 <= n, KnownNat n) => Dimension n

-- | Construct 1st dimension
--
-- @since 0.2.4
pattern Dim1 :: Dimension 1
pattern $mDim1 :: forall {r}. Dimension 1 -> ((# #) -> r) -> ((# #) -> r) -> r
$bDim1 :: Dimension 1
Dim1 = DimN

-- | Construct 2nd dimension
--
-- @since 0.2.4
pattern Dim2 :: Dimension 2
pattern $mDim2 :: forall {r}. Dimension 2 -> ((# #) -> r) -> ((# #) -> r) -> r
$bDim2 :: Dimension 2
Dim2 = DimN

-- | Construct 3rd dimension
--
-- @since 0.2.4
pattern Dim3 :: Dimension 3
pattern $mDim3 :: forall {r}. Dimension 3 -> ((# #) -> r) -> ((# #) -> r) -> r
$bDim3 :: Dimension 3
Dim3 = DimN

-- | Construct 4th dimension
--
-- @since 0.2.4
pattern Dim4 :: Dimension 4
pattern $mDim4 :: forall {r}. Dimension 4 -> ((# #) -> r) -> ((# #) -> r) -> r
$bDim4 :: Dimension 4
Dim4 = DimN

-- | Construct 5th dimension
--
-- @since 0.2.4
pattern Dim5 :: Dimension 5
pattern $mDim5 :: forall {r}. Dimension 5 -> ((# #) -> r) -> ((# #) -> r) -> r
$bDim5 :: Dimension 5
Dim5 = DimN

-- | A type level constraint that ensures index is indeed valid and that supplied dimension can be
-- safely used with it.
--
-- @since 0.2.4
type IsIndexDimension ix n = (1 <= n, n <= Dimensions ix, Index ix, KnownNat n)

-- | This type family will always point to a type for a dimension that is one lower than the type
-- argument.
--
-- @since 0.1.0
type family Lower ix :: Type

type family ReportInvalidDim (dims :: Nat) (n :: Nat) isNotZero isLess :: Bool where
  ReportInvalidDim dims n True True = True
  ReportInvalidDim dims n True False =
    TypeError
      ( Text "Dimension "
          :<>: ShowType n
          :<>: Text " is higher than "
          :<>: Text "the maximum expected "
          :<>: ShowType dims
      )
  ReportInvalidDim dims n False isLess =
    TypeError (Text "Zero dimensional indices are not supported")

type family IsDimValid ix n :: Bool where
  IsDimValid ix n = ReportInvalidDim (Dimensions ix) n (1 <=? n) (n <=? Dimensions ix)

-- | This is bread and butter of multi-dimensional array indexing. It is unlikely that any of the
-- functions in this class will be useful to a regular user, unless general algorithms are being
-- implemented that do span multiple dimensions.
class
  ( Eq ix
  , Ord ix
  , Show ix
  , NFData ix
  , Typeable ix
  , Eq (Lower ix)
  , Ord (Lower ix)
  , Show (Lower ix)
  , NFData (Lower ix)
  , KnownNat (Dimensions ix)
  ) =>
  Index ix
  where
  -- | Type level information on how many dimensions this index has.
  --
  -- @since 0.2.0
  type Dimensions ix :: Nat

  -- | What is the dimensionality of this index.
  --
  -- @since 0.2.0
  dimensions :: proxy ix -> Dim

  -- | Total number of elements in an array of this size.
  --
  -- @since 0.1.0
  totalElem :: Sz ix -> Int

  -- | Prepend a dimension to the index
  --
  -- @since 0.1.0
  consDim :: Int -> Lower ix -> ix

  -- | Take a dimension from the index from the outside
  --
  -- @since 0.1.0
  unconsDim :: ix -> (Int, Lower ix)

  -- | Apppend a dimension to the index
  --
  -- @since 0.1.0
  snocDim :: Lower ix -> Int -> ix

  -- | Take a dimension from the index from the inside
  --
  -- @since 0.1.0
  unsnocDim :: ix -> (Lower ix, Int)

  -- | Pull out value at specified dimension from the index, thus also lowering it dimensionality.
  --
  -- @since 0.2.5
  pullOutDimM :: MonadThrow m => ix -> Dim -> m (Int, Lower ix)

  -- | Insert a dimension into the index
  insertDimM :: MonadThrow m => Lower ix -> Dim -> Int -> m ix

  -- | Extract the value index has at specified dimension.
  --
  -- @since 0.3.0
  getDimM :: MonadThrow m => ix -> Dim -> m Int
  getDimM ix
ix Dim
dim = (Int, ix) -> Int
forall a b. (a, b) -> a
fst ((Int, ix) -> Int) -> m (Int, ix) -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ix -> Dim -> (Int -> Int) -> m (Int, ix)
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> (Int -> Int) -> m (Int, ix)
forall (m :: * -> *).
MonadThrow m =>
ix -> Dim -> (Int -> Int) -> m (Int, ix)
modifyDimM ix
ix Dim
dim Int -> Int
forall a. a -> a
id
  {-# INLINE [1] getDimM #-}

  -- | Set the value for an index at specified dimension.
  --
  -- @since 0.3.0
  setDimM :: MonadThrow m => ix -> Dim -> Int -> m ix
  setDimM ix
ix Dim
dim Int
i = (Int, ix) -> ix
forall a b. (a, b) -> b
snd ((Int, ix) -> ix) -> m (Int, ix) -> m ix
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ix -> Dim -> (Int -> Int) -> m (Int, ix)
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> (Int -> Int) -> m (Int, ix)
forall (m :: * -> *).
MonadThrow m =>
ix -> Dim -> (Int -> Int) -> m (Int, ix)
modifyDimM ix
ix Dim
dim (Int -> Int -> Int
forall a b. a -> b -> a
const Int
i)
  {-# INLINE [1] setDimM #-}

  -- | Update the value for an index at specified dimension and return the old value as
  -- well as the updated index.
  --
  -- @since 0.4.1
  modifyDimM :: MonadThrow m => ix -> Dim -> (Int -> Int) -> m (Int, ix)
  modifyDimM ix
ix Dim
dim Int -> Int
f = do
    Int
i <- ix -> Dim -> m Int
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m Int
forall (m :: * -> *). MonadThrow m => ix -> Dim -> m Int
getDimM ix
ix Dim
dim
    ix
ix' <- ix -> Dim -> Int -> m ix
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
forall (m :: * -> *). MonadThrow m => ix -> Dim -> Int -> m ix
setDimM ix
ix Dim
dim (Int -> Int
f Int
i)
    (Int, ix) -> m (Int, ix)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
i, ix
ix')
  {-# INLINE [1] modifyDimM #-}

  -- | Lift an `Int` to any index by replicating the value as many times as there are dimensions.
  --
  -- @since 0.1.0
  pureIndex :: Int -> ix

  -- | Zip together two indices with a function
  --
  -- @since 0.1.0
  liftIndex2 :: (Int -> Int -> Int) -> ix -> ix -> ix

  -- | Map a function over an index
  --
  -- @since 0.1.0
  liftIndex :: (Int -> Int) -> ix -> ix
  liftIndex Int -> Int
f = (Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (\Int
_ Int
i -> Int -> Int
f Int
i) (Int -> ix
forall ix. Index ix => Int -> ix
pureIndex Int
0)
  {-# INLINE [1] liftIndex #-}

  -- | Perform a left fold over the index
  foldlIndex :: (a -> Int -> a) -> a -> ix -> a
  default foldlIndex
    :: Index (Lower ix)
    => (a -> Int -> a)
    -> a
    -> ix
    -> a
  foldlIndex a -> Int -> a
f !a
acc !ix
ix = (a -> Int -> a) -> a -> Lower ix -> a
forall ix a. Index ix => (a -> Int -> a) -> a -> ix -> a
forall a. (a -> Int -> a) -> a -> Lower ix -> a
foldlIndex a -> Int -> a
f (a -> Int -> a
f a
acc Int
i0) Lower ix
ixL
    where
      !(Int
i0, Lower ix
ixL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ix
  {-# INLINE [1] foldlIndex #-}

  -- TODO: implement in terms of foldlIndex and pull out of the class

  -- | Check whether index is positive and is within the size.
  --
  -- @since 0.1.0
  isSafeIndex
    :: Sz ix
    -- ^ Size
    -> ix
    -- ^ Index
    -> Bool
  default isSafeIndex
    :: Index (Lower ix)
    => Sz ix
    -> ix
    -> Bool
  isSafeIndex Sz ix
sz !ix
ix = Sz Int -> Int -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz Int
n0 Int
i0 Bool -> Bool -> Bool
&& Sz (Lower ix) -> Lower ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz (Lower ix)
szL Lower ix
ixL
    where
      !(Sz Int
n0, 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
i0, Lower ix
ixL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ix
  {-# INLINE [1] isSafeIndex #-}

  -- | Convert linear index from size and index
  --
  -- @since 0.1.0
  toLinearIndex
    :: Sz ix
    -- ^ Size
    -> ix
    -- ^ Index
    -> Ix1
  default toLinearIndex :: Index (Lower ix) => Sz ix -> ix -> Ix1
  toLinearIndex (SafeSz ix
sz) !ix
ix = Sz (Lower ix) -> Lower ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex (Lower ix -> Sz (Lower ix)
forall ix. ix -> Sz ix
SafeSz Lower ix
szL) Lower ix
ixL Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i
    where
      !(Lower ix
szL, Int
n) = ix -> (Lower ix, Int)
forall ix. Index ix => ix -> (Lower ix, Int)
unsnocDim ix
sz
      !(Lower ix
ixL, Int
i) = ix -> (Lower ix, Int)
forall ix. Index ix => ix -> (Lower ix, Int)
unsnocDim ix
ix
  {-# INLINE [1] toLinearIndex #-}

  -- | Convert linear index from size and index with an accumulator. Currently is useless and will
  -- likely be removed in future versions.
  --
  -- @since 0.1.0
  toLinearIndexAcc :: Ix1 -> ix -> ix -> Ix1
  default toLinearIndexAcc :: Index (Lower ix) => Ix1 -> ix -> ix -> Ix1
  toLinearIndexAcc !Int
acc !ix
sz !ix
ix = Int -> Lower ix -> Lower ix -> Int
forall ix. Index ix => Int -> ix -> ix -> Int
toLinearIndexAcc (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) Lower ix
szL Lower ix
ixL
    where
      !(Int
n, Lower ix
szL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
sz
      !(Int
i, Lower ix
ixL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ix
  {-# INLINE [1] toLinearIndexAcc #-}

  -- | Compute an index from size and linear index
  --
  -- @since 0.1.0
  fromLinearIndex :: Sz ix -> Ix1 -> ix
  default fromLinearIndex :: Index (Lower ix) => Sz ix -> Ix1 -> ix
  fromLinearIndex (SafeSz ix
sz) !Int
k = Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
q Lower ix
ixL
    where
      !(!Int
q, !Lower ix
ixL) = Lower ix -> Int -> (Int, Lower ix)
forall ix. Index ix => ix -> Int -> (Int, ix)
fromLinearIndexAcc ((Int, Lower ix) -> Lower ix
forall a b. (a, b) -> b
snd (ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
sz)) Int
k
  {-# INLINE [1] fromLinearIndex #-}

  -- | Compute an index from size and linear index using an accumulator, thus trying to optimize for
  -- tail recursion while getting the index computed.
  --
  -- @since 0.1.0
  fromLinearIndexAcc :: ix -> Ix1 -> (Int, ix)
  default fromLinearIndexAcc :: Index (Lower ix) => ix -> Ix1 -> (Ix1, ix)
  fromLinearIndexAcc !ix
ix' !Int
k = (Int
q, Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
r Lower ix
ixL)
    where
      !(!Int
m, !Lower ix
ix) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ix'
      !(!Int
kL, !Lower ix
ixL) = Lower ix -> Int -> (Int, Lower ix)
forall ix. Index ix => ix -> Int -> (Int, ix)
fromLinearIndexAcc Lower ix
ix Int
k
      !(!Int
q, !Int
r) = Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
quotRem Int
kL Int
m
  {-# INLINE [1] fromLinearIndexAcc #-}

  -- | A way to make sure index is withing the bounds for the supplied size. Takes two functions
  -- that will be invoked whenever index (2nd arg) is outsize the supplied size (1st arg)
  --
  -- @since 0.1.0
  repairIndex
    :: Sz ix
    -- ^ Size
    -> ix
    -- ^ Index
    -> (Sz Int -> Int -> Int)
    -- ^ Repair when below zero
    -> (Sz Int -> Int -> Int)
    -- ^ Repair when higher than size
    -> ix
  default repairIndex
    :: Index (Lower ix)
    => Sz ix
    -> ix
    -> (Sz Int -> Int -> Int)
    -> (Sz Int -> Int -> Int)
    -> ix
  repairIndex Sz ix
sz !ix
ix Sz Int -> Int -> Int
rBelow Sz Int -> Int -> Int
rOver =
    Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim (Sz Int
-> Int -> (Sz Int -> Int -> Int) -> (Sz Int -> Int -> Int) -> Int
forall ix.
Index ix =>
Sz ix
-> ix -> (Sz Int -> Int -> Int) -> (Sz Int -> Int -> Int) -> ix
repairIndex Sz Int
n Int
i Sz Int -> Int -> Int
rBelow Sz Int -> Int -> Int
rOver) (Sz (Lower ix)
-> Lower ix
-> (Sz Int -> Int -> Int)
-> (Sz Int -> Int -> Int)
-> Lower ix
forall ix.
Index ix =>
Sz ix
-> ix -> (Sz Int -> Int -> Int) -> (Sz Int -> Int -> Int) -> ix
repairIndex Sz (Lower ix)
szL Lower ix
ixL Sz Int -> Int -> Int
rBelow Sz Int -> Int -> Int
rOver)
    where
      !(Sz Int
n, 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
i, Lower ix
ixL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ix
  {-# INLINE [1] repairIndex #-}

  -- | This function is what makes it possible to iterate over an array of any dimension.
  --
  -- @since 0.1.0
  iterM
    :: Monad m
    => ix
    -- ^ Start index
    -> ix
    -- ^ End index
    -> ix
    -- ^ Increment
    -> (Int -> Int -> Bool)
    -- ^ Continue iterating while predicate is True (eg. until end of row)
    -> a
    -- ^ Initial value for an accumulator
    -> (ix -> a -> m a)
    -- ^ Accumulator function
    -> m a
  default iterM
    :: (Index (Lower 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 -> 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
s (Int -> Int -> Bool
`cond` Int
e) (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 ->
      Lower ix
-> Lower ix
-> Lower ix
-> (Int -> Int -> Bool)
-> a
-> (Lower ix -> a -> m a)
-> m a
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix
-> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a
forall (m :: * -> *) a.
Monad m =>
Lower ix
-> Lower ix
-> Lower ix
-> (Int -> Int -> Bool)
-> a
-> (Lower ix -> a -> m a)
-> m a
iterM Lower ix
sIxL Lower ix
eIxL Lower ix
incIxL Int -> Int -> Bool
cond a
acc0 ((Lower ix -> a -> m a) -> m a) -> (Lower ix -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \ !Lower ix
ix -> ix -> a -> m a
f (Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
i Lower ix
ix)
    where
      !(Int
s, Lower ix
sIxL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
sIx
      !(Int
e, Lower ix
eIxL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
eIx
      !(Int
inc, Lower ix
incIxL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
incIx
  {-# INLINE iterM #-}

  iterRowMajorST
    :: Int
    -- ^ Scheduler multiplying factor. Must be positive
    -> Scheduler s a
    -- ^ Scheduler to use
    -> ix
    -- ^ Start index
    -> ix
    -- ^ Stride
    -> Sz ix
    -- ^ Size
    -> a
    -- ^ Initial accumulator
    -> (a -> ST s (a, a))
    -- ^ Function that splits accumulator for each scheduled job.
    -> (ix -> a -> ST s a)
    -- ^ Action
    -> ST s a
  default iterRowMajorST
    :: Index (Lower ix)
    => Int
    -> Scheduler s a
    -> ix
    -> ix
    -> Sz ix
    -> a
    -> (a -> ST s (a, a))
    -> (ix -> a -> ST s a)
    -> ST s a
  iterRowMajorST !Int
fact Scheduler s a
scheduler ix
ixStart ix
ixStride Sz ix
sz a
initAcc a -> ST s (a, a)
splitAcc ix -> a -> ST s a
f = do
    let !(SafeSz Int
n, szL :: Sz (Lower ix)
szL@(SafeSz Lower ix
nL)) = Sz ix -> (Sz Int, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz Sz ix
sz
    if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
      then do
        let !(!Int
start, !Lower ix
ixL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStart
            !(!Int
stride, !Lower ix
sL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStride
        if Scheduler s a -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s a
scheduler Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 Bool -> Bool -> Bool
&& Int
fact Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Scheduler s a -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s a
scheduler Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
fact
          then do
            let !newFact :: Int
newFact = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
fact Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
n)
            Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> a -> ST s a)
-> ST s a
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
start (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
stride) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
stride) a
initAcc ((Int -> a -> ST s a) -> ST s a) -> (Int -> a -> ST s a) -> ST s a
forall a b. (a -> b) -> a -> b
$ \Int
j a
acc ->
              Int
-> Scheduler s a
-> Lower ix
-> Lower ix
-> Sz (Lower ix)
-> a
-> (a -> ST s (a, a))
-> (Lower ix -> a -> ST s a)
-> ST s a
forall ix s a.
Index ix =>
Int
-> Scheduler s a
-> ix
-> ix
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (ix -> a -> ST s a)
-> ST s a
forall s a.
Int
-> Scheduler s a
-> Lower ix
-> Lower ix
-> Sz (Lower ix)
-> a
-> (a -> ST s (a, a))
-> (Lower ix -> a -> ST s a)
-> ST s a
iterRowMajorST Int
newFact Scheduler s a
scheduler Lower ix
ixL Lower ix
sL Sz (Lower ix)
szL a
acc a -> ST s (a, a)
splitAcc (ix -> a -> ST s a
f (ix -> a -> ST s a) -> (Lower ix -> ix) -> Lower ix -> a -> ST s a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
j)
          else Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> a
-> (a -> ST s (a, a))
-> (Int -> Int -> Int -> Int -> a -> ST s a)
-> ST s a
forall s a b.
Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> b
-> (b -> ST s (b, b))
-> (Int -> Int -> Int -> Int -> b -> ST s a)
-> ST s b
splitWorkWithFactorST Int
fact Scheduler s a
scheduler Int
start Int
stride Int
n a
initAcc a -> ST s (a, a)
splitAcc ((Int -> Int -> Int -> Int -> a -> ST s a) -> ST s a)
-> (Int -> Int -> Int -> Int -> a -> ST s a) -> ST s a
forall a b. (a -> b) -> a -> b
$
            \Int
_ Int
_ Int
chunkStartAdj Int
chunkStopAdj a
acc ->
              Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> a -> ST s a)
-> ST s a
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
chunkStartAdj (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
chunkStopAdj) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
stride) a
acc ((Int -> a -> ST s a) -> ST s a) -> (Int -> a -> ST s a) -> ST s a
forall a b. (a -> b) -> a -> b
$ \Int
j a
a ->
                Lower ix
-> Lower ix
-> Lower ix
-> (Int -> Int -> Bool)
-> a
-> (Lower ix -> a -> ST s a)
-> ST s a
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix
-> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a
forall (m :: * -> *) a.
Monad m =>
Lower ix
-> Lower ix
-> Lower ix
-> (Int -> Int -> Bool)
-> a
-> (Lower ix -> a -> m a)
-> m a
iterM Lower ix
ixL Lower ix
nL Lower ix
sL Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) a
a (ix -> a -> ST s a
f (ix -> a -> ST s a) -> (Lower ix -> ix) -> Lower ix -> a -> ST s a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
j)
      else a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
initAcc
  {-# INLINE iterRowMajorST #-}

  -- | Similar to `iterM`, but no restriction on a Monad.
  --
  -- @since 1.0.2
  iterF :: ix -> ix -> ix -> (Int -> Int -> Bool) -> f a -> (ix -> f a -> f a) -> f a
  default iterF
    :: (Index (Lower ix))
    => ix
    -> ix
    -> ix
    -> (Int -> Int -> Bool)
    -> f a
    -> (ix -> f a -> f a)
    -> f a
  iterF !ix
sIx !ix
eIx !ix
incIx Int -> Int -> Bool
cond f a
initAct ix -> f a -> f a
f =
    Int
-> (Int -> Bool)
-> (Int -> Int)
-> f a
-> (Int -> f a -> f a)
-> f a
forall (f :: * -> *) a.
Int
-> (Int -> Bool)
-> (Int -> Int)
-> f a
-> (Int -> f a -> f a)
-> f a
loopF Int
s (Int -> Int -> Bool
`cond` Int
e) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
inc) f a
initAct ((Int -> f a -> f a) -> f a) -> (Int -> f a -> f a) -> f a
forall a b. (a -> b) -> a -> b
$ \ !Int
i f a
g ->
      Lower ix
-> Lower ix
-> Lower ix
-> (Int -> Int -> Bool)
-> f a
-> (Lower ix -> f a -> f a)
-> f a
forall ix (f :: * -> *) a.
Index ix =>
ix
-> ix
-> ix
-> (Int -> Int -> Bool)
-> f a
-> (ix -> f a -> f a)
-> f a
forall (f :: * -> *) a.
Lower ix
-> Lower ix
-> Lower ix
-> (Int -> Int -> Bool)
-> f a
-> (Lower ix -> f a -> f a)
-> f a
iterF Lower ix
sIxL Lower ix
eIxL Lower ix
incIxL Int -> Int -> Bool
cond f a
g (\ !Lower ix
ix -> ix -> f a -> f a
f (Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
i Lower ix
ix))
    where
      !(Int
s, Lower ix
sIxL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
sIx
      !(Int
e, Lower ix
eIxL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
eIx
      !(Int
inc, Lower ix
incIxL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
incIx
  {-# INLINE iterF #-}

  -- | A single step in iteration
  --
  -- @since 0.1.0
  stepNextMF :: ix -> ix -> ix -> (Int -> Int -> Bool) -> (Maybe ix -> f a) -> f a
  default stepNextMF
    :: (Index (Lower ix))
    => ix
    -> ix
    -> ix
    -> (Int -> Int -> Bool)
    -> (Maybe ix -> f a)
    -> f a
  stepNextMF !ix
sIx !ix
eIx !ix
incIx Int -> Int -> Bool
cond Maybe ix -> f a
f =
    Int -> (Int -> Bool) -> (Int -> Int) -> (Maybe Int -> f a) -> f a
forall (f :: * -> *) a.
Int -> (Int -> Bool) -> (Int -> Int) -> (Maybe Int -> f a) -> f a
nextMaybeF Int
s (Int -> Int -> Bool
`cond` Int
e) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
inc) ((Maybe Int -> f a) -> f a) -> (Maybe Int -> f a) -> f a
forall a b. (a -> b) -> a -> b
$ \ !Maybe Int
mni ->
      Lower ix
-> Lower ix
-> Lower ix
-> (Int -> Int -> Bool)
-> (Maybe (Lower ix) -> f a)
-> f a
forall ix (f :: * -> *) a.
Index ix =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (Maybe ix -> f a) -> f a
forall (f :: * -> *) a.
Lower ix
-> Lower ix
-> Lower ix
-> (Int -> Int -> Bool)
-> (Maybe (Lower ix) -> f a)
-> f a
stepNextMF Lower ix
sIxL Lower ix
eIxL Lower ix
incIxL Int -> Int -> Bool
cond ((Maybe (Lower ix) -> f a) -> f a)
-> (Maybe (Lower ix) -> f a) -> f a
forall a b. (a -> b) -> a -> b
$ \Maybe (Lower ix)
mIxN ->
        Maybe ix -> f a
f (Maybe ix -> f a) -> Maybe ix -> f a
forall a b. (a -> b) -> a -> b
$!
          case Maybe (Lower ix)
mIxN of
            Just Lower ix
ixN -> ix -> Maybe ix
forall a. a -> Maybe a
Just (ix -> Maybe ix) -> ix -> Maybe ix
forall a b. (a -> b) -> a -> b
$! Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
s Lower ix
ixN
            Maybe (Lower ix)
Nothing ->
              case Maybe Int
mni of
                Just Int
ni -> ix -> Maybe ix
forall a. a -> Maybe a
Just (ix -> Maybe ix) -> ix -> Maybe ix
forall a b. (a -> b) -> a -> b
$! Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
ni (Int -> Lower ix
forall ix. Index ix => Int -> ix
pureIndex Int
0)
                Maybe Int
Nothing -> Maybe ix
forall a. Maybe a
Nothing
    where
      !(Int
s, Lower ix
sIxL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
sIx
      !(Int
e, Lower ix
eIxL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
eIx
      !(Int
inc, Lower ix
incIxL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
incIx
  {-# INLINE stepNextMF #-}

  iterTargetRowMajorA_
    :: Applicative f
    => Int
    -- ^ Target linear index accumulator
    -> Int
    -- ^ Target linear index start
    -> Sz ix
    -- ^ Target size
    -> ix
    -- ^ Source start index
    -> ix
    -- ^ Source stride
    -> (Ix1 -> ix -> f a)
    -- ^ Action that accepts a linear index of the target,
    -- multi-dimensional index of the source and accumulator
    -> f ()
  default iterTargetRowMajorA_
    :: (Applicative f, Index (Lower ix))
    => Int
    -> Int
    -> Sz ix
    -> ix
    -> ix
    -> (Ix1 -> ix -> f a)
    -> f ()
  iterTargetRowMajorA_ !Int
iAcc !Int
iStart Sz ix
szRes ix
ixStart ix
ixStride Int -> ix -> f a
f = do
    let !(SafeSz Int
nRes, !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
szRes
        !(!Int
start, !Lower ix
ixL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStart
        !(!Int
stride, !Lower ix
sL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStride
    Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> (Int -> Int -> f ())
-> f ()
forall (f :: * -> *) a.
Applicative f =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> (Int -> Int -> f a)
-> f ()
iloopA_ (Int
iAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
nRes) Int
start (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
nRes Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
stride) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
stride) ((Int -> Int -> f ()) -> f ()) -> (Int -> Int -> f ()) -> f ()
forall a b. (a -> b) -> a -> b
$ \Int
k Int
j ->
      Int
-> Int
-> Sz (Lower ix)
-> Lower ix
-> Lower ix
-> (Int -> Lower ix -> f a)
-> f ()
forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
Int -> Int -> Sz ix -> ix -> ix -> (Int -> ix -> f a) -> f ()
forall (f :: * -> *) a.
Applicative f =>
Int
-> Int
-> Sz (Lower ix)
-> Lower ix
-> Lower ix
-> (Int -> Lower ix -> f a)
-> f ()
iterTargetRowMajorA_ Int
k Int
iStart Sz (Lower ix)
szL Lower ix
ixL Lower ix
sL ((Int -> Lower ix -> f a) -> f ())
-> (Int -> Lower ix -> f a) -> f ()
forall a b. (a -> b) -> a -> b
$ \Int
i Lower ix
jl -> Int -> ix -> f a
f Int
i (Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
j Lower ix
jl)
  {-# INLINE iterTargetRowMajorA_ #-}

  iterTargetRowMajorAccM
    :: Monad m
    => Int
    -- ^ Target linear index accumulator
    -> Int
    -- ^ Target linear index start
    -> Sz ix
    -- ^ Target size
    -> ix
    -- ^ Source start index
    -> ix
    -- ^ Source stride
    -> a
    -- ^ Accumulator
    -> (Ix1 -> ix -> a -> m a)
    -- ^ Action that accepts a linear index of the target,
    -- multi-dimensional index of the source and accumulator
    -> m a
  default iterTargetRowMajorAccM
    :: (Monad m, Index (Lower ix))
    => Int
    -> Int
    -> Sz ix
    -> ix
    -> ix
    -> a
    -> (Ix1 -> ix -> a -> m a)
    -> m a
  iterTargetRowMajorAccM !Int
iAcc !Int
iStart Sz ix
szRes ix
ixStart ix
ixStride a
initAcc Int -> ix -> a -> m a
f = do
    let !(SafeSz Int
nRes, !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
szRes
        !(!Int
start, !Lower ix
ixL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStart
        !(!Int
stride, !Lower ix
sL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStride
    Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM (Int
iAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
nRes) Int
start (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
nRes Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
stride) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
stride) a
initAcc ((Int -> Int -> a -> m a) -> m a)
-> (Int -> Int -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \Int
k Int
j a
acc ->
      Int
-> Int
-> Sz (Lower ix)
-> Lower ix
-> Lower ix
-> a
-> (Int -> Lower ix -> a -> m a)
-> m a
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
Int
-> Int -> Sz ix -> ix -> ix -> a -> (Int -> ix -> a -> m a) -> m a
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> Sz (Lower ix)
-> Lower ix
-> Lower ix
-> a
-> (Int -> Lower ix -> a -> m a)
-> m a
iterTargetRowMajorAccM Int
k Int
iStart Sz (Lower ix)
szL Lower ix
ixL Lower ix
sL a
acc ((Int -> Lower ix -> a -> m a) -> m a)
-> (Int -> Lower ix -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \Int
i Lower ix
jl -> Int -> ix -> a -> m a
f Int
i (Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
j Lower ix
jl)
  {-# INLINE iterTargetRowMajorAccM #-}

  iterTargetRowMajorAccST
    :: Int
    -- ^ Linear index accumulator
    -> Int
    -- ^ Scheduler multiplying factor. Must be positive
    -> Scheduler s a
    -- ^ Scheduler to use
    -> Int
    -- ^ Target linear index start
    -> Sz ix
    -- ^ Target size
    -> ix
    -- ^ Source start index
    -> ix
    -- ^ Source stride
    -> a
    -- ^ Initial accumulator
    -> (a -> ST s (a, a))
    -- ^ Function that splits accumulator for each scheduled job.
    -> (Ix1 -> ix -> a -> ST s a)
    -- ^ Action
    -> ST s a
  default iterTargetRowMajorAccST
    :: Index (Lower ix)
    => Int
    -> Int
    -> Scheduler s a
    -> Int
    -> Sz ix
    -> ix
    -> ix
    -> a
    -> (a -> ST s (a, a))
    -> (Ix1 -> ix -> a -> ST s a)
    -> ST s a
  iterTargetRowMajorAccST !Int
iAcc !Int
fact Scheduler s a
scheduler Int
iStart Sz ix
sz ix
ixStart ix
ixStride a
initAcc a -> ST s (a, a)
splitAcc Int -> ix -> a -> ST s a
f = do
    let !(SafeSz Int
n, Sz (Lower ix)
nL) = Sz ix -> (Sz Int, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz Sz ix
sz
    if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
      then do
        let !(!Int
start, !Lower ix
ixL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStart
            !(!Int
stride, !Lower ix
sL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStride
            !iAccL :: Int
iAccL = Int
iAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n
        if Scheduler s a -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s a
scheduler Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 Bool -> Bool -> Bool
&& Int
fact Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Scheduler s a -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s a
scheduler Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
fact
          then do
            let newFact :: Int
newFact = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
fact Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
n)
            Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> ST s a)
-> ST s a
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM Int
iAccL Int
start (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
stride) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
stride) a
initAcc ((Int -> Int -> a -> ST s a) -> ST s a)
-> (Int -> Int -> a -> ST s a) -> ST s a
forall a b. (a -> b) -> a -> b
$ \Int
k Int
j a
acc -> do
              Int
-> Int
-> Scheduler s a
-> Int
-> Sz (Lower ix)
-> Lower ix
-> Lower ix
-> a
-> (a -> ST s (a, a))
-> (Int -> Lower ix -> a -> ST s a)
-> ST s a
forall ix s a.
Index ix =>
Int
-> Int
-> Scheduler s a
-> Int
-> Sz ix
-> ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s a
forall s a.
Int
-> Int
-> Scheduler s a
-> Int
-> Sz (Lower ix)
-> Lower ix
-> Lower ix
-> a
-> (a -> ST s (a, a))
-> (Int -> Lower ix -> a -> ST s a)
-> ST s a
iterTargetRowMajorAccST Int
k Int
newFact Scheduler s a
scheduler Int
iStart Sz (Lower ix)
nL Lower ix
ixL Lower ix
sL a
acc a -> ST s (a, a)
splitAcc ((Int -> Lower ix -> a -> ST s a) -> ST s a)
-> (Int -> Lower ix -> a -> ST s a) -> ST s a
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                Int -> ix -> a -> ST s a
f Int
i (ix -> a -> ST s a) -> (Lower ix -> ix) -> Lower ix -> a -> ST s a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
j
          else Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> a
-> (a -> ST s (a, a))
-> (Int -> Int -> Int -> Int -> a -> ST s a)
-> ST s a
forall s a b.
Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> b
-> (b -> ST s (b, b))
-> (Int -> Int -> Int -> Int -> b -> ST s a)
-> ST s b
splitWorkWithFactorST Int
fact Scheduler s a
scheduler Int
start Int
stride Int
n a
initAcc a -> ST s (a, a)
splitAcc ((Int -> Int -> Int -> Int -> a -> ST s a) -> ST s a)
-> (Int -> Int -> Int -> Int -> a -> ST s a) -> ST s a
forall a b. (a -> b) -> a -> b
$
            \Int
chunkStart Int
_ Int
chunkStartAdj Int
chunkStopAdj a
acc ->
              Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> ST s a)
-> ST s a
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM (Int
iAccL Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkStart) Int
chunkStartAdj (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
chunkStopAdj) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
stride) a
acc ((Int -> Int -> a -> ST s a) -> ST s a)
-> (Int -> Int -> a -> ST s a) -> ST s a
forall a b. (a -> b) -> a -> b
$ \Int
k Int
j a
a ->
                Int
-> Int
-> Sz (Lower ix)
-> Lower ix
-> Lower ix
-> a
-> (Int -> Lower ix -> a -> ST s a)
-> ST s a
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
Int
-> Int -> Sz ix -> ix -> ix -> a -> (Int -> ix -> a -> m a) -> m a
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> Sz (Lower ix)
-> Lower ix
-> Lower ix
-> a
-> (Int -> Lower ix -> a -> m a)
-> m a
iterTargetRowMajorAccM Int
k Int
iStart Sz (Lower ix)
nL Lower ix
ixL Lower ix
sL a
a ((Int -> Lower ix -> a -> ST s a) -> ST s a)
-> (Int -> Lower ix -> a -> ST s a) -> ST s a
forall a b. (a -> b) -> a -> b
$ \Int
i -> Int -> ix -> a -> ST s a
f Int
i (ix -> a -> ST s a) -> (Lower ix -> ix) -> Lower ix -> a -> ST s a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
j
      else a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
initAcc
  {-# INLINE iterTargetRowMajorAccST #-}

  iterTargetRowMajorAccST_
    :: Int
    -- ^ Index accumulator
    -> Int
    -- ^ Scheduler multiplying factor. Must be positive
    -> Scheduler s ()
    -- ^ Scheduler to use
    -> Int
    -- ^ Target linear start index
    -> Sz ix
    -- ^ Target size
    -> ix
    -- ^ Source start index
    -> ix
    -- ^ Source stride
    -> a
    -- ^ Initial accumulator
    -> (a -> ST s (a, a))
    -- ^ Function that splits accumulator for each scheduled job.
    -> (Ix1 -> ix -> a -> ST s a)
    -- ^ Action
    -> ST s ()
  default iterTargetRowMajorAccST_
    :: Index (Lower ix)
    => Int
    -> Int
    -> Scheduler s ()
    -> Int
    -> Sz ix
    -> ix
    -> ix
    -> a
    -> (a -> ST s (a, a))
    -> (Ix1 -> ix -> a -> ST s a)
    -> ST s ()
  iterTargetRowMajorAccST_ !Int
iAcc !Int
fact Scheduler s ()
scheduler Int
iStart Sz ix
sz ix
ixStart ix
ixStride a
initAcc a -> ST s (a, a)
splitAcc Int -> ix -> a -> ST s a
f = do
    let !(SafeSz Int
n, 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
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
      let !(!Int
start, !Lower ix
ixL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStart
          !(!Int
stride, !Lower ix
sL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStride
          !iAccL :: Int
iAccL = Int
iAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n
      if Scheduler s () -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 Bool -> Bool -> Bool
&& Int
fact Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Scheduler s () -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
fact
        then do
          let !newFact :: Int
newFact = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
fact Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
n)
          ST s a -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s a -> ST s ()) -> ST s a -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> ST s a)
-> ST s a
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM Int
iAccL Int
start (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
stride) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
stride) a
initAcc ((Int -> Int -> a -> ST s a) -> ST s a)
-> (Int -> Int -> a -> ST s a) -> ST s a
forall a b. (a -> b) -> a -> b
$ \Int
k Int
j a
acc -> do
            (a
accCur, a
accNext) <- a -> ST s (a, a)
splitAcc a
acc
            Scheduler s () -> ST s () -> ST s ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
              Int
-> Int
-> Scheduler s ()
-> Int
-> Sz (Lower ix)
-> Lower ix
-> Lower ix
-> a
-> (a -> ST s (a, a))
-> (Int -> Lower ix -> a -> ST s a)
-> ST s ()
forall ix s a.
Index ix =>
Int
-> Int
-> Scheduler s ()
-> Int
-> Sz ix
-> ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s ()
forall s a.
Int
-> Int
-> Scheduler s ()
-> Int
-> Sz (Lower ix)
-> Lower ix
-> Lower ix
-> a
-> (a -> ST s (a, a))
-> (Int -> Lower ix -> a -> ST s a)
-> ST s ()
iterTargetRowMajorAccST_ Int
k Int
newFact Scheduler s ()
scheduler Int
iStart Sz (Lower ix)
szL Lower ix
ixL Lower ix
sL a
accCur a -> ST s (a, a)
splitAcc ((Int -> Lower ix -> a -> ST s a) -> ST s ())
-> (Int -> Lower ix -> a -> ST s a) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                Int -> ix -> a -> ST s a
f Int
i (ix -> a -> ST s a) -> (Lower ix -> ix) -> Lower ix -> a -> ST s a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
j
            a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
accNext
        else ST s a -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s a -> ST s ()) -> ST s a -> ST s ()
forall a b. (a -> b) -> a -> b
$
          Int
-> Scheduler s ()
-> Int
-> Int
-> Int
-> a
-> (a -> ST s (a, a))
-> (Int -> Int -> Int -> Int -> a -> ST s ())
-> ST s a
forall s a b.
Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> b
-> (b -> ST s (b, b))
-> (Int -> Int -> Int -> Int -> b -> ST s a)
-> ST s b
splitWorkWithFactorST Int
fact Scheduler s ()
scheduler Int
start Int
stride Int
n a
initAcc a -> ST s (a, a)
splitAcc ((Int -> Int -> Int -> Int -> a -> ST s ()) -> ST s a)
-> (Int -> Int -> Int -> Int -> a -> ST s ()) -> ST s a
forall a b. (a -> b) -> a -> b
$
            \Int
chunkStart Int
_ Int
chunkStartAdj Int
chunkStopAdj a
acc ->
              ST s a -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s a -> ST s ()) -> ST s a -> ST s ()
forall a b. (a -> b) -> a -> b
$
                Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> ST s a)
-> ST s a
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM (Int
iAccL Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkStart) Int
chunkStartAdj (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
chunkStopAdj) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
stride) a
acc ((Int -> Int -> a -> ST s a) -> ST s a)
-> (Int -> Int -> a -> ST s a) -> ST s a
forall a b. (a -> b) -> a -> b
$ \Int
k Int
j a
a ->
                  Int
-> Int
-> Sz (Lower ix)
-> Lower ix
-> Lower ix
-> a
-> (Int -> Lower ix -> a -> ST s a)
-> ST s a
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
Int
-> Int -> Sz ix -> ix -> ix -> a -> (Int -> ix -> a -> m a) -> m a
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> Sz (Lower ix)
-> Lower ix
-> Lower ix
-> a
-> (Int -> Lower ix -> a -> m a)
-> m a
iterTargetRowMajorAccM Int
k Int
iStart Sz (Lower ix)
szL Lower ix
ixL Lower ix
sL a
a ((Int -> Lower ix -> a -> ST s a) -> ST s a)
-> (Int -> Lower ix -> a -> ST s a) -> ST s a
forall a b. (a -> b) -> a -> b
$ \Int
i -> Int -> ix -> a -> ST s a
f Int
i (ix -> a -> ST s a) -> (Lower ix -> ix) -> Lower ix -> a -> ST s a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
j
  {-# INLINE iterTargetRowMajorAccST_ #-}

-- | Zero-dimension, i.e. a scalar. Can't really be used directly as there is no instance of
-- `Index` for it, and is included for completeness.
data Ix0 = Ix0 deriving (Ix0 -> Ix0 -> Bool
(Ix0 -> Ix0 -> Bool) -> (Ix0 -> Ix0 -> Bool) -> Eq Ix0
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Ix0 -> Ix0 -> Bool
== :: Ix0 -> Ix0 -> Bool
$c/= :: Ix0 -> Ix0 -> Bool
/= :: Ix0 -> Ix0 -> Bool
Eq, Eq Ix0
Eq Ix0 =>
(Ix0 -> Ix0 -> Ordering)
-> (Ix0 -> Ix0 -> Bool)
-> (Ix0 -> Ix0 -> Bool)
-> (Ix0 -> Ix0 -> Bool)
-> (Ix0 -> Ix0 -> Bool)
-> (Ix0 -> Ix0 -> Ix0)
-> (Ix0 -> Ix0 -> Ix0)
-> Ord Ix0
Ix0 -> Ix0 -> Bool
Ix0 -> Ix0 -> Ordering
Ix0 -> Ix0 -> Ix0
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Ix0 -> Ix0 -> Ordering
compare :: Ix0 -> Ix0 -> Ordering
$c< :: Ix0 -> Ix0 -> Bool
< :: Ix0 -> Ix0 -> Bool
$c<= :: Ix0 -> Ix0 -> Bool
<= :: Ix0 -> Ix0 -> Bool
$c> :: Ix0 -> Ix0 -> Bool
> :: Ix0 -> Ix0 -> Bool
$c>= :: Ix0 -> Ix0 -> Bool
>= :: Ix0 -> Ix0 -> Bool
$cmax :: Ix0 -> Ix0 -> Ix0
max :: Ix0 -> Ix0 -> Ix0
$cmin :: Ix0 -> Ix0 -> Ix0
min :: Ix0 -> Ix0 -> Ix0
Ord, Int -> Ix0 -> ShowS
[Ix0] -> ShowS
Ix0 -> String
(Int -> Ix0 -> ShowS)
-> (Ix0 -> String) -> ([Ix0] -> ShowS) -> Show Ix0
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Ix0 -> ShowS
showsPrec :: Int -> Ix0 -> ShowS
$cshow :: Ix0 -> String
show :: Ix0 -> String
$cshowList :: [Ix0] -> ShowS
showList :: [Ix0] -> ShowS
Show)

instance NFData Ix0 where
  rnf :: Ix0 -> ()
rnf Ix0
Ix0 = ()

-- | A type synonym for 1-dimensional index, i.e. `Int`.
--
-- >>> 5 :: Ix1
-- 5
--
-- @since 0.1.0
type Ix1 = Int

-- | This is a very handy pattern synonym to indicate that any arbitrary `Integral` literal is an
-- `Int`, e.g. a 1-dimensional index: @(Ix1 5) == (5 :: Int)@
--
-- >>> Ix1 5
-- 5
-- >>> :t Ix1 5
-- Ix1 5 :: Ix1
--
-- @since 0.1.0
pattern Ix1 :: Int -> Ix1
pattern $mIx1 :: forall {r}. Int -> (Int -> r) -> ((# #) -> r) -> r
$bIx1 :: Int -> Int
Ix1 i = i

{-# COMPLETE Ix1 #-}

type instance Lower Int = Ix0

instance Index Ix1 where
  type Dimensions Ix1 = 1
  dimensions :: forall (proxy :: * -> *). proxy Int -> Dim
dimensions proxy Int
_ = Dim
1
  {-# INLINE [1] dimensions #-}
  totalElem :: Sz Int -> Int
totalElem = Sz Int -> Int
forall ix. Sz ix -> ix
unSz
  {-# INLINE [1] totalElem #-}
  isSafeIndex :: Sz Int -> Int -> Bool
isSafeIndex (SafeSz Int
k) !Int
i = Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
i Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k
  {-# INLINE [1] isSafeIndex #-}
  toLinearIndex :: Sz Int -> Int -> Int
toLinearIndex Sz Int
_ = Int -> Int
forall a. a -> a
id
  {-# INLINE [1] toLinearIndex #-}
  toLinearIndexAcc :: Int -> Int -> Int -> Int
toLinearIndexAcc !Int
acc Int
m Int
i = Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i
  {-# INLINE [1] toLinearIndexAcc #-}
  fromLinearIndex :: Sz Int -> Int -> Int
fromLinearIndex Sz Int
_ = Int -> Int
forall a. a -> a
id
  {-# INLINE [1] fromLinearIndex #-}
  fromLinearIndexAcc :: Int -> Int -> (Int, Int)
fromLinearIndexAcc Int
n Int
k = Int
k Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Int
n
  {-# INLINE [1] fromLinearIndexAcc #-}
  repairIndex :: Sz Int
-> Int -> (Sz Int -> Int -> Int) -> (Sz Int -> Int -> Int) -> Int
repairIndex k :: Sz Int
k@(SafeSz Int
ksz) !Int
i Sz Int -> Int -> Int
rBelow Sz Int -> Int -> Int
rOver
    | Int
ksz Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = IndexException -> Int
forall a e. Exception e => e -> a
throw (IndexException -> Int) -> IndexException -> Int
forall a b. (a -> b) -> a -> b
$ Int -> IndexException
forall ix. Index ix => ix -> IndexException
IndexZeroException Int
ksz
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = Sz Int -> Int -> Int
rBelow Sz Int
k Int
i
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
ksz = Sz Int -> Int -> Int
rOver Sz Int
k Int
i
    | Bool
otherwise = Int
i
  {-# INLINE [1] repairIndex #-}
  consDim :: Int -> Lower Int -> Int
consDim Int
i Lower Int
_ = Int
i
  {-# INLINE [1] consDim #-}
  unconsDim :: Int -> (Int, Lower Int)
unconsDim Int
i = (Int
i, Ix0
Lower Int
Ix0)
  {-# INLINE [1] unconsDim #-}
  snocDim :: Lower Int -> Int -> Int
snocDim Lower Int
_ Int
i = Int
i
  {-# INLINE [1] snocDim #-}
  unsnocDim :: Int -> (Lower Int, Int)
unsnocDim Int
i = (Ix0
Lower Int
Ix0, Int
i)
  {-# INLINE [1] unsnocDim #-}
  getDimM :: forall (m :: * -> *). MonadThrow m => Int -> Dim -> m Int
getDimM Int
ix Dim
1 = Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
ix
  getDimM Int
ix Dim
d = IndexException -> m Int
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (IndexException -> m Int) -> IndexException -> m Int
forall a b. (a -> b) -> a -> b
$ Int -> Dim -> IndexException
forall ix.
(NFData ix, Eq ix, Show ix, Typeable ix) =>
ix -> Dim -> IndexException
IndexDimensionException Int
ix Dim
d
  {-# INLINE [1] getDimM #-}
  setDimM :: forall (m :: * -> *). MonadThrow m => Int -> Dim -> Int -> m Int
setDimM Int
_ Dim
1 Int
ix = Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
ix
  setDimM Int
ix Dim
d Int
_ = IndexException -> m Int
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (IndexException -> m Int) -> IndexException -> m Int
forall a b. (a -> b) -> a -> b
$ Int -> Dim -> IndexException
forall ix.
(NFData ix, Eq ix, Show ix, Typeable ix) =>
ix -> Dim -> IndexException
IndexDimensionException Int
ix Dim
d
  {-# INLINE [1] setDimM #-}
  modifyDimM :: forall (m :: * -> *).
MonadThrow m =>
Int -> Dim -> (Int -> Int) -> m (Int, Int)
modifyDimM Int
ix Dim
1 Int -> Int
f = (Int, Int) -> m (Int, Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
ix, Int -> Int
f Int
ix)
  modifyDimM Int
ix Dim
d Int -> Int
_ = IndexException -> m (Int, Int)
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (IndexException -> m (Int, Int)) -> IndexException -> m (Int, Int)
forall a b. (a -> b) -> a -> b
$ Int -> Dim -> IndexException
forall ix.
(NFData ix, Eq ix, Show ix, Typeable ix) =>
ix -> Dim -> IndexException
IndexDimensionException Int
ix Dim
d
  {-# INLINE [1] modifyDimM #-}
  pullOutDimM :: forall (m :: * -> *).
MonadThrow m =>
Int -> Dim -> m (Int, Lower Int)
pullOutDimM Int
ix Dim
1 = (Int, Ix0) -> m (Int, Ix0)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
ix, Ix0
Ix0)
  pullOutDimM Int
ix Dim
d = IndexException -> m (Int, Lower Int)
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (IndexException -> m (Int, Lower Int))
-> IndexException -> m (Int, Lower Int)
forall a b. (a -> b) -> a -> b
$ Int -> Dim -> IndexException
forall ix.
(NFData ix, Eq ix, Show ix, Typeable ix) =>
ix -> Dim -> IndexException
IndexDimensionException Int
ix Dim
d
  {-# INLINE [1] pullOutDimM #-}
  insertDimM :: forall (m :: * -> *).
MonadThrow m =>
Lower Int -> Dim -> Int -> m Int
insertDimM Ix0
Lower Int
Ix0 Dim
1 Int
i = Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i
  insertDimM Lower Int
ix Dim
d Int
_ = IndexException -> m Int
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (IndexException -> m Int) -> IndexException -> m Int
forall a b. (a -> b) -> a -> b
$ Ix0 -> Dim -> IndexException
forall ix.
(NFData ix, Eq ix, Show ix, Typeable ix) =>
ix -> Dim -> IndexException
IndexDimensionException Ix0
Lower Int
ix Dim
d
  {-# INLINE [1] insertDimM #-}
  pureIndex :: Int -> Int
pureIndex Int
i = Int
i
  {-# INLINE [1] pureIndex #-}
  liftIndex :: (Int -> Int) -> Int -> Int
liftIndex Int -> Int
f = Int -> Int
f
  {-# INLINE [1] liftIndex #-}
  liftIndex2 :: (Int -> Int -> Int) -> Int -> Int -> Int
liftIndex2 Int -> Int -> Int
f = Int -> Int -> Int
f
  {-# INLINE [1] liftIndex2 #-}
  foldlIndex :: forall a. (a -> Int -> a) -> a -> Int -> a
foldlIndex a -> Int -> a
f = a -> Int -> a
f
  {-# INLINE [1] foldlIndex #-}
  iterM :: forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> Int
-> (Int -> Int -> Bool)
-> a
-> (Int -> a -> m a)
-> m a
iterM Int
k0 Int
k1 Int
inc Int -> Int -> Bool
cond = 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)
  {-# INLINE iterM #-}
  iterF :: forall (f :: * -> *) a.
Int
-> Int
-> Int
-> (Int -> Int -> Bool)
-> f a
-> (Int -> f a -> f a)
-> f a
iterF Int
k0 Int
k1 Int
inc Int -> Int -> Bool
cond = Int
-> (Int -> Bool)
-> (Int -> Int)
-> f a
-> (Int -> f a -> f a)
-> f a
forall (f :: * -> *) a.
Int
-> (Int -> Bool)
-> (Int -> Int)
-> f a
-> (Int -> f a -> f a)
-> f a
loopF Int
k0 (Int -> Int -> Bool
`cond` Int
k1) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
inc)
  {-# INLINE iterF #-}
  stepNextMF :: forall (f :: * -> *) a.
Int
-> Int -> Int -> (Int -> Int -> Bool) -> (Maybe Int -> f a) -> f a
stepNextMF Int
k0 Int
k1 Int
inc Int -> Int -> Bool
cond = Int -> (Int -> Bool) -> (Int -> Int) -> (Maybe Int -> f a) -> f a
forall (f :: * -> *) a.
Int -> (Int -> Bool) -> (Int -> Int) -> (Maybe Int -> f a) -> f a
nextMaybeF Int
k0 (Int -> Int -> Bool
`cond` Int
k1) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
inc)
  {-# INLINE stepNextMF #-}

  iterRowMajorST :: forall s a.
Int
-> Scheduler s a
-> Int
-> Int
-> Sz Int
-> a
-> (a -> ST s (a, a))
-> (Int -> a -> ST s a)
-> ST s a
iterRowMajorST Int
fact Scheduler s a
scheduler Int
start Int
step Sz Int
n =
    Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> a
-> (a -> ST s (a, a))
-> (Int -> a -> ST s a)
-> ST s a
forall s a.
Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> a
-> (a -> ST s (a, a))
-> (Int -> a -> ST s a)
-> ST s a
iterLinearAccST Int
fact Scheduler s a
scheduler Int
start Int
step (Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
n)
  {-# INLINE iterRowMajorST #-}

  iterTargetRowMajorA_ :: forall (f :: * -> *) a.
Applicative f =>
Int -> Int -> Sz Int -> Int -> Int -> (Int -> Int -> f a) -> f ()
iterTargetRowMajorA_ Int
iAcc Int
iStart (SafeSz Int
nRes) Int
start Int
stride =
    Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> (Int -> Int -> f a)
-> f ()
forall (f :: * -> *) a.
Applicative f =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> (Int -> Int -> f a)
-> f ()
iloopA_ (Int
iAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
nRes Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
iStart) Int
start (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
nRes Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
stride) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
stride)
  {-# INLINE iterTargetRowMajorA_ #-}

  iterTargetRowMajorAccM :: forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> Sz Int
-> Int
-> Int
-> a
-> (Int -> Int -> a -> m a)
-> m a
iterTargetRowMajorAccM Int
iAcc Int
iStart (SafeSz Int
nRes) Int
start Int
stride =
    Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM (Int
iAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
nRes Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
iStart) Int
start (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
nRes Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
stride) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
stride)
  {-# INLINE iterTargetRowMajorAccM #-}

  iterTargetRowMajorAccST :: forall s a.
Int
-> Int
-> Scheduler s a
-> Int
-> Sz Int
-> Int
-> Int
-> a
-> (a -> ST s (a, a))
-> (Int -> Int -> a -> ST s a)
-> ST s a
iterTargetRowMajorAccST Int
iAcc Int
fact Scheduler s a
scheduler Int
iStart Sz Int
sz Int
start Int
stride a
initAcc a -> ST s (a, a)
splitAcc Int -> Int -> a -> ST s a
action = do
    let !n :: Int
n = Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
sz
        !iAccL :: Int
iAccL = Int
iStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
iAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n
    Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> a
-> (a -> ST s (a, a))
-> (Int -> Int -> Int -> Int -> a -> ST s a)
-> ST s a
forall s a b.
Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> b
-> (b -> ST s (b, b))
-> (Int -> Int -> Int -> Int -> b -> ST s a)
-> ST s b
splitWorkWithFactorST Int
fact Scheduler s a
scheduler Int
start Int
stride Int
n a
initAcc a -> ST s (a, a)
splitAcc ((Int -> Int -> Int -> Int -> a -> ST s a) -> ST s a)
-> (Int -> Int -> Int -> Int -> a -> ST s a) -> ST s a
forall a b. (a -> b) -> a -> b
$
      \Int
chunkStart Int
_ Int
chunkStartAdj Int
chunkStopAdj a
acc ->
        Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> ST s a)
-> ST s a
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM (Int
iAccL Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkStart) Int
chunkStartAdj (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
chunkStopAdj) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
stride) a
acc Int -> Int -> a -> ST s a
action
  {-# INLINE iterTargetRowMajorAccST #-}

  iterTargetRowMajorAccST_ :: forall s a.
Int
-> Int
-> Scheduler s ()
-> Int
-> Sz Int
-> Int
-> Int
-> a
-> (a -> ST s (a, a))
-> (Int -> Int -> a -> ST s a)
-> ST s ()
iterTargetRowMajorAccST_ Int
iAcc Int
fact Scheduler s ()
scheduler Int
iStart Sz Int
sz Int
start Int
stride a
initAcc a -> ST s (a, a)
splitAcc Int -> Int -> a -> ST s a
action = do
    let !n :: Int
n = Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
sz
        !iAccL :: Int
iAccL = Int
iStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
iAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n
    ST s a -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s a -> ST s ()) -> ST s a -> ST s ()
forall a b. (a -> b) -> a -> b
$
      Int
-> Scheduler s ()
-> Int
-> Int
-> Int
-> a
-> (a -> ST s (a, a))
-> (Int -> Int -> Int -> Int -> a -> ST s ())
-> ST s a
forall s a b.
Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> b
-> (b -> ST s (b, b))
-> (Int -> Int -> Int -> Int -> b -> ST s a)
-> ST s b
splitWorkWithFactorST Int
fact Scheduler s ()
scheduler Int
start Int
stride Int
n a
initAcc a -> ST s (a, a)
splitAcc ((Int -> Int -> Int -> Int -> a -> ST s ()) -> ST s a)
-> (Int -> Int -> Int -> Int -> a -> ST s ()) -> ST s a
forall a b. (a -> b) -> a -> b
$
        \Int
chunkStart Int
_ Int
chunkStartAdj Int
chunkStopAdj a
acc ->
          ST s a -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s a -> ST s ()) -> ST s a -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> ST s a)
-> ST s a
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM (Int
iAccL Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkStart) Int
chunkStartAdj (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
chunkStopAdj) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
stride) a
acc Int -> Int -> a -> ST s a
action
  {-# INLINE iterTargetRowMajorAccST_ #-}

-- | Same as `iterM`, but don't bother with accumulator and return value.
--
-- @since 0.1.0
iterM_ :: (Index ix, Monad m) => ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ :: forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ ix
sIx ix
eIx ix
incIx Int -> Int -> Bool
cond ix -> m a
f = ix
-> ix
-> ix
-> (Int -> Int -> Bool)
-> ()
-> (ix -> () -> m ())
-> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix
-> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a
forall (m :: * -> *) a.
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 () ((ix -> () -> m ()) -> m ()) -> (ix -> () -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !ix
ix !()
a -> ix -> m a
f ix
ix m a -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
a
{-# INLINE iterM_ #-}
{-# DEPRECATED iterM_ "In favor of more lax `iterA_`" #-}

-- | Same as `iterM`, Iterate over a region with specific step, but using
-- `Applicative` instead of a `Monad` and don't bother with accumulator or return value.
--
-- @since 1.0.2
iterA_
  :: forall ix f a
   . (Index ix, Applicative f)
  => ix
  -- ^ Starting index
  -> ix
  -- ^ Ending index (not included)
  -> ix
  -- ^ Stepping index
  -> (Int -> Int -> Bool)
  -- ^ Continuation function. Loop will stop on `False`
  -> (ix -> f a)
  -- ^ Action applied to an index. Result is ignored.
  -> f ()
iterA_ :: forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> f a) -> f ()
iterA_ ix
sIx ix
eIx ix
incIx Int -> Int -> Bool
cond ix -> f a
f =
  ix
-> ix
-> ix
-> (Int -> Int -> Bool)
-> f ()
-> (ix -> f () -> f ())
-> f ()
forall ix (f :: * -> *) a.
Index ix =>
ix
-> ix
-> ix
-> (Int -> Int -> Bool)
-> f a
-> (ix -> f a -> f a)
-> f a
forall (f :: * -> *) a.
ix
-> ix
-> ix
-> (Int -> Int -> Bool)
-> f a
-> (ix -> f a -> f a)
-> f a
iterF ix
sIx ix
eIx ix
incIx Int -> Int -> Bool
cond (() -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) ((ix -> f () -> f ()) -> f ()) -> (ix -> f () -> f ()) -> f ()
forall a b. (a -> b) -> a -> b
$ \ix
ix f ()
go -> ix -> f a
f ix
ix f a -> f () -> f ()
forall a b. f a -> f b -> f b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> f ()
go
{-# INLINE iterA_ #-}

-- | Exceptions that get thrown when there is a problem with an index, size or dimension.
--
-- @since 0.3.0
data IndexException where
  -- | Index contains a zero value along one of the dimensions.
  IndexZeroException :: Index ix => !ix -> IndexException
  -- | Dimension is out of reach.
  IndexDimensionException :: (NFData ix, Eq ix, Show ix, Typeable ix) => !ix -> !Dim -> IndexException
  -- | Index is out of bounds.
  IndexOutOfBoundsException :: Index ix => !(Sz ix) -> !ix -> IndexException

instance Show IndexException where
  show :: IndexException -> String
show (IndexZeroException ix
ix) = String
"IndexZeroException: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> ix -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
1 ix
ix String
""
  show (IndexDimensionException ix
ix Dim
dim) =
    String
"IndexDimensionException: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> Dim -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
1 Dim
dim String
" for " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> ix -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
1 ix
ix String
""
  show (IndexOutOfBoundsException Sz ix
sz ix
ix) =
    String
"IndexOutOfBoundsException: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> ix -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
1 ix
ix String
" is not safe for " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> Sz ix -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
1 Sz ix
sz String
""
  showsPrec :: Int -> IndexException -> ShowS
showsPrec Int
n IndexException
exc = Int -> ShowS -> ShowS
showsPrecWrapped Int
n (IndexException -> String
forall a. Show a => a -> String
show IndexException
exc String -> ShowS
forall a. [a] -> [a] -> [a]
++)

instance Eq IndexException where
  IndexException
e1 == :: IndexException -> IndexException -> Bool
== IndexException
e2 =
    case (IndexException
e1, IndexException
e2) of
      (IndexZeroException ix
i1, IndexZeroException ix
i2t)
        | Just ix
i2 <- ix -> Maybe ix
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast ix
i2t -> ix
i1 ix -> ix -> Bool
forall a. Eq a => a -> a -> Bool
== ix
i2
      (IndexDimensionException ix
i1 Dim
d1, IndexDimensionException ix
i2t Dim
d2)
        | Just ix
i2 <- ix -> Maybe ix
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast ix
i2t -> ix
i1 ix -> ix -> Bool
forall a. Eq a => a -> a -> Bool
== ix
i2 Bool -> Bool -> Bool
&& Dim
d1 Dim -> Dim -> Bool
forall a. Eq a => a -> a -> Bool
== Dim
d2
      (IndexOutOfBoundsException Sz ix
sz1 ix
i1, IndexOutOfBoundsException Sz ix
sz2t ix
i2t)
        | Just ix
i2 <- ix -> Maybe ix
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast ix
i2t
        , Just Sz ix
sz2 <- Sz ix -> Maybe (Sz ix)
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t ->
            Sz ix
sz1 Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix
sz2 Bool -> Bool -> Bool
&& ix
i1 ix -> ix -> Bool
forall a. Eq a => a -> a -> Bool
== ix
i2
      (IndexException, IndexException)
_ -> Bool
False

instance NFData IndexException where
  rnf :: IndexException -> ()
rnf =
    \case
      IndexZeroException ix
i -> ix -> ()
forall a. NFData a => a -> ()
rnf ix
i
      IndexDimensionException ix
i Dim
d -> ix
i ix -> () -> ()
forall a b. NFData a => a -> b -> b
`deepseq` Dim -> ()
forall a. NFData a => a -> ()
rnf Dim
d
      IndexOutOfBoundsException Sz ix
sz ix
i -> Sz ix
sz Sz ix -> () -> ()
forall a b. NFData a => a -> b -> b
`deepseq` ix -> ()
forall a. NFData a => a -> ()
rnf ix
i

instance Exception IndexException

-- | Exception that indicates an issue with an array size.
--
-- @since 0.3.0
data SizeException where
  -- | Two sizes are expected to be equal along some or all dimensions, but they are not.
  SizeMismatchException :: Index ix => !(Sz ix) -> !(Sz ix) -> SizeException
  -- | Total number of elements does not match between the two sizes.
  SizeElementsMismatchException :: (Index ix, Index ix') => !(Sz ix) -> !(Sz ix') -> SizeException
  -- | Described subregion is too big for the specified size.
  SizeSubregionException :: Index ix => !(Sz ix) -> !ix -> !(Sz ix) -> SizeException
  -- | An array with the size cannot contain any elements.
  SizeEmptyException :: Index ix => !(Sz ix) -> SizeException
  -- | Total number of elements is too large resulting in overflow.
  --
  -- @since 0.6.0
  SizeOverflowException :: Index ix => !(Sz ix) -> SizeException
  -- | At least one dimensions contain a negative value.
  --
  -- @since 0.6.0
  SizeNegativeException :: Index ix => !(Sz ix) -> SizeException

instance Eq SizeException where
  SizeException
e1 == :: SizeException -> SizeException -> Bool
== SizeException
e2 =
    case (SizeException
e1, SizeException
e2) of
      (SizeMismatchException Sz ix
sz1 Sz ix
sz1', SizeMismatchException Sz ix
sz2t Sz ix
sz2t')
        | Just Sz ix
sz2 <- Sz ix -> Maybe (Sz ix)
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t
        , Just Sz ix
sz2' <- Sz ix -> Maybe (Sz ix)
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t' ->
            Sz ix
sz1 Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix
sz2 Bool -> Bool -> Bool
&& Sz ix
sz1' Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix
sz2'
      (SizeElementsMismatchException Sz ix
sz1 Sz ix'
sz1', SizeElementsMismatchException Sz ix
sz2t Sz ix'
sz2t')
        | Just Sz ix
sz2 <- Sz ix -> Maybe (Sz ix)
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t
        , Just Sz ix'
sz2' <- Sz ix' -> Maybe (Sz ix')
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix'
sz2t' ->
            Sz ix
sz1 Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix
sz2 Bool -> Bool -> Bool
&& Sz ix'
sz1' Sz ix' -> Sz ix' -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix'
sz2'
      (SizeSubregionException Sz ix
sz1 ix
i1 Sz ix
sz1', SizeSubregionException Sz ix
sz2t ix
i2t Sz ix
sz2t')
        | Just Sz ix
sz2 <- Sz ix -> Maybe (Sz ix)
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t
        , Just ix
i2 <- ix -> Maybe ix
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast ix
i2t
        , Just Sz ix
sz2' <- Sz ix -> Maybe (Sz ix)
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t' ->
            Sz ix
sz1 Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix
sz2 Bool -> Bool -> Bool
&& ix
i1 ix -> ix -> Bool
forall a. Eq a => a -> a -> Bool
== ix
i2 Bool -> Bool -> Bool
&& Sz ix
sz1' Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix
sz2'
      (SizeEmptyException Sz ix
sz1, SizeEmptyException Sz ix
sz2t)
        | Just Sz ix
sz2 <- Sz ix -> Maybe (Sz ix)
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t -> Sz ix
sz1 Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix
sz2
      (SizeOverflowException Sz ix
sz1, SizeOverflowException Sz ix
sz2t)
        | Just Sz ix
sz2 <- Sz ix -> Maybe (Sz ix)
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t -> Sz ix
sz1 Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix
sz2
      (SizeNegativeException Sz ix
sz1, SizeNegativeException Sz ix
sz2t)
        | Just Sz ix
sz2 <- Sz ix -> Maybe (Sz ix)
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t -> Sz ix
sz1 Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix
sz2
      (SizeException, SizeException)
_ -> Bool
False

instance NFData SizeException where
  rnf :: SizeException -> ()
rnf =
    \case
      SizeMismatchException Sz ix
sz Sz ix
sz' -> Sz ix
sz Sz ix -> () -> ()
forall a b. NFData a => a -> b -> b
`deepseq` Sz ix -> ()
forall a. NFData a => a -> ()
rnf Sz ix
sz'
      SizeElementsMismatchException Sz ix
sz Sz ix'
sz' -> Sz ix
sz Sz ix -> () -> ()
forall a b. NFData a => a -> b -> b
`deepseq` Sz ix' -> ()
forall a. NFData a => a -> ()
rnf Sz ix'
sz'
      SizeSubregionException Sz ix
sz ix
i Sz ix
sz' -> Sz ix
sz Sz ix -> () -> ()
forall a b. NFData a => a -> b -> b
`deepseq` ix
i ix -> () -> ()
forall a b. NFData a => a -> b -> b
`deepseq` Sz ix -> ()
forall a. NFData a => a -> ()
rnf Sz ix
sz'
      SizeEmptyException Sz ix
sz -> Sz ix -> ()
forall a. NFData a => a -> ()
rnf Sz ix
sz
      SizeOverflowException Sz ix
sz -> Sz ix -> ()
forall a. NFData a => a -> ()
rnf Sz ix
sz
      SizeNegativeException Sz ix
sz -> Sz ix -> ()
forall a. NFData a => a -> ()
rnf Sz ix
sz

instance Exception SizeException

instance Show SizeException where
  show :: SizeException -> String
show (SizeMismatchException Sz ix
sz Sz ix
sz') =
    String
"SizeMismatchException: (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Sz ix -> String
forall a. Show a => a -> String
show Sz ix
sz String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") vs (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Sz ix -> String
forall a. Show a => a -> String
show Sz ix
sz' String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
  show (SizeElementsMismatchException Sz ix
sz Sz ix'
sz') =
    String
"SizeElementsMismatchException: (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Sz ix -> String
forall a. Show a => a -> String
show Sz ix
sz String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") vs (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Sz ix' -> String
forall a. Show a => a -> String
show Sz ix'
sz' String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
  show (SizeSubregionException Sz ix
sz' ix
ix Sz ix
sz) =
    String
"SizeSubregionException: ("
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Sz ix -> String
forall a. Show a => a -> String
show Sz ix
sz'
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") is to small for "
      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
" ("
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Sz ix -> String
forall a. Show a => a -> String
show Sz ix
sz
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
  show (SizeEmptyException Sz ix
sz) =
    String
"SizeEmptyException: (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Sz ix -> String
forall a. Show a => a -> String
show Sz ix
sz String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") corresponds to an empty array"
  show (SizeOverflowException Sz ix
sz) =
    String
"SizeOverflowException: (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Sz ix -> String
forall a. Show a => a -> String
show Sz ix
sz String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") is too big"
  show (SizeNegativeException Sz ix
sz) =
    String
"SizeNegativeException: (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Sz ix -> String
forall a. Show a => a -> String
show Sz ix
sz String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") contains negative value"
  showsPrec :: Int -> SizeException -> ShowS
showsPrec Int
n SizeException
exc = Int -> ShowS -> ShowS
showsPrecWrapped Int
n (SizeException -> String
forall a. Show a => a -> String
show SizeException
exc String -> ShowS
forall a. [a] -> [a] -> [a]
++)

-- | Exception that can happen upon conversion of a ragged type array into the rectangular kind. Which
-- means conversion from lists is susceptible to this exception.
--
-- @since 0.3.0
data ShapeException
  = -- | Across a specific dimension there was not enough elements for the supplied size
    DimTooShortException !Dim !(Sz Ix1) !(Sz Ix1)
  | -- | Across a specific dimension there was too many elements for the supplied size
    DimTooLongException !Dim !(Sz Ix1) !(Sz Ix1)
  | -- | Expected an empty size, but the shape was not empty.
    ShapeNonEmpty
  deriving (ShapeException -> ShapeException -> Bool
(ShapeException -> ShapeException -> Bool)
-> (ShapeException -> ShapeException -> Bool) -> Eq ShapeException
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ShapeException -> ShapeException -> Bool
== :: ShapeException -> ShapeException -> Bool
$c/= :: ShapeException -> ShapeException -> Bool
/= :: ShapeException -> ShapeException -> Bool
Eq)

instance Show ShapeException where
  showsPrec :: Int -> ShapeException -> ShowS
showsPrec Int
n =
    \case
      DimTooShortException Dim
d Sz Int
sz Sz Int
sz' -> String -> Dim -> Sz Int -> Sz Int -> ShowS
showsShapeExc String
"DimTooShortException" Dim
d Sz Int
sz Sz Int
sz'
      DimTooLongException Dim
d Sz Int
sz Sz Int
sz' -> String -> Dim -> Sz Int -> Sz Int -> ShowS
showsShapeExc String
"DimTooLongException" Dim
d Sz Int
sz Sz Int
sz'
      ShapeException
ShapeNonEmpty -> (String
"ShapeNonEmpty" String -> ShowS
forall a. [a] -> [a] -> [a]
++)
    where
      showsShapeExc :: String -> Dim -> Sz Int -> Sz Int -> ShowS
showsShapeExc String
tyName Dim
d Sz Int
sz Sz Int
sz' =
        Int -> ShowS -> ShowS
showsPrecWrapped
          Int
n
          ( (String
tyName String -> ShowS
forall a. [a] -> [a] -> [a]
++)
              ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
" for " String -> ShowS
forall a. [a] -> [a] -> [a]
++)
              ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dim -> ShowS
forall a. Show a => a -> ShowS
shows Dim
d
              ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
": expected (" String -> ShowS
forall a. [a] -> [a] -> [a]
++)
              ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz Int -> ShowS
forall a. Show a => a -> ShowS
shows Sz Int
sz
              ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"), got (" String -> ShowS
forall a. [a] -> [a] -> [a]
++)
              ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz Int -> ShowS
forall a. Show a => a -> ShowS
shows Sz Int
sz'
              ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
")" String -> ShowS
forall a. [a] -> [a] -> [a]
++)
          )

instance Exception ShapeException

showsPrecWrapped :: Int -> ShowS -> ShowS
showsPrecWrapped :: Int -> ShowS -> ShowS
showsPrecWrapped Int
n ShowS
inner
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 = ShowS
inner
  | Bool
otherwise = (Char
'(' Char -> ShowS
forall a. a -> [a] -> [a]
:) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
inner ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
")" String -> ShowS
forall a. [a] -> [a] -> [a]
++)