{-# LANGUAGE CPP                    #-}
{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE ExplicitForAll         #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE MagicHash              #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE PolyKinds              #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE UnboxedTuples          #-}
module Numeric.DataFrame.Internal.PrimArray
  ( PrimArray (..), CumulDims (..)
  , cumulDims, cdTotalDim, cdTotalDim#
  , cdIx, cdIxSub, getOffAndSteps, getOffAndStepsSub
  , cdIxM, getOffAndStepsM, getOffAndStepsSubM
  , ixOff, unsafeFromFlatList, getSteps, fromSteps
  , withArrayContent, fromElems, broadcast
  ) where

import Control.Arrow  ((***))
import Data.Monoid    as Mon (Monoid (..))
import Data.Semigroup as Sem (Semigroup (..))
import GHC.Base       hiding (foldr)
 -- (ByteArray#, Int (..), Int#, RuntimeRep (..), Type,
 --                           Word (..), inline, runRW#, touch#, word2Int#, (+#))
import GHC.Exts           (TYPE)
import Numeric.Dimensions
import Numeric.PrimBytes


-- | Given @Dims ns@, @CumulativeDims@ is a list of length @Length ns + 1@;
--   which cumulative @totalDim@ accumulated on the right.
--   In particular, its first element is @totalDim ds@,
--   its last element is always is always @1@.
newtype CumulDims = CumulDims { CumulDims -> [Word]
unCumulDims :: [Word] }
  deriving Int -> CumulDims -> ShowS
[CumulDims] -> ShowS
CumulDims -> String
(Int -> CumulDims -> ShowS)
-> (CumulDims -> String)
-> ([CumulDims] -> ShowS)
-> Show CumulDims
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [CumulDims] -> ShowS
$cshowList :: [CumulDims] -> ShowS
show :: CumulDims -> String
$cshow :: CumulDims -> String
showsPrec :: Int -> CumulDims -> ShowS
$cshowsPrec :: Int -> CumulDims -> ShowS
Show

instance Sem.Semigroup CumulDims where
    CumulDims [Word]
as <> :: CumulDims -> CumulDims -> CumulDims
<> CumulDims [Word]
bs = [Word] -> CumulDims
CumulDims ([Word] -> CumulDims) -> [Word] -> CumulDims
forall a b. (a -> b) -> a -> b
$ (Word -> Word) -> [Word] -> [Word]
forall a b. (a -> b) -> [a] -> [b]
map ([Word] -> Word
forall a. [a] -> a
head [Word]
bs Word -> Word -> Word
forall a. Num a => a -> a -> a
*) ([Word] -> [Word]
forall a. [a] -> [a]
init [Word]
as) [Word] -> [Word] -> [Word]
forall a. [a] -> [a] -> [a]
++ [Word]
bs

instance Mon.Monoid CumulDims where
    mempty :: CumulDims
mempty = [Word] -> CumulDims
CumulDims [Word
1]
#if !(MIN_VERSION_base(4,11,0))
    mappend = (<>)
#endif


-- | Calculate cumulative dims
cumulDims :: Dims (ns :: [k]) -> CumulDims
cumulDims :: Dims ns -> CumulDims
cumulDims = [Word] -> CumulDims
CumulDims ([Word] -> CumulDims)
-> (Dims ns -> [Word]) -> Dims ns -> CumulDims
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word -> [Word] -> [Word]) -> (Word, [Word]) -> [Word]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (:)
          ((Word, [Word]) -> [Word])
-> (Dims ns -> (Word, [Word])) -> Dims ns -> [Word]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word -> (Word, [Word]) -> (Word, [Word]))
-> (Word, [Word]) -> [Word] -> (Word, [Word])
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Word
d (Word
c, [Word]
cs) -> (Word
cWord -> Word -> Word
forall a. Num a => a -> a -> a
*Word
d,Word
cWord -> [Word] -> [Word]
forall a. a -> [a] -> [a]
:[Word]
cs)) (Word
1, []) ([Word] -> (Word, [Word]))
-> (Dims ns -> [Word]) -> Dims ns -> (Word, [Word])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dims ns -> [Word]
forall k (xs :: [k]). Dims xs -> [Word]
listDims

-- | Get the total number of elements
cdTotalDim :: CumulDims -> Word
cdTotalDim :: CumulDims -> Word
cdTotalDim ~(CumulDims ~(Word
n:[Word]
_)) = Word
n

cdTotalDim# :: CumulDims -> Int#
cdTotalDim# :: CumulDims -> Int#
cdTotalDim# ~(CumulDims ~(Word
n:[Word]
_)) = case Word
n of W# Word#
w -> Word# -> Int#
word2Int# Word#
w

-- | Calculate offset of an Idxs
--
--   Note, you can take offset of subspace with CumulDims of larger space
--     - very convenient!
--
--   If any of the dims in @ns@ is unknown (@n ~ XN m@),
--   then this function is unsafe and can throw an `OutOfDimBounds` exception.
--   Otherwise, its safety is guaranteed by the type system.
cdIx :: CumulDims -> Idxs ns -> Int
cdIx :: CumulDims -> Idxs ns -> Int
cdIx CumulDims
steps
  = Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Int) -> (Idxs ns -> Word) -> Idxs ns -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word, [Word]) -> Word
forall a b. (a, b) -> a
fst ((Word, [Word]) -> Word)
-> (Idxs ns -> (Word, [Word])) -> Idxs ns -> Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Word -> [Word] -> [Word] -> (Word, [Word])
getOffAndSteps' Word
0 Word
0 (CumulDims -> [Word]
unCumulDims CumulDims
steps) ([Word] -> (Word, [Word]))
-> (Idxs ns -> [Word]) -> Idxs ns -> (Word, [Word])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Idxs ns -> [Word]
forall k (ds :: [k]). Idxs ds -> [Word]
listIdxs

-- | Calculate offset of an Idxs.
--
--   Also check if the last index plus dimVal of subN is not bigger than the
--   corresponding dim inside CumulDims; throw an `OutOfDimBounds` otherwise.
--
--   If any of the dims in @ns@ is unknown (@n ~ XN m@),
--   then this function is unsafe and can throw an `OutOfDimBounds` exception.
--   Otherwise, its safety is guaranteed by the type system.
cdIxSub :: CumulDims -> Idxs (ns +: idxN) -> Dim subN -> Int
cdIxSub :: CumulDims -> Idxs (ns +: idxN) -> Dim subN -> Int
cdIxSub CumulDims
steps Idxs (ns +: idxN)
idxs Dim subN
d
  = Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral
  (Word -> Int) -> ([Word] -> Word) -> [Word] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word, [Word]) -> Word
forall a b. (a, b) -> a
fst ((Word, [Word]) -> Word)
-> ([Word] -> (Word, [Word])) -> [Word] -> Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Word -> [Word] -> [Word] -> (Word, [Word])
getOffAndSteps' (Dim subN -> Word
forall k (x :: k). Dim x -> Word
dimVal Dim subN
d) Word
0 (CumulDims -> [Word]
unCumulDims CumulDims
steps) ([Word] -> Int) -> [Word] -> Int
forall a b. (a -> b) -> a -> b
$ Idxs (ns +: idxN) -> [Word]
forall k (ds :: [k]). Idxs ds -> [Word]
listIdxs Idxs (ns +: idxN)
idxs

-- | Calculate offset of an Idxs and return remaining CumulDims.
--
--   If any of the dims in @ns@ is unknown (@n ~ XN m@),
--   then this function is unsafe and can throw an `OutOfDimBounds` exception.
--   Otherwise, its safety is guaranteed by the type system.
getOffAndSteps :: Int -- ^ Initial offset
               -> CumulDims -> Idxs ns -> (Int, CumulDims)
getOffAndSteps :: Int -> CumulDims -> Idxs ns -> (Int, CumulDims)
getOffAndSteps Int
off0 CumulDims
steps
  = (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Int)
-> ([Word] -> CumulDims) -> (Word, [Word]) -> (Int, CumulDims)
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** [Word] -> CumulDims
CumulDims)
  ((Word, [Word]) -> (Int, CumulDims))
-> (Idxs ns -> (Word, [Word])) -> Idxs ns -> (Int, CumulDims)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Word -> [Word] -> [Word] -> (Word, [Word])
getOffAndSteps' Word
0 (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
off0) (CumulDims -> [Word]
unCumulDims CumulDims
steps) ([Word] -> (Word, [Word]))
-> (Idxs ns -> [Word]) -> Idxs ns -> (Word, [Word])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Idxs ns -> [Word]
forall k (ds :: [k]). Idxs ds -> [Word]
listIdxs

-- | Calculate offset of an Idxs and return remaining CumulDims.
--
--   Also check if the last index plus dimVal of subN is not bigger than the
--   corresponding dim inside CumulDims; throw an `OutOfDimBounds` otherwise.
--
--   If any of the dims in @ns@ is unknown (@n ~ XN m@),
--   then this function is unsafe and can throw an `OutOfDimBounds` exception.
--   Otherwise, its safety is guaranteed by the type system.
getOffAndStepsSub :: Int -- ^ Initial offset
                  -> CumulDims -> Idxs (ns +: idxN)
                  -> Dim subN -> (Int, CumulDims)
getOffAndStepsSub :: Int
-> CumulDims -> Idxs (ns +: idxN) -> Dim subN -> (Int, CumulDims)
getOffAndStepsSub Int
off0 CumulDims
steps Idxs (ns +: idxN)
idxs Dim subN
d
  = (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Int)
-> ([Word] -> CumulDims) -> (Word, [Word]) -> (Int, CumulDims)
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** [Word] -> CumulDims
CumulDims)
  ((Word, [Word]) -> (Int, CumulDims))
-> ([Word] -> (Word, [Word])) -> [Word] -> (Int, CumulDims)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Word -> [Word] -> [Word] -> (Word, [Word])
getOffAndSteps' (Dim subN -> Word
forall k (x :: k). Dim x -> Word
dimVal Dim subN
d) (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
off0) (CumulDims -> [Word]
unCumulDims CumulDims
steps)
  ([Word] -> (Int, CumulDims)) -> [Word] -> (Int, CumulDims)
forall a b. (a -> b) -> a -> b
$ Idxs (ns +: idxN) -> [Word]
forall k (ds :: [k]). Idxs ds -> [Word]
listIdxs Idxs (ns +: idxN)
idxs


getOffAndSteps' :: Word -> Word -> [Word] -> [Word] -> (Word, [Word])
getOffAndSteps' :: Word -> Word -> [Word] -> [Word] -> (Word, [Word])
getOffAndSteps' Word
sub0 Word
off0 [Word]
steps0 [Word]
is0 = Word -> Word -> [Word] -> [Word] -> (Word, [Word])
go Word
sub0 Word
off0 [Word]
steps0 [Word]
is0
  where
#ifndef UNSAFE_INDICES
    ds0 :: [Word]
ds0 = case CumulDims -> SomeDims
fromSteps ([Word] -> CumulDims
CumulDims [Word]
steps0) of SomeDims Dims ns
x -> Dims ns -> [Word]
forall k (xs :: [k]). Dims xs -> [Word]
listDims Dims ns
x
#endif
    go :: Word -> Word -> [Word] -> [Word] -> (Word, [Word])
    go :: Word -> Word -> [Word] -> [Word] -> (Word, [Word])
go Word
0 Word
off [Word]
steps [] = (Word
off, [Word]
steps)
    go Word
sub Word
off ~(steps :: [Word]
steps@(Word
s:[Word]
_)) [] = (Word
off, Word
subWord -> Word -> Word
forall a. Num a => a -> a -> a
*Word
s Word -> [Word] -> [Word]
forall a. a -> [a] -> [a]
: [Word]
steps)
#ifndef UNSAFE_INDICES
    go Word
sub Word
_ ~(Word
bs:(Word
s:[Word]
_)) [Word
i]
      | Word
b <- Word -> Word -> Word
forall a. Integral a => a -> a -> a
quot Word
bs Word
s
      , Word
sub Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
> Word
0 Bool -> Bool -> Bool
&& Word
i Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
sub Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
> Word
b
      = String
-> Word
-> Word
-> Maybe Word
-> Maybe ([Word], [Word])
-> (Word, [Word])
forall i a.
Integral i =>
String -> i -> Word -> Maybe Word -> Maybe ([Word], [Word]) -> a
outOfDimBoundsNoCallStack
          String
"{Calculating SubDataFrame offset}"
          Word
i Word
b (Word -> Maybe Word
forall a. a -> Maybe a
Just Word
sub) (([Word], [Word]) -> Maybe ([Word], [Word])
forall a. a -> Maybe a
Just ([Word]
ds0, [Word]
is0))
    go Word
_ Word
_ ~(Word
bs:(Word
s:[Word]
_)) (Word
i:[Word]
_)
      | Word
b <- Word -> Word -> Word
forall a. Integral a => a -> a -> a
quot Word
bs Word
s
      , Word
i Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
>= Word
b
      = String
-> Word
-> Word
-> Maybe Word
-> Maybe ([Word], [Word])
-> (Word, [Word])
forall i a.
Integral i =>
String -> i -> Word -> Maybe Word -> Maybe ([Word], [Word]) -> a
outOfDimBoundsNoCallStack
          String
"{Calculating SubDataFrame offset}"
          Word
i Word
b Maybe Word
forall a. Maybe a
Nothing (([Word], [Word]) -> Maybe ([Word], [Word])
forall a. a -> Maybe a
Just ([Word]
ds0, [Word]
is0))
#endif
    go Word
sub Word
off ~(Word
_:steps :: [Word]
steps@(Word
s:[Word]
_)) (Word
i:[Word]
ixs)
      = Word -> Word -> [Word] -> [Word] -> (Word, [Word])
go Word
sub (Word
off Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
iWord -> Word -> Word
forall a. Num a => a -> a -> a
*Word
s) [Word]
steps [Word]
ixs
{-# INLINE getOffAndSteps' #-}



-- | Same as `cdIx`, but safe; returns @Nothing@ if out of bounds.
cdIxM :: CumulDims -> Idxs ns -> Maybe Int
cdIxM :: CumulDims -> Idxs ns -> Maybe Int
cdIxM CumulDims
steps
  = ((Word, [Word]) -> Int) -> Maybe (Word, [Word]) -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Int) -> ((Word, [Word]) -> Word) -> (Word, [Word]) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word, [Word]) -> Word
forall a b. (a, b) -> a
fst)
  (Maybe (Word, [Word]) -> Maybe Int)
-> (Idxs ns -> Maybe (Word, [Word])) -> Idxs ns -> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Word -> [Word] -> [Word] -> Maybe (Word, [Word])
getOffAndStepsM' Word
0 Word
0 (CumulDims -> [Word]
unCumulDims CumulDims
steps) ([Word] -> Maybe (Word, [Word]))
-> (Idxs ns -> [Word]) -> Idxs ns -> Maybe (Word, [Word])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Idxs ns -> [Word]
forall k (ds :: [k]). Idxs ds -> [Word]
listIdxs

-- | Same as `getOffAndSteps`, but safe; returns @Nothing@ if out of bounds.
--   Trims the first (slicing) dimension of the returned CumulDims to fit
--   the original dataframe if necessary.
getOffAndStepsM :: Int -- ^ Initial offset
                -> CumulDims -> Idxs ns -> Maybe (Int, CumulDims)
getOffAndStepsM :: Int -> CumulDims -> Idxs ns -> Maybe (Int, CumulDims)
getOffAndStepsM Int
off0 CumulDims
steps
  = ((Word, [Word]) -> (Int, CumulDims))
-> Maybe (Word, [Word]) -> Maybe (Int, CumulDims)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Int)
-> ([Word] -> CumulDims) -> (Word, [Word]) -> (Int, CumulDims)
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** [Word] -> CumulDims
CumulDims)
  (Maybe (Word, [Word]) -> Maybe (Int, CumulDims))
-> (Idxs ns -> Maybe (Word, [Word]))
-> Idxs ns
-> Maybe (Int, CumulDims)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Word -> [Word] -> [Word] -> Maybe (Word, [Word])
getOffAndStepsM' Word
0 (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
off0) (CumulDims -> [Word]
unCumulDims CumulDims
steps) ([Word] -> Maybe (Word, [Word]))
-> (Idxs ns -> [Word]) -> Idxs ns -> Maybe (Word, [Word])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Idxs ns -> [Word]
forall k (ds :: [k]). Idxs ds -> [Word]
listIdxs

-- | Same as `getOffAndStepsSub`, but safe; returns @Nothing@ if out of bounds.
--   Trims the first (slicing) dimension of the returned CumulDims to fit
--   the original dataframe if necessary.
getOffAndStepsSubM :: Int -- ^ Initial offset
                  -> CumulDims -> Idxs (ns +: idxN)
                  -> Dim subN -> Maybe (Int, CumulDims)
getOffAndStepsSubM :: Int
-> CumulDims
-> Idxs (ns +: idxN)
-> Dim subN
-> Maybe (Int, CumulDims)
getOffAndStepsSubM Int
off0 CumulDims
steps Idxs (ns +: idxN)
idxs Dim subN
d
  = ((Word, [Word]) -> (Int, CumulDims))
-> Maybe (Word, [Word]) -> Maybe (Int, CumulDims)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Int)
-> ([Word] -> CumulDims) -> (Word, [Word]) -> (Int, CumulDims)
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** [Word] -> CumulDims
CumulDims)
  (Maybe (Word, [Word]) -> Maybe (Int, CumulDims))
-> ([Word] -> Maybe (Word, [Word]))
-> [Word]
-> Maybe (Int, CumulDims)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Word -> [Word] -> [Word] -> Maybe (Word, [Word])
getOffAndStepsM' (Dim subN -> Word
forall k (x :: k). Dim x -> Word
dimVal Dim subN
d) (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
off0) (CumulDims -> [Word]
unCumulDims CumulDims
steps)
  ([Word] -> Maybe (Int, CumulDims))
-> [Word] -> Maybe (Int, CumulDims)
forall a b. (a -> b) -> a -> b
$ Idxs (ns +: idxN) -> [Word]
forall k (ds :: [k]). Idxs ds -> [Word]
listIdxs Idxs (ns +: idxN)
idxs

getOffAndStepsM' :: Word -> Word -> [Word] -> [Word] -> Maybe (Word, [Word])
getOffAndStepsM' :: Word -> Word -> [Word] -> [Word] -> Maybe (Word, [Word])
getOffAndStepsM' = Word -> Word -> [Word] -> [Word] -> Maybe (Word, [Word])
go
  where
    go :: Word -> Word -> [Word] -> [Word] -> Maybe (Word, [Word])
    go :: Word -> Word -> [Word] -> [Word] -> Maybe (Word, [Word])
go Word
_   Word
off [Word]
steps [] = (Word, [Word]) -> Maybe (Word, [Word])
forall a. a -> Maybe a
Just (Word
off, [Word]
steps)
    go Word
sub Word
off ~(Word
bs:steps :: [Word]
steps@(Word
s:[Word]
_)) (Word
i:[Word]
ixs)
      | Word
is Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
>= Word
bs       = Maybe (Word, [Word])
forall a. Maybe a
Nothing
      | Bool -> Bool
not ([Word] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Word]
ixs) = Word -> Word -> [Word] -> [Word] -> Maybe (Word, [Word])
go Word
sub (Word
off Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
is) [Word]
steps [Word]
ixs
      | Word
sub Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
0       = (Word, [Word]) -> Maybe (Word, [Word])
forall a. a -> Maybe a
Just (Word
off Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
is, [Word]
steps)
      | Bool
otherwise      = (Word, [Word]) -> Maybe (Word, [Word])
forall a. a -> Maybe a
Just (Word
off Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
is, Word -> Word -> Word
forall a. Ord a => a -> a -> a
min (Word
bs Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
is) (Word
subWord -> Word -> Word
forall a. Num a => a -> a -> a
*Word
s) Word -> [Word] -> [Word]
forall a. a -> [a] -> [a]
: [Word]
steps)
      where
        is :: Word
is = Word
iWord -> Word -> Word
forall a. Num a => a -> a -> a
*Word
s
{-# INLINE getOffAndStepsM' #-}

-- | Try to get @CumulDims@ from an array,
--   and create it using @Dims@ if failed.
getSteps :: PrimArray t a => Dims (ns :: [k]) -> a -> CumulDims
getSteps :: Dims ns -> a -> CumulDims
getSteps Dims ns
dds = (t -> CumulDims)
-> (CumulDims -> Int# -> ByteArray# -> CumulDims) -> a -> CumulDims
forall t a r.
PrimArray t a =>
(t -> r) -> (CumulDims -> Int# -> ByteArray# -> r) -> a -> r
withArrayContent (CumulDims -> t -> CumulDims
forall a b. a -> b -> a
const (CumulDims -> t -> CumulDims) -> CumulDims -> t -> CumulDims
forall a b. (a -> b) -> a -> b
$ Dims ns -> CumulDims
forall k (ns :: [k]). Dims ns -> CumulDims
cumulDims Dims ns
dds) (\CumulDims
cd Int#
_ ByteArray#
_ -> CumulDims
cd)
{-# INLINE getSteps #-}

-- | Get @Dims@ by "de-accumulating" @CumulDims@.
fromSteps :: CumulDims -> SomeDims
fromSteps :: CumulDims -> SomeDims
fromSteps = [Word] -> SomeDims
someDimsVal ([Word] -> SomeDims)
-> (CumulDims -> [Word]) -> CumulDims -> SomeDims
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Word] -> [Word]
f ([Word] -> [Word]) -> (CumulDims -> [Word]) -> CumulDims -> [Word]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CumulDims -> [Word]
unCumulDims
  where
    -- ignore last value, which is always 1
    f :: [Word] -> [Word]
    f :: [Word] -> [Word]
f []       = []
    f [Word
_]      = []
    f [Word
n,Word
_]    = [Word
n]
    f (Word
a:Word
b:[Word]
cs) = Word
a Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
b Word -> [Word] -> [Word]
forall a. a -> [a] -> [a]
: [Word] -> [Word]
f (Word
bWord -> [Word] -> [Word]
forall a. a -> [a] -> [a]
:[Word]
cs)
{-# INLINE fromSteps #-}

class PrimBytes t => PrimArray t a | a -> t where
    -- | Broadcast element into array
    --
    --   Warning: do not use this function at the call site; use `broadcast`
    --            instead. Otherwise you will miss some rewrite rules.
    broadcast# :: t -> a
    -- | Index an array given an offset
    ix# :: Int# -> a -> t
    ix# Int#
i = (t -> t) -> (CumulDims -> Int# -> ByteArray# -> t) -> a -> t
forall t a r.
PrimArray t a =>
(t -> r) -> (CumulDims -> Int# -> ByteArray# -> r) -> a -> r
withArrayContent t -> t
forall a. a -> a
id (\CumulDims
_ Int#
off ByteArray#
arr -> ByteArray# -> Int# -> t
forall a. PrimBytes a => ByteArray# -> Int# -> a
indexArray ByteArray#
arr (Int#
off Int# -> Int# -> Int#
+# Int#
i))
    {-# INLINE ix# #-}
    -- | Generate an array using an accumulator funtion
    gen# :: CumulDims
            -- ^ Dimensionality of the result array;
            --   Be careful! @ns@ depends on @a@, but this is not reflected in
            --   types and is not checked at runtime.
         -> (s -> (# s, t #))
         -> s -> (# s, a #)
    -- | update a single element in an array given an offset
    upd# :: CumulDims
            -- ^ Dimensionality of the result array;
            --   Be careful! @ns@ depends on @a@, but this is not reflected in
            --   types and is not checked at runtime.
         -> Int# -> t -> a -> a

    -- | If the array is represented as a single broadcasted value, return this
    --   this value. Otherwise, return the full array content:
    --    @CumulDims@, array offset (elements), byte array with the content.
    --
    --   Warning: never use this function directly. Use `withArrayContent` instead.
    --            There is a bug in GHC 8.6, such that certain optimizations
    --            (probably, instance specialization/rewrite rules) break the code,
    --            which is only observable at runtime. The effect is that the
    --            content of a `ByteArray#` becomes a garbage. The workaround is
    --            to use a non-inlinable wrapper to disable these optimizations.
    --            In addition, the wrapper function has some rewrite rules, which
    --            can potentially improve performance with other GHC versions.
    withArrayContent# :: forall (rep :: RuntimeRep) (r :: TYPE rep)
                       . (t -> r)
                      -> (CumulDims -> Int# -> ByteArray# -> r)
                      -> a -> r

    -- | Offset of an array as a number of elements
    offsetElems :: a -> Int#
    offsetElems a
a = (t -> (Int# -> Int#) -> Int#)
-> (CumulDims -> Int# -> ByteArray# -> (Int# -> Int#) -> Int#)
-> a
-> (Int# -> Int#)
-> Int#
forall t a r.
PrimArray t a =>
(t -> r) -> (CumulDims -> Int# -> ByteArray# -> r) -> a -> r
withArrayContent (\t
_ Int# -> Int#
f -> Int# -> Int#
f Int#
0#) ((Int# -> ByteArray# -> (Int# -> Int#) -> Int#)
-> CumulDims -> Int# -> ByteArray# -> (Int# -> Int#) -> Int#
forall a b. a -> b -> a
const ((Int# -> ByteArray# -> (Int# -> Int#) -> Int#)
 -> CumulDims -> Int# -> ByteArray# -> (Int# -> Int#) -> Int#)
-> (Int# -> ByteArray# -> (Int# -> Int#) -> Int#)
-> CumulDims
-> Int#
-> ByteArray#
-> (Int# -> Int#)
-> Int#
forall a b. (a -> b) -> a -> b
$ \Int#
o ByteArray#
_ Int# -> Int#
f -> Int# -> Int#
f Int#
o) a
a (\Int#
i -> Int#
i)
    {-# INLINE offsetElems #-}

    -- | Normally, this returns a cumulative @totalDim@s.
    --   However, if a particular implementation does not have the dimensionality
    --   information, it cannot return @CumulDims@;
    --   In this case, it is a sign that all elements of an array are same.
    --   Thus, it is possible to return the single element value instead.
    --
    --   Note, this function returns the only unique element only if it is
    --   a such by construction (there is no equality checks involved).
    uniqueOrCumulDims :: a -> Either t CumulDims
    uniqueOrCumulDims = (t -> Either t CumulDims)
-> (CumulDims -> Int# -> ByteArray# -> Either t CumulDims)
-> a
-> Either t CumulDims
forall t a r.
PrimArray t a =>
(t -> r) -> (CumulDims -> Int# -> ByteArray# -> r) -> a -> r
withArrayContent t -> Either t CumulDims
forall a b. a -> Either a b
Left (\CumulDims
cd Int#
_ ByteArray#
_ -> CumulDims -> Either t CumulDims
forall a b. b -> Either a b
Right CumulDims
cd)
    {-# INLINE uniqueOrCumulDims #-}

    -- | Define an array by its offset and cumulative dims in a ByteArray.
    --   Both offset and dims are given in element number (not in bytes).
    --
    --   Warning: never use this function directly. Use `fromElems` instead.
    --            There is a bug in GHC 8.6, such that certain optimizations
    --            (probably, instance specialization/rewrite rules) break the code,
    --            which is only observable at runtime. The effect is that the
    --            content of a `ByteArray#` becomes a garbage. The workaround is
    --            to use a non-inlinable wrapper to disable these optimizations.
    --            In addition, the wrapper function has some rewrite rules, which
    --            can potentially improve performance with other GHC versions.
    fromElems# :: CumulDims -> Int# -> ByteArray# -> a

{-# WARNING fromElems# "Please, use fromElems instead." #-}
{-# WARNING withArrayContent# "Please, use withArrayContent instead." #-}

-- | Index array by an integer offset (starting from 0).
ixOff :: PrimArray t a => Int -> a -> t
ixOff :: Int -> a -> t
ixOff (I# Int#
i) = Int# -> a -> t
forall t a. PrimArray t a => Int# -> a -> t
ix# Int#
i

-- | Construct an array from a flat list and @Dims@;
--   Be careful! @ns@ depends on @a@, but this is not reflected in
--   types and is not checked at runtime.
unsafeFromFlatList :: PrimArray t a => Dims ns -> t -> [t] -> a
unsafeFromFlatList :: Dims ns -> t -> [t] -> a
unsafeFromFlatList Dims ns
ds t
x0 [t]
vs = case CumulDims -> ([t] -> (# [t], t #)) -> [t] -> (# [t], a #)
forall t a s.
PrimArray t a =>
CumulDims -> (s -> (# s, t #)) -> s -> (# s, a #)
gen# (Dims ns -> CumulDims
forall k (ns :: [k]). Dims ns -> CumulDims
cumulDims Dims ns
ds) [t] -> (# [t], t #)
f [t]
vs of (# [t]
_, a
r #) -> a
r
  where
    f :: [t] -> (# [t], t #)
f []     = (# [], t
x0 #)
    f (t
x:[t]
xs) = (# [t]
xs, t
x #)

-- | If the array is represented as a single broadcasted value, return this
--   this value. Otherwise, return the full array content:
--    @CumulDims@, array offset (elements), byte array with the content.
withArrayContent :: forall (t :: Type) (a :: Type)
                           (rep :: RuntimeRep) (r :: TYPE rep)
                  . PrimArray t a
                 => (t -> r)
                 -> (CumulDims -> Int# -> ByteArray# -> r)
                 -> a -> r
withArrayContent :: (t -> r) -> (CumulDims -> Int# -> ByteArray# -> r) -> a -> r
withArrayContent = (t -> r) -> (CumulDims -> Int# -> ByteArray# -> r) -> a -> r
forall t a r.
PrimArray t a =>
(t -> r) -> (CumulDims -> Int# -> ByteArray# -> r) -> a -> r
withArrayContent#
#if __GLASGOW_HASKELL__ == 806
{-# NOINLINE withArrayContent #-}
#else
{-# INLINE[1] withArrayContent #-}
#endif

-- | Define an array by its offset and cumulative dims in a ByteArray.
--   Both offset and dims are given in element number (not in bytes).
--
--   It is better to use this function instead of @fromBytes@ to avoid
--   recalculating @CumulDims@ for implementations that require it.
fromElems :: forall (t :: Type) (a :: Type)
           . PrimArray t a => CumulDims -> Int# -> ByteArray# -> a
fromElems :: CumulDims -> Int# -> ByteArray# -> a
fromElems = CumulDims -> Int# -> ByteArray# -> a
forall t a. PrimArray t a => CumulDims -> Int# -> ByteArray# -> a
fromElems#
#if __GLASGOW_HASKELL__ == 806
{-# NOINLINE fromElems #-}
#else
{-# INLINE[1] fromElems #-}
#endif

-- | Broadcast element into array
broadcast :: forall (t :: Type) (a :: Type)
           . PrimArray t a => t -> a
broadcast :: t -> a
broadcast = t -> a
forall t a. PrimArray t a => t -> a
broadcast#
{-# INLINE[1] broadcast #-}

{-# RULES

"withArrayContent/id"
  withArrayContent broadcast fromElems = id

"withArrayContent+fromElems" forall f g cd off ba .
  withArrayContent f g (fromElems cd off ba) = g cd off ba

"withArrayContent+broadcast" forall f g e .
  withArrayContent f g (broadcast e) = f e
  #-}