{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE UnboxedTuples #-}
module Numeric.DataFrame.Internal.PrimArray
( PrimArray (..), CumulDims (..)
, cumulDims, cdTotalDim, cdTotalDim#, cdIx
, ixOff, unsafeFromFlatList, getSteps, fromSteps
) where
import Data.Monoid as Mon (Monoid (..))
import Data.Semigroup as Sem (Semigroup (..))
import GHC.Base (ByteArray#, Int (..), Int#, Word (..), word2Int#)
import Numeric.Dimensions
import Numeric.PrimBytes
newtype CumulDims = CumulDims { unCumulDims :: [Word] }
instance Sem.Semigroup CumulDims where
CumulDims as <> CumulDims bs = CumulDims $ map (head bs *) (init as) ++ bs
instance Mon.Monoid CumulDims where
mempty = CumulDims [1]
#if !(MIN_VERSION_base(4,11,0))
mappend = (<>)
#endif
cumulDims :: Dims (ns :: [k]) -> CumulDims
cumulDims = CumulDims . uncurry (:)
. foldr (\d (c, cs) -> (c*d,c:cs)) (1, []) . listDims
cdTotalDim :: CumulDims -> Word
cdTotalDim ~(CumulDims ~(n:_)) = n
cdTotalDim# :: CumulDims -> Int#
cdTotalDim# ~(CumulDims ~(n:_)) = case n of W# w -> word2Int# w
cdIx :: CumulDims -> Idxs ns -> Int
cdIx ~(CumulDims ~(_:steps))
= fromIntegral . sum . zipWith (*) steps . listIdxs
getSteps :: PrimArray t a => Dims (ns :: [k]) -> a -> CumulDims
getSteps dds df = case uniqueOrCumulDims df of
Left _ -> cumulDims dds
Right ds -> ds
{-# INLINE getSteps #-}
fromSteps :: CumulDims -> SomeDims
fromSteps = someDimsVal . f . unCumulDims
where
f :: [Word] -> [Word]
f [] = []
f [_] = []
f [n,_] = [n]
f (a:b:cs) = a `quot` b : f (b:cs)
{-# INLINE fromSteps #-}
class PrimBytes t => PrimArray t a | a -> t where
broadcast :: t -> a
ix# :: Int# -> a -> t
gen# :: CumulDims
-> (s -> (# s, t #))
-> s -> (# s, a #)
upd# :: CumulDims
-> Int# -> t -> a -> a
arrayContent# :: a -> (# t | (# CumulDims, Int#, ByteArray# #) #)
offsetElems :: a -> Int#
offsetElems a = case arrayContent# a of
(# _ | #) -> 0#
(# | (# _, o, _ #) #) -> o
{-# INLINE offsetElems #-}
uniqueOrCumulDims :: a -> Either t CumulDims
uniqueOrCumulDims a = case arrayContent# a of
(# x | #) -> Left x
(# | (# cd, _, _ #) #) -> Right cd
{-# INLINE uniqueOrCumulDims #-}
fromElems :: CumulDims -> Int# -> ByteArray# -> a
ixOff :: PrimArray t a => Int -> a -> t
ixOff (I# i) = ix# i
unsafeFromFlatList :: PrimArray t a => Dims ns -> t -> [t] -> a
unsafeFromFlatList ds x0 vs = case gen# (cumulDims ds) f vs of (# _, r #) -> r
where
f [] = (# [], x0 #)
f (x:xs) = (# xs, x #)