{-# 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 (
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
newtype Sz ix
=
SafeSz ix
deriving (Sz ix -> Sz ix -> Bool
forall ix. Eq ix => Sz ix -> Sz ix -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Sz ix -> Sz ix -> Bool
$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
Eq, 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
min :: Sz ix -> Sz ix -> Sz ix
$cmin :: forall ix. Ord ix => Sz ix -> Sz ix -> Sz ix
max :: Sz ix -> Sz ix -> Sz ix
$cmax :: forall ix. Ord ix => Sz ix -> Sz ix -> Sz ix
>= :: 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
$c< :: forall ix. Ord ix => Sz ix -> Sz ix -> Bool
compare :: Sz ix -> Sz ix -> Ordering
$ccompare :: forall ix. Ord ix => Sz ix -> Sz ix -> Ordering
Ord, Sz ix -> ()
forall ix. NFData ix => Sz ix -> ()
forall a. (a -> ()) -> NFData a
rnf :: Sz ix -> ()
$crnf :: forall ix. NFData ix => Sz ix -> ()
NFData)
pattern Sz :: Index ix => ix -> Sz ix
pattern $bSz :: forall ix. Index ix => ix -> Sz ix
$mSz :: forall {r} {ix}.
Index ix =>
Sz ix -> (ix -> r) -> ((# #) -> r) -> r
Sz ix <- SafeSz ix
where
Sz ix
ix = forall ix. ix -> Sz ix
SafeSz (forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (forall a. Ord a => a -> a -> a
max Int
0) ix
ix)
{-# COMPLETE Sz #-}
pattern Sz1 :: Ix1 -> Sz Ix1
pattern $bSz1 :: Int -> Sz Int
$mSz1 :: forall {r}. Sz Int -> (Int -> r) -> ((# #) -> r) -> r
Sz1 ix <- SafeSz ix
where
Sz1 Int
ix = forall ix. ix -> Sz ix
SafeSz (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 = forall ix. ix -> Sz ix
SafeSz forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
uniformRM (forall ix. Index ix => Int -> ix
pureIndex Int
0, forall ix. Index ix => Int -> ix
pureIndex 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 = forall ix. ix -> Sz ix
SafeSz forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
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 forall a. [a] -> [a] -> [a]
++)
where
str :: String
str =
String
"Sz"
forall a. [a] -> [a] -> [a]
++ case Dim -> Int
unDim (forall ix (proxy :: * -> *). Index ix => proxy ix -> Dim
dimensions Sz ix
sz) of
Int
1 -> String
"1 " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show ix
usz
Int
_ -> String
" (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> ShowS
shows ix
usz String
")"
instance (Num ix, Index ix) => Num (Sz ix) where
+ :: Sz ix -> Sz ix -> Sz ix
(+) Sz ix
x Sz ix
y = forall ix. Index ix => ix -> Sz ix
Sz (coerce :: forall a b. Coercible a b => a -> b
coerce Sz ix
x forall a. Num a => a -> a -> a
+ coerce :: forall a b. Coercible a b => a -> b
coerce Sz ix
y)
{-# INLINE (+) #-}
(-) Sz ix
x Sz ix
y = forall ix. Index ix => ix -> Sz ix
Sz (coerce :: forall a b. Coercible a b => a -> b
coerce Sz ix
x forall a. Num a => a -> a -> a
- coerce :: 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 = forall ix. Index ix => ix -> Sz ix
Sz (coerce :: forall a b. Coercible a b => a -> b
coerce Sz ix
x forall a. Num a => a -> a -> a
* coerce :: 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 forall a. Eq a => a -> a -> Bool
== forall ix. Index ix => Sz ix
zeroSz = Sz ix
x
| Bool
otherwise =
forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$
String
"Attempted to negate: "
forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Sz ix
x
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 = forall ix. ix -> Sz ix
SafeSz (forall a. Num a => a -> a
signum (coerce :: forall a b. Coercible a b => a -> b
coerce Sz ix
x))
{-# INLINE signum #-}
fromInteger :: Integer -> Sz ix
fromInteger = forall ix. Index ix => ix -> Sz ix
Sz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Num a => Integer -> a
fromInteger
{-# INLINE fromInteger #-}
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
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i forall a. Ord a => a -> a -> Bool
< Int
0) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ forall ix. Index ix => Sz ix -> SizeException
SizeNegativeException (forall ix. ix -> Sz ix
SafeSz ix
ix)
let acc' :: Int
acc' = Int
i forall a. Num a => a -> a -> a
* Int
acc
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
acc' forall a. Eq a => a -> a -> Bool
/= Int
0 Bool -> Bool -> Bool
&& Int
acc' forall a. Ord a => a -> a -> Bool
< Int
acc) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ forall ix. Index ix => Sz ix -> SizeException
SizeOverflowException (forall ix. ix -> Sz ix
SafeSz ix
ix)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
acc'
forall ix. Index ix => ix -> Sz ix
Sz ix
ix forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall ix a. Index ix => (a -> Int -> a) -> a -> ix -> a
foldlIndex (\m Int
acc Int
i -> m Int
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> Int -> m Int
guardNegativeOverflow Int
i) (forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
1) ix
ix
{-# INLINE mkSzM #-}
unSz :: Sz ix -> ix
unSz :: forall ix. Sz ix -> ix
unSz (SafeSz ix
ix) = ix
ix
{-# INLINE unSz #-}
zeroSz :: Index ix => Sz ix
zeroSz :: forall ix. Index ix => Sz ix
zeroSz = forall ix. ix -> Sz ix
SafeSz (forall ix. Index ix => Int -> ix
pureIndex Int
0)
{-# INLINE zeroSz #-}
oneSz :: Index ix => Sz ix
oneSz :: forall ix. Index ix => Sz ix
oneSz = forall ix. ix -> Sz ix
SafeSz (forall ix. Index ix => Int -> ix
pureIndex Int
1)
{-# INLINE oneSz #-}
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) = forall ix. Index ix => ix -> Sz ix
Sz (forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex Int -> Int
f ix
ix)
{-# INLINE liftSz #-}
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 = forall ix. Index ix => ix -> Sz ix
Sz (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
f (coerce :: forall a b. Coercible a b => a -> b
coerce Sz ix
sz1) (coerce :: forall a b. Coercible a b => a -> b
coerce Sz ix
sz2))
{-# INLINE liftSz2 #-}
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) = forall ix. ix -> Sz ix
SafeSz (forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
i Lower ix
ix)
{-# INLINE consSz #-}
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) = forall ix. ix -> Sz ix
SafeSz (forall ix. Index ix => Lower ix -> Int -> ix
snocDim Lower ix
i Int
ix)
{-# INLINE snocSz #-}
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) = forall ix. ix -> Sz ix
SafeSz forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
sz Dim
dim Int
sz1
{-# INLINE setSzM #-}
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) = forall ix. ix -> Sz ix
SafeSz forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
Lower ix -> Dim -> Int -> m ix
insertDimM Lower ix
sz Dim
dim Int
sz1
{-# INLINE insertSzM #-}
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) = coerce :: forall a b. Coercible a b => a -> b
coerce (forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
sz)
{-# INLINE unconsSz #-}
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) = coerce :: forall a b. Coercible a b => a -> b
coerce (forall ix. Index ix => ix -> (Lower ix, Int)
unsnocDim ix
sz)
{-# INLINE unsnocSz #-}
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) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap coerce :: forall a b. Coercible a b => a -> b
coerce forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m (Int, Lower ix)
pullOutDimM ix
sz
{-# INLINE pullOutSzM #-}
newtype Dim = Dim {Dim -> Int
unDim :: Int} deriving (Dim -> Dim -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Dim -> Dim -> Bool
$c/= :: Dim -> Dim -> Bool
== :: Dim -> Dim -> Bool
$c== :: Dim -> Dim -> Bool
Eq, Eq 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
min :: Dim -> Dim -> Dim
$cmin :: Dim -> Dim -> Dim
max :: Dim -> Dim -> Dim
$cmax :: Dim -> Dim -> Dim
>= :: Dim -> Dim -> Bool
$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
compare :: Dim -> Dim -> Ordering
$ccompare :: Dim -> Dim -> Ordering
Ord, Integer -> Dim
Dim -> Dim
Dim -> Dim -> Dim
forall a.
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> a)
-> (a -> a)
-> (Integer -> a)
-> Num a
fromInteger :: Integer -> Dim
$cfromInteger :: Integer -> Dim
signum :: Dim -> Dim
$csignum :: Dim -> Dim
abs :: Dim -> Dim
$cabs :: Dim -> Dim
negate :: Dim -> Dim
$cnegate :: Dim -> Dim
* :: Dim -> Dim -> Dim
$c* :: Dim -> Dim -> Dim
- :: Dim -> Dim -> Dim
$c- :: Dim -> Dim -> Dim
+ :: Dim -> Dim -> Dim
$c+ :: Dim -> Dim -> Dim
Num, Num Dim
Ord Dim
Dim -> Rational
forall a. Num a -> Ord a -> (a -> Rational) -> Real a
toRational :: Dim -> Rational
$ctoRational :: Dim -> Rational
Real, Enum Dim
Real 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
toInteger :: Dim -> Integer
$ctoInteger :: Dim -> Integer
divMod :: Dim -> Dim -> (Dim, Dim)
$cdivMod :: Dim -> Dim -> (Dim, Dim)
quotRem :: Dim -> Dim -> (Dim, Dim)
$cquotRem :: Dim -> Dim -> (Dim, Dim)
mod :: Dim -> Dim -> Dim
$cmod :: Dim -> Dim -> Dim
div :: Dim -> Dim -> Dim
$cdiv :: Dim -> Dim -> Dim
rem :: Dim -> Dim -> Dim
$crem :: Dim -> Dim -> Dim
quot :: Dim -> Dim -> Dim
$cquot :: Dim -> Dim -> Dim
Integral, Int -> Dim
Dim -> Int
Dim -> [Dim]
Dim -> Dim
Dim -> Dim -> [Dim]
Dim -> Dim -> Dim -> [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
enumFromThenTo :: Dim -> Dim -> Dim -> [Dim]
$cenumFromThenTo :: Dim -> Dim -> Dim -> [Dim]
enumFromTo :: Dim -> Dim -> [Dim]
$cenumFromTo :: Dim -> Dim -> [Dim]
enumFromThen :: Dim -> Dim -> [Dim]
$cenumFromThen :: Dim -> Dim -> [Dim]
enumFrom :: Dim -> [Dim]
$cenumFrom :: Dim -> [Dim]
fromEnum :: Dim -> Int
$cfromEnum :: Dim -> Int
toEnum :: Int -> Dim
$ctoEnum :: Int -> Dim
pred :: Dim -> Dim
$cpred :: Dim -> Dim
succ :: Dim -> Dim
$csucc :: Dim -> Dim
Enum, Dim -> ()
forall a. (a -> ()) -> NFData a
rnf :: Dim -> ()
$crnf :: Dim -> ()
NFData)
instance Show Dim where
show :: Dim -> String
show (Dim Int
d) = String
"(Dim " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
d 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
uniformRM (Int
1, 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
uniformRM (coerce :: forall a b. Coercible a b => a -> b
coerce (Dim, Dim)
r) g
g
instance Random Dim
data Dimension (n :: Nat) where
DimN :: (1 <= n, KnownNat n) => Dimension n
pattern Dim1 :: Dimension 1
pattern $bDim1 :: Dimension 1
$mDim1 :: forall {r}. Dimension 1 -> ((# #) -> r) -> ((# #) -> r) -> r
Dim1 = DimN
pattern Dim2 :: Dimension 2
pattern $bDim2 :: Dimension 2
$mDim2 :: forall {r}. Dimension 2 -> ((# #) -> r) -> ((# #) -> r) -> r
Dim2 = DimN
pattern Dim3 :: Dimension 3
pattern $bDim3 :: Dimension 3
$mDim3 :: forall {r}. Dimension 3 -> ((# #) -> r) -> ((# #) -> r) -> r
Dim3 = DimN
pattern Dim4 :: Dimension 4
pattern $bDim4 :: Dimension 4
$mDim4 :: forall {r}. Dimension 4 -> ((# #) -> r) -> ((# #) -> r) -> r
Dim4 = DimN
pattern Dim5 :: Dimension 5
pattern $bDim5 :: Dimension 5
$mDim5 :: forall {r}. Dimension 5 -> ((# #) -> r) -> ((# #) -> r) -> r
Dim5 = DimN
type IsIndexDimension ix n = (1 <= n, n <= Dimensions ix, Index ix, KnownNat n)
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)
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 Dimensions ix :: Nat
dimensions :: proxy ix -> Dim
totalElem :: Sz ix -> Int
consDim :: Int -> Lower ix -> ix
unconsDim :: ix -> (Int, Lower ix)
snocDim :: Lower ix -> Int -> ix
unsnocDim :: ix -> (Lower ix, Int)
pullOutDimM :: MonadThrow m => ix -> Dim -> m (Int, Lower ix)
insertDimM :: MonadThrow m => Lower ix -> Dim -> Int -> m ix
getDimM :: MonadThrow m => ix -> Dim -> m Int
getDimM ix
ix Dim
dim = forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> (Int -> Int) -> m (Int, ix)
modifyDimM ix
ix Dim
dim forall a. a -> a
id
{-# INLINE [1] getDimM #-}
setDimM :: MonadThrow m => ix -> Dim -> Int -> m ix
setDimM ix
ix Dim
dim Int
i = forall a b. (a, b) -> b
snd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> (Int -> Int) -> m (Int, ix)
modifyDimM ix
ix Dim
dim (forall a b. a -> b -> a
const Int
i)
{-# INLINE [1] setDimM #-}
modifyDimM :: MonadThrow m => ix -> Dim -> (Int -> Int) -> m (Int, ix)
modifyDimM ix
ix Dim
dim Int -> Int
f = do
Int
i <- forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> m Int
getDimM ix
ix Dim
dim
ix
ix' <- forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM ix
ix Dim
dim (Int -> Int
f Int
i)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
i, ix
ix')
{-# INLINE [1] modifyDimM #-}
pureIndex :: Int -> ix
liftIndex2 :: (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex :: (Int -> Int) -> ix -> ix
liftIndex Int -> Int
f = forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (\Int
_ Int
i -> Int -> Int
f Int
i) (forall ix. Index ix => Int -> ix
pureIndex Int
0)
{-# INLINE [1] liftIndex #-}
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 = forall ix a. Index ix => (a -> Int -> a) -> a -> ix -> a
foldlIndex a -> Int -> a
f (a -> Int -> a
f a
acc Int
i0) Lower ix
ixL
where
!(Int
i0, Lower ix
ixL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ix
{-# INLINE [1] foldlIndex #-}
isSafeIndex
:: Sz ix
-> ix
-> Bool
default isSafeIndex
:: Index (Lower ix)
=> Sz ix
-> ix
-> Bool
isSafeIndex Sz ix
sz !ix
ix = forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz Int
n0 Int
i0 Bool -> Bool -> 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) = forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz Sz ix
sz
!(Int
i0, Lower ix
ixL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ix
{-# INLINE [1] isSafeIndex #-}
toLinearIndex
:: Sz ix
-> ix
-> Ix1
default toLinearIndex :: Index (Lower ix) => Sz ix -> ix -> Ix1
toLinearIndex (SafeSz ix
sz) !ix
ix = forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex (forall ix. ix -> Sz ix
SafeSz Lower ix
szL) Lower ix
ixL forall a. Num a => a -> a -> a
* Int
n forall a. Num a => a -> a -> a
+ Int
i
where
!(Lower ix
szL, Int
n) = forall ix. Index ix => ix -> (Lower ix, Int)
unsnocDim ix
sz
!(Lower ix
ixL, Int
i) = forall ix. Index ix => ix -> (Lower ix, Int)
unsnocDim ix
ix
{-# INLINE [1] toLinearIndex #-}
toLinearIndexAcc :: Ix1 -> ix -> ix -> Ix1
default toLinearIndexAcc :: Index (Lower ix) => Ix1 -> ix -> ix -> Ix1
toLinearIndexAcc !Int
acc !ix
sz !ix
ix = forall ix. Index ix => Int -> ix -> ix -> Int
toLinearIndexAcc (Int
acc forall a. Num a => a -> a -> a
* Int
n forall a. Num a => a -> a -> a
+ Int
i) Lower ix
szL Lower ix
ixL
where
!(Int
n, Lower ix
szL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
sz
!(Int
i, Lower ix
ixL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ix
{-# INLINE [1] toLinearIndexAcc #-}
fromLinearIndex :: Sz ix -> Ix1 -> ix
default fromLinearIndex :: Index (Lower ix) => Sz ix -> Ix1 -> ix
fromLinearIndex (SafeSz ix
sz) !Int
k = forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
q Lower ix
ixL
where
!(!Int
q, !Lower ix
ixL) = forall ix. Index ix => ix -> Int -> (Int, ix)
fromLinearIndexAcc (forall a b. (a, b) -> b
snd (forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
sz)) Int
k
{-# INLINE [1] fromLinearIndex #-}
fromLinearIndexAcc :: ix -> Ix1 -> (Int, ix)
default fromLinearIndexAcc :: Index (Lower ix) => ix -> Ix1 -> (Ix1, ix)
fromLinearIndexAcc !ix
ix' !Int
k = (Int
q, forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
r Lower ix
ixL)
where
!(!Int
m, !Lower ix
ix) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ix'
!(!Int
kL, !Lower ix
ixL) = forall ix. Index ix => ix -> Int -> (Int, ix)
fromLinearIndexAcc Lower ix
ix Int
k
!(!Int
q, !Int
r) = forall a. Integral a => a -> a -> (a, a)
quotRem Int
kL Int
m
{-# INLINE [1] fromLinearIndexAcc #-}
repairIndex
:: Sz ix
-> ix
-> (Sz Int -> Int -> Int)
-> (Sz Int -> Int -> Int)
-> 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 =
forall ix. Index ix => Int -> Lower ix -> ix
consDim (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) (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) = forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz Sz ix
sz
!(Int
i, Lower ix
ixL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ix
{-# INLINE [1] repairIndex #-}
iterM
:: Monad m
=> ix
-> ix
-> ix
-> (Int -> Int -> Bool)
-> a
-> (ix -> a -> m a)
-> 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 =
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) (forall a. Num a => a -> a -> a
+ Int
inc) a
acc forall a b. (a -> b) -> a -> b
$ \ !Int
i !a
acc0 ->
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix
-> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a
iterM Lower ix
sIxL Lower ix
eIxL Lower ix
incIxL Int -> Int -> Bool
cond a
acc0 forall a b. (a -> b) -> a -> b
$ \ !Lower ix
ix -> ix -> a -> m a
f (forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
i Lower ix
ix)
where
!(Int
s, Lower ix
sIxL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
sIx
!(Int
e, Lower ix
eIxL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
eIx
!(Int
inc, Lower ix
incIxL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
incIx
{-# INLINE iterM #-}
iterRowMajorST
:: Int
-> Scheduler s a
-> ix
-> ix
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (ix -> a -> ST s a)
-> 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)) = forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz Sz ix
sz
if Int
n forall a. Ord a => a -> a -> Bool
> Int
0
then do
let !(!Int
start, !Lower ix
ixL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStart
!(!Int
stride, !Lower ix
sL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStride
if forall s a. Scheduler s a -> Int
numWorkers Scheduler s a
scheduler forall a. Ord a => a -> a -> Bool
> Int
1 Bool -> Bool -> Bool
&& Int
fact forall a. Ord a => a -> a -> Bool
> Int
1 Bool -> Bool -> Bool
&& Int
n forall a. Ord a => a -> a -> Bool
< forall s a. Scheduler s a -> Int
numWorkers Scheduler s a
scheduler forall a. Num a => a -> a -> a
* Int
fact
then do
let !newFact :: Int
newFact = Int
1 forall a. Num a => a -> a -> a
+ (Int
fact forall a. Integral a => a -> a -> a
`quot` Int
n)
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
start (forall a. Ord a => a -> a -> Bool
< Int
start forall a. Num a => a -> a -> a
+ Int
n forall a. Num a => a -> a -> a
* Int
stride) (forall a. Num a => a -> a -> a
+ Int
stride) a
initAcc forall a b. (a -> b) -> a -> b
$ \Int
j a
acc ->
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
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 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
j)
else 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 forall a b. (a -> b) -> a -> b
$
\Int
_ Int
_ Int
chunkStartAdj Int
chunkStopAdj a
acc ->
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
chunkStartAdj (forall a. Ord a => a -> a -> Bool
< Int
chunkStopAdj) (forall a. Num a => a -> a -> a
+ Int
stride) a
acc forall a b. (a -> b) -> a -> b
$ \Int
j a
a ->
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix
-> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a
iterM Lower ix
ixL Lower ix
nL Lower ix
sL forall a. Ord a => a -> a -> Bool
(<) a
a (ix -> a -> ST s a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
j)
else forall (f :: * -> *) a. Applicative f => a -> f a
pure a
initAcc
{-# INLINE iterRowMajorST #-}
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 =
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) (forall a. Num a => a -> a -> a
+ Int
inc) f a
initAct forall a b. (a -> b) -> a -> b
$ \ !Int
i f a
g ->
forall ix (f :: * -> *) a.
Index ix =>
ix
-> ix
-> ix
-> (Int -> Int -> Bool)
-> f a
-> (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 (forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
i Lower ix
ix))
where
!(Int
s, Lower ix
sIxL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
sIx
!(Int
e, Lower ix
eIxL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
eIx
!(Int
inc, Lower ix
incIxL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
incIx
{-# INLINE iterF #-}
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 =
forall (f :: * -> *) a.
Int -> (Int -> Bool) -> (Int -> Int) -> (Maybe Int -> f a) -> f a
nextMaybeF Int
s (Int -> Int -> Bool
`cond` Int
e) (forall a. Num a => a -> a -> a
+ Int
inc) forall a b. (a -> b) -> a -> b
$ \ !Maybe Int
mni ->
forall ix (f :: * -> *) a.
Index ix =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (Maybe ix -> f a) -> f a
stepNextMF Lower ix
sIxL Lower ix
eIxL Lower ix
incIxL Int -> Int -> Bool
cond forall a b. (a -> b) -> a -> b
$ \Maybe (Lower ix)
mIxN ->
Maybe ix -> f a
f forall a b. (a -> b) -> a -> b
$!
case Maybe (Lower ix)
mIxN of
Just Lower ix
ixN -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$! 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 -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$! forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
ni (forall ix. Index ix => Int -> ix
pureIndex Int
0)
Maybe Int
Nothing -> forall a. Maybe a
Nothing
where
!(Int
s, Lower ix
sIxL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
sIx
!(Int
e, Lower ix
eIxL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
eIx
!(Int
inc, Lower ix
incIxL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
incIx
{-# INLINE stepNextMF #-}
iterTargetRowMajorA_
:: Applicative f
=> Int
-> Int
-> Sz ix
-> ix
-> ix
-> (Ix1 -> ix -> f a)
-> 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) = forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz Sz ix
szRes
!(!Int
start, !Lower ix
ixL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStart
!(!Int
stride, !Lower ix
sL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStride
forall (f :: * -> *) a.
Applicative f =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> (Int -> Int -> f a)
-> f ()
iloopA_ (Int
iAcc forall a. Num a => a -> a -> a
* Int
nRes) Int
start (forall a. Ord a => a -> a -> Bool
< Int
start forall a. Num a => a -> a -> a
+ Int
nRes forall a. Num a => a -> a -> a
* Int
stride) (forall a. Num a => a -> a -> a
+ Int
stride) forall a b. (a -> b) -> a -> b
$ \Int
k Int
j ->
forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
Int -> Int -> Sz ix -> ix -> ix -> (Int -> ix -> f a) -> f ()
iterTargetRowMajorA_ Int
k Int
iStart Sz (Lower ix)
szL Lower ix
ixL Lower ix
sL forall a b. (a -> b) -> a -> b
$ \Int
i Lower ix
jl -> Int -> ix -> f a
f Int
i (forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
j Lower ix
jl)
{-# INLINE iterTargetRowMajorA_ #-}
iterTargetRowMajorAccM
:: Monad m
=> Int
-> Int
-> Sz ix
-> ix
-> ix
-> a
-> (Ix1 -> ix -> a -> m a)
-> 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) = forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz Sz ix
szRes
!(!Int
start, !Lower ix
ixL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStart
!(!Int
stride, !Lower ix
sL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStride
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM (Int
iAcc forall a. Num a => a -> a -> a
* Int
nRes) Int
start (forall a. Ord a => a -> a -> Bool
< Int
start forall a. Num a => a -> a -> a
+ Int
nRes forall a. Num a => a -> a -> a
* Int
stride) (forall a. Num a => a -> a -> a
+ Int
stride) a
initAcc forall a b. (a -> b) -> a -> b
$ \Int
k Int
j a
acc ->
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
Int
-> Int -> Sz ix -> ix -> ix -> a -> (Int -> ix -> a -> m a) -> m a
iterTargetRowMajorAccM Int
k Int
iStart Sz (Lower ix)
szL Lower ix
ixL Lower ix
sL a
acc forall a b. (a -> b) -> a -> b
$ \Int
i Lower ix
jl -> Int -> ix -> a -> m a
f Int
i (forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
j Lower ix
jl)
{-# INLINE iterTargetRowMajorAccM #-}
iterTargetRowMajorAccST
:: 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
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) = forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz Sz ix
sz
if Int
n forall a. Ord a => a -> a -> Bool
> Int
0
then do
let !(!Int
start, !Lower ix
ixL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStart
!(!Int
stride, !Lower ix
sL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStride
!iAccL :: Int
iAccL = Int
iAcc forall a. Num a => a -> a -> a
* Int
n
if forall s a. Scheduler s a -> Int
numWorkers Scheduler s a
scheduler forall a. Ord a => a -> a -> Bool
> Int
1 Bool -> Bool -> Bool
&& Int
fact forall a. Ord a => a -> a -> Bool
> Int
1 Bool -> Bool -> Bool
&& Int
n forall a. Ord a => a -> a -> Bool
< forall s a. Scheduler s a -> Int
numWorkers Scheduler s a
scheduler forall a. Num a => a -> a -> a
* Int
fact
then do
let newFact :: Int
newFact = Int
1 forall a. Num a => a -> a -> a
+ (Int
fact forall a. Integral a => a -> a -> a
`quot` Int
n)
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM Int
iAccL Int
start (forall a. Ord a => a -> a -> Bool
< Int
start forall a. Num a => a -> a -> a
+ Int
n forall a. Num a => a -> a -> a
* Int
stride) (forall a. Num a => a -> a -> a
+ Int
stride) a
initAcc forall a b. (a -> b) -> a -> b
$ \Int
k Int
j a
acc -> do
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
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 forall a b. (a -> b) -> a -> b
$ \Int
i ->
Int -> ix -> a -> ST s a
f Int
i forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
j
else 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 forall a b. (a -> b) -> a -> b
$
\Int
chunkStart Int
_ Int
chunkStartAdj Int
chunkStopAdj a
acc ->
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM (Int
iAccL forall a. Num a => a -> a -> a
+ Int
chunkStart) Int
chunkStartAdj (forall a. Ord a => a -> a -> Bool
< Int
chunkStopAdj) (forall a. Num a => a -> a -> a
+ Int
stride) a
acc forall a b. (a -> b) -> a -> b
$ \Int
k Int
j a
a ->
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
Int
-> Int -> Sz ix -> ix -> ix -> a -> (Int -> ix -> a -> m a) -> m a
iterTargetRowMajorAccM Int
k Int
iStart Sz (Lower ix)
nL Lower ix
ixL Lower ix
sL a
a forall a b. (a -> b) -> a -> b
$ \Int
i -> Int -> ix -> a -> ST s a
f Int
i forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
j
else forall (f :: * -> *) a. Applicative f => a -> f a
pure a
initAcc
{-# INLINE iterTargetRowMajorAccST #-}
iterTargetRowMajorAccST_
:: Int
-> Int
-> Scheduler s ()
-> Int
-> Sz ix
-> ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> 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) = forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz Sz ix
sz
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n forall a. Ord a => a -> a -> Bool
> Int
0) forall a b. (a -> b) -> a -> b
$ do
let !(!Int
start, !Lower ix
ixL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStart
!(!Int
stride, !Lower ix
sL) = forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
ixStride
!iAccL :: Int
iAccL = Int
iAcc forall a. Num a => a -> a -> a
* Int
n
if forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler forall a. Ord a => a -> a -> Bool
> Int
1 Bool -> Bool -> Bool
&& Int
fact forall a. Ord a => a -> a -> Bool
> Int
1 Bool -> Bool -> Bool
&& Int
n forall a. Ord a => a -> a -> Bool
< forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler forall a. Num a => a -> a -> a
* Int
fact
then do
let !newFact :: Int
newFact = Int
1 forall a. Num a => a -> a -> a
+ (Int
fact forall a. Integral a => a -> a -> a
`quot` Int
n)
forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM Int
iAccL Int
start (forall a. Ord a => a -> a -> Bool
< Int
n forall a. Num a => a -> a -> a
* Int
stride) (forall a. Num a => a -> a -> a
+ Int
stride) a
initAcc 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
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$
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 ()
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 forall a b. (a -> b) -> a -> b
$ \Int
i ->
Int -> ix -> a -> ST s a
f Int
i forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
j
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
accNext
else forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$
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 forall a b. (a -> b) -> a -> b
$
\Int
chunkStart Int
_ Int
chunkStartAdj Int
chunkStopAdj a
acc ->
forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM (Int
iAccL forall a. Num a => a -> a -> a
+ Int
chunkStart) Int
chunkStartAdj (forall a. Ord a => a -> a -> Bool
< Int
chunkStopAdj) (forall a. Num a => a -> a -> a
+ Int
stride) a
acc forall a b. (a -> b) -> a -> b
$ \Int
k Int
j a
a ->
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
Int
-> Int -> Sz ix -> ix -> ix -> a -> (Int -> ix -> a -> m a) -> m a
iterTargetRowMajorAccM Int
k Int
iStart Sz (Lower ix)
szL Lower ix
ixL Lower ix
sL a
a forall a b. (a -> b) -> a -> b
$ \Int
i -> Int -> ix -> a -> ST s a
f Int
i forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
j
{-# INLINE iterTargetRowMajorAccST_ #-}
data Ix0 = Ix0 deriving (Ix0 -> Ix0 -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Ix0 -> Ix0 -> Bool
$c/= :: Ix0 -> Ix0 -> Bool
== :: Ix0 -> Ix0 -> Bool
$c== :: Ix0 -> Ix0 -> Bool
Eq, Eq 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
min :: Ix0 -> Ix0 -> Ix0
$cmin :: Ix0 -> Ix0 -> Ix0
max :: Ix0 -> Ix0 -> Ix0
$cmax :: Ix0 -> Ix0 -> Ix0
>= :: Ix0 -> Ix0 -> Bool
$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
compare :: Ix0 -> Ix0 -> Ordering
$ccompare :: Ix0 -> Ix0 -> Ordering
Ord, Int -> Ix0 -> ShowS
[Ix0] -> ShowS
Ix0 -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Ix0] -> ShowS
$cshowList :: [Ix0] -> ShowS
show :: Ix0 -> String
$cshow :: Ix0 -> String
showsPrec :: Int -> Ix0 -> ShowS
$cshowsPrec :: Int -> Ix0 -> ShowS
Show)
instance NFData Ix0 where
rnf :: Ix0 -> ()
rnf Ix0
Ix0 = ()
type Ix1 = Int
pattern Ix1 :: Int -> Ix1
pattern $bIx1 :: Int -> Int
$mIx1 :: forall {r}. Int -> (Int -> r) -> ((# #) -> r) -> r
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 = forall ix. Sz ix -> ix
unSz
{-# INLINE [1] totalElem #-}
isSafeIndex :: Sz Int -> Int -> Bool
isSafeIndex (SafeSz Int
k) !Int
i = Int
0 forall a. Ord a => a -> a -> Bool
<= Int
i Bool -> Bool -> Bool
&& Int
i forall a. Ord a => a -> a -> Bool
< Int
k
{-# INLINE [1] isSafeIndex #-}
toLinearIndex :: Sz Int -> Int -> Int
toLinearIndex Sz Int
_ = forall a. a -> a
id
{-# INLINE [1] toLinearIndex #-}
toLinearIndexAcc :: Int -> Int -> Int -> Int
toLinearIndexAcc !Int
acc Int
m Int
i = Int
acc forall a. Num a => a -> a -> a
* Int
m forall a. Num a => a -> a -> a
+ Int
i
{-# INLINE [1] toLinearIndexAcc #-}
fromLinearIndex :: Sz Int -> Int -> Int
fromLinearIndex Sz Int
_ = forall a. a -> a
id
{-# INLINE [1] fromLinearIndex #-}
fromLinearIndexAcc :: Int -> Int -> (Int, Int)
fromLinearIndexAcc Int
n Int
k = Int
k 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 forall a. Ord a => a -> a -> Bool
<= Int
0 = forall a e. Exception e => e -> a
throw forall a b. (a -> b) -> a -> b
$ forall ix. Index ix => ix -> IndexException
IndexZeroException Int
ksz
| Int
i forall a. Ord a => a -> a -> Bool
< Int
0 = Sz Int -> Int -> Int
rBelow Sz Int
k Int
i
| Int
i 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
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
Ix0, Int
i)
{-# INLINE [1] unsnocDim #-}
getDimM :: forall (m :: * -> *). MonadThrow m => Int -> Dim -> m Int
getDimM Int
ix Dim
1 = forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
ix
getDimM Int
ix Dim
d = forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ 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 = forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
ix
setDimM Int
ix Dim
d Int
_ = forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ 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 = forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
ix, Int -> Int
f Int
ix)
modifyDimM Int
ix Dim
d Int -> Int
_ = forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ 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 = forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
ix, Ix0
Ix0)
pullOutDimM Int
ix Dim
d = forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ 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 = forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i
insertDimM Lower Int
ix Dim
d Int
_ = forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ forall ix.
(NFData ix, Eq ix, Show ix, Typeable ix) =>
ix -> Dim -> IndexException
IndexDimensionException 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 = forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
k0 (Int -> Int -> Bool
`cond` Int
k1) (forall a. Num a => a -> a -> a
+ Int
inc)
{-# 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 = 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) (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 = forall (f :: * -> *) a.
Int -> (Int -> Bool) -> (Int -> Int) -> (Maybe Int -> f a) -> f a
nextMaybeF Int
k0 (Int -> Int -> Bool
`cond` Int
k1) (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 =
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 (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 =
forall (f :: * -> *) a.
Applicative f =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> (Int -> Int -> f a)
-> f ()
iloopA_ (Int
iAcc forall a. Num a => a -> a -> a
* Int
nRes forall a. Num a => a -> a -> a
+ Int
iStart) Int
start (forall a. Ord a => a -> a -> Bool
< Int
start forall a. Num a => a -> a -> a
+ Int
nRes forall a. Num a => a -> a -> a
* Int
stride) (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 =
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM (Int
iAcc forall a. Num a => a -> a -> a
* Int
nRes forall a. Num a => a -> a -> a
+ Int
iStart) Int
start (forall a. Ord a => a -> a -> Bool
< Int
start forall a. Num a => a -> a -> a
+ Int
nRes forall a. Num a => a -> a -> a
* Int
stride) (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 = forall ix. Sz ix -> ix
unSz Sz Int
sz
!iAccL :: Int
iAccL = Int
iStart forall a. Num a => a -> a -> a
+ Int
iAcc forall a. Num a => a -> a -> a
* Int
n
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 forall a b. (a -> b) -> a -> b
$
\Int
chunkStart Int
_ Int
chunkStartAdj Int
chunkStopAdj a
acc ->
forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM (Int
iAccL forall a. Num a => a -> a -> a
+ Int
chunkStart) Int
chunkStartAdj (forall a. Ord a => a -> a -> Bool
< Int
chunkStopAdj) (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 = forall ix. Sz ix -> ix
unSz Sz Int
sz
!iAccL :: Int
iAccL = Int
iStart forall a. Num a => a -> a -> a
+ Int
iAcc forall a. Num a => a -> a -> a
* Int
n
forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$
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 forall a b. (a -> b) -> a -> b
$
\Int
chunkStart Int
_ Int
chunkStartAdj Int
chunkStopAdj a
acc ->
forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM (Int
iAccL forall a. Num a => a -> a -> a
+ Int
chunkStart) Int
chunkStartAdj (forall a. Ord a => a -> a -> Bool
< Int
chunkStopAdj) (forall a. Num a => a -> a -> a
+ Int
stride) a
acc Int -> Int -> a -> ST s a
action
{-# INLINE iterTargetRowMajorAccST_ #-}
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 = forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix
-> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a
iterM ix
sIx ix
eIx ix
incIx Int -> Int -> Bool
cond () forall a b. (a -> b) -> a -> b
$ \ !ix
ix !()
a -> ix -> m a
f ix
ix forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
a
{-# INLINE iterM_ #-}
{-# DEPRECATED iterM_ "In favor of more lax `iterA_`" #-}
iterA_
:: forall ix f a
. (Index ix, Applicative f)
=> ix
-> ix
-> ix
-> (Int -> Int -> Bool)
-> (ix -> f a)
-> 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 =
forall ix (f :: * -> *) a.
Index 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 (forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) forall a b. (a -> b) -> a -> b
$ \ix
ix f ()
go -> ix -> f a
f ix
ix forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> f ()
go
{-# INLINE iterA_ #-}
data IndexException where
IndexZeroException :: Index ix => !ix -> IndexException
IndexDimensionException :: (NFData ix, Eq ix, Show ix, Typeable ix) => !ix -> !Dim -> IndexException
IndexOutOfBoundsException :: Index ix => !(Sz ix) -> !ix -> IndexException
instance Show IndexException where
show :: IndexException -> String
show (IndexZeroException ix
ix) = String
"IndexZeroException: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => Int -> a -> ShowS
showsPrec Int
1 ix
ix String
""
show (IndexDimensionException ix
ix Dim
dim) =
String
"IndexDimensionException: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => Int -> a -> ShowS
showsPrec Int
1 Dim
dim String
" for " forall a. [a] -> [a] -> [a]
++ forall a. Show a => Int -> a -> ShowS
showsPrec Int
1 ix
ix String
""
show (IndexOutOfBoundsException Sz ix
sz ix
ix) =
String
"IndexOutOfBoundsException: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => Int -> a -> ShowS
showsPrec Int
1 ix
ix String
" is not safe for " forall a. [a] -> [a] -> [a]
++ 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 (forall a. Show a => a -> String
show IndexException
exc 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 <- forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast ix
i2t -> ix
i1 forall a. Eq a => a -> a -> Bool
== ix
i2
(IndexDimensionException ix
i1 Dim
d1, IndexDimensionException ix
i2t Dim
d2)
| Just ix
i2 <- forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast ix
i2t -> ix
i1 forall a. Eq a => a -> a -> Bool
== ix
i2 Bool -> Bool -> Bool
&& Dim
d1 forall a. Eq a => a -> a -> Bool
== Dim
d2
(IndexOutOfBoundsException Sz ix
sz1 ix
i1, IndexOutOfBoundsException Sz ix
sz2t ix
i2t)
| Just ix
i2 <- forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast ix
i2t
, Just Sz ix
sz2 <- forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t ->
Sz ix
sz1 forall a. Eq a => a -> a -> Bool
== Sz ix
sz2 Bool -> Bool -> Bool
&& ix
i1 forall a. Eq a => a -> a -> Bool
== ix
i2
(IndexException, IndexException)
_ -> Bool
False
instance NFData IndexException where
rnf :: IndexException -> ()
rnf =
\case
IndexZeroException ix
i -> forall a. NFData a => a -> ()
rnf ix
i
IndexDimensionException ix
i Dim
d -> ix
i forall a b. NFData a => a -> b -> b
`deepseq` forall a. NFData a => a -> ()
rnf Dim
d
IndexOutOfBoundsException Sz ix
sz ix
i -> Sz ix
sz forall a b. NFData a => a -> b -> b
`deepseq` forall a. NFData a => a -> ()
rnf ix
i
instance Exception IndexException
data SizeException where
SizeMismatchException :: Index ix => !(Sz ix) -> !(Sz ix) -> SizeException
SizeElementsMismatchException :: (Index ix, Index ix') => !(Sz ix) -> !(Sz ix') -> SizeException
SizeSubregionException :: Index ix => !(Sz ix) -> !ix -> !(Sz ix) -> SizeException
SizeEmptyException :: Index ix => !(Sz ix) -> SizeException
SizeOverflowException :: Index ix => !(Sz ix) -> SizeException
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 <- forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t
, Just Sz ix
sz2' <- forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t' ->
Sz ix
sz1 forall a. Eq a => a -> a -> Bool
== Sz ix
sz2 Bool -> Bool -> Bool
&& Sz ix
sz1' 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 <- forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t
, Just Sz ix'
sz2' <- forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix'
sz2t' ->
Sz ix
sz1 forall a. Eq a => a -> a -> Bool
== Sz ix
sz2 Bool -> Bool -> Bool
&& Sz ix'
sz1' 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 <- forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t
, Just ix
i2 <- forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast ix
i2t
, Just Sz ix
sz2' <- forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t' ->
Sz ix
sz1 forall a. Eq a => a -> a -> Bool
== Sz ix
sz2 Bool -> Bool -> Bool
&& ix
i1 forall a. Eq a => a -> a -> Bool
== ix
i2 Bool -> Bool -> Bool
&& Sz ix
sz1' forall a. Eq a => a -> a -> Bool
== Sz ix
sz2'
(SizeEmptyException Sz ix
sz1, SizeEmptyException Sz ix
sz2t)
| Just Sz ix
sz2 <- forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t -> Sz ix
sz1 forall a. Eq a => a -> a -> Bool
== Sz ix
sz2
(SizeOverflowException Sz ix
sz1, SizeOverflowException Sz ix
sz2t)
| Just Sz ix
sz2 <- forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t -> Sz ix
sz1 forall a. Eq a => a -> a -> Bool
== Sz ix
sz2
(SizeNegativeException Sz ix
sz1, SizeNegativeException Sz ix
sz2t)
| Just Sz ix
sz2 <- forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast Sz ix
sz2t -> Sz ix
sz1 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 forall a b. NFData a => a -> b -> b
`deepseq` forall a. NFData a => a -> ()
rnf Sz ix
sz'
SizeElementsMismatchException Sz ix
sz Sz ix'
sz' -> Sz ix
sz forall a b. NFData a => a -> b -> b
`deepseq` forall a. NFData a => a -> ()
rnf Sz ix'
sz'
SizeSubregionException Sz ix
sz ix
i Sz ix
sz' -> Sz ix
sz forall a b. NFData a => a -> b -> b
`deepseq` ix
i forall a b. NFData a => a -> b -> b
`deepseq` forall a. NFData a => a -> ()
rnf Sz ix
sz'
SizeEmptyException Sz ix
sz -> forall a. NFData a => a -> ()
rnf Sz ix
sz
SizeOverflowException Sz ix
sz -> forall a. NFData a => a -> ()
rnf Sz ix
sz
SizeNegativeException Sz ix
sz -> 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: (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Sz ix
sz forall a. [a] -> [a] -> [a]
++ String
") vs (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Sz ix
sz' forall a. [a] -> [a] -> [a]
++ String
")"
show (SizeElementsMismatchException Sz ix
sz Sz ix'
sz') =
String
"SizeElementsMismatchException: (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Sz ix
sz forall a. [a] -> [a] -> [a]
++ String
") vs (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Sz ix'
sz' forall a. [a] -> [a] -> [a]
++ String
")"
show (SizeSubregionException Sz ix
sz' ix
ix Sz ix
sz) =
String
"SizeSubregionException: ("
forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Sz ix
sz'
forall a. [a] -> [a] -> [a]
++ String
") is to small for "
forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show ix
ix
forall a. [a] -> [a] -> [a]
++ String
" ("
forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Sz ix
sz
forall a. [a] -> [a] -> [a]
++ String
")"
show (SizeEmptyException Sz ix
sz) =
String
"SizeEmptyException: (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Sz ix
sz forall a. [a] -> [a] -> [a]
++ String
") corresponds to an empty array"
show (SizeOverflowException Sz ix
sz) =
String
"SizeOverflowException: (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Sz ix
sz forall a. [a] -> [a] -> [a]
++ String
") is too big"
show (SizeNegativeException Sz ix
sz) =
String
"SizeNegativeException: (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Sz ix
sz forall a. [a] -> [a] -> [a]
++ String
") contains negative value"
showsPrec :: Int -> SizeException -> ShowS
showsPrec Int
n SizeException
exc = Int -> ShowS -> ShowS
showsPrecWrapped Int
n (forall a. Show a => a -> String
show SizeException
exc forall a. [a] -> [a] -> [a]
++)
data ShapeException
=
DimTooShortException !Dim !(Sz Ix1) !(Sz Ix1)
|
DimTooLongException !Dim !(Sz Ix1) !(Sz Ix1)
|
ShapeNonEmpty
deriving (ShapeException -> ShapeException -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ShapeException -> ShapeException -> Bool
$c/= :: ShapeException -> ShapeException -> Bool
== :: ShapeException -> ShapeException -> Bool
$c== :: 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" 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 forall a. [a] -> [a] -> [a]
++)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
" for " forall a. [a] -> [a] -> [a]
++)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> ShowS
shows Dim
d
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
": expected (" forall a. [a] -> [a] -> [a]
++)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> ShowS
shows Sz Int
sz
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"), got (" forall a. [a] -> [a] -> [a]
++)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> ShowS
shows Sz Int
sz'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
")" forall a. [a] -> [a] -> [a]
++)
)
instance Exception ShapeException
showsPrecWrapped :: Int -> ShowS -> ShowS
showsPrecWrapped :: Int -> ShowS -> ShowS
showsPrecWrapped Int
n ShowS
inner
| Int
n forall a. Ord a => a -> a -> Bool
< Int
1 = ShowS
inner
| Bool
otherwise = (Char
'(' forall a. a -> [a] -> [a]
:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
inner forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
")" forall a. [a] -> [a] -> [a]
++)