{-# 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)
import GHC.Exts (TYPE)
import Numeric.Dimensions
import Numeric.PrimBytes
newtype CumulDims = CumulDims { unCumulDims :: [Word] }
deriving Show
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 steps
= fromIntegral . fst . getOffAndSteps' 0 0 (unCumulDims steps) . listIdxs
cdIxSub :: CumulDims -> Idxs (ns +: idxN) -> Dim subN -> Int
cdIxSub steps idxs d
= fromIntegral
. fst . getOffAndSteps' (dimVal d) 0 (unCumulDims steps) $ listIdxs idxs
getOffAndSteps :: Int
-> CumulDims -> Idxs ns -> (Int, CumulDims)
getOffAndSteps off0 steps
= (fromIntegral *** CumulDims)
. getOffAndSteps' 0 (fromIntegral off0) (unCumulDims steps) . listIdxs
getOffAndStepsSub :: Int
-> CumulDims -> Idxs (ns +: idxN)
-> Dim subN -> (Int, CumulDims)
getOffAndStepsSub off0 steps idxs d
= (fromIntegral *** CumulDims)
. getOffAndSteps' (dimVal d) (fromIntegral off0) (unCumulDims steps)
$ listIdxs idxs
getOffAndSteps' :: Word -> Word -> [Word] -> [Word] -> (Word, [Word])
getOffAndSteps' sub0 off0 steps0 is0 = go sub0 off0 steps0 is0
where
#ifndef UNSAFE_INDICES
ds0 = case fromSteps (CumulDims steps0) of SomeDims x -> listDims x
#endif
go :: Word -> Word -> [Word] -> [Word] -> (Word, [Word])
go 0 off steps [] = (off, steps)
go sub off ~(steps@(s:_)) [] = (off, sub*s : steps)
#ifndef UNSAFE_INDICES
go sub _ ~(bs:(s:_)) [i]
| b <- quot bs s
, sub > 0 && i + sub > b
= outOfDimBoundsNoCallStack
"{Calculating SubDataFrame offset}"
i b (Just sub) (Just (ds0, is0))
go _ _ ~(bs:(s:_)) (i:_)
| b <- quot bs s
, i >= b
= outOfDimBoundsNoCallStack
"{Calculating SubDataFrame offset}"
i b Nothing (Just (ds0, is0))
#endif
go sub off ~(_:steps@(s:_)) (i:ixs)
= go sub (off + i*s) steps ixs
{-# INLINE getOffAndSteps' #-}
cdIxM :: CumulDims -> Idxs ns -> Maybe Int
cdIxM steps
= fmap (fromIntegral . fst)
. getOffAndStepsM' 0 0 (unCumulDims steps) . listIdxs
getOffAndStepsM :: Int
-> CumulDims -> Idxs ns -> Maybe (Int, CumulDims)
getOffAndStepsM off0 steps
= fmap (fromIntegral *** CumulDims)
. getOffAndStepsM' 0 (fromIntegral off0) (unCumulDims steps) . listIdxs
getOffAndStepsSubM :: Int
-> CumulDims -> Idxs (ns +: idxN)
-> Dim subN -> Maybe (Int, CumulDims)
getOffAndStepsSubM off0 steps idxs d
= fmap (fromIntegral *** CumulDims)
. getOffAndStepsM' (dimVal d) (fromIntegral off0) (unCumulDims steps)
$ listIdxs idxs
getOffAndStepsM' :: Word -> Word -> [Word] -> [Word] -> Maybe (Word, [Word])
getOffAndStepsM' = go
where
go :: Word -> Word -> [Word] -> [Word] -> Maybe (Word, [Word])
go _ off steps [] = Just (off, steps)
go sub off ~(bs:steps@(s:_)) (i:ixs)
| is >= bs = Nothing
| not (null ixs) = go sub (off + is) steps ixs
| sub == 0 = Just (off + is, steps)
| otherwise = Just (off + is, min (bs - is) (sub*s) : steps)
where
is = i*s
{-# INLINE getOffAndStepsM' #-}
getSteps :: PrimArray t a => Dims (ns :: [k]) -> a -> CumulDims
getSteps dds = withArrayContent (const $ cumulDims dds) (\cd _ _ -> cd)
{-# 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
ix# i = withArrayContent id (\_ off arr -> indexArray arr (off +# i))
{-# INLINE ix# #-}
gen# :: CumulDims
-> (s -> (# s, t #))
-> s -> (# s, a #)
upd# :: CumulDims
-> Int# -> t -> a -> a
withArrayContent# :: forall (rep :: RuntimeRep) (r :: TYPE rep)
. (t -> r)
-> (CumulDims -> Int# -> ByteArray# -> r)
-> a -> r
offsetElems :: a -> Int#
offsetElems a = withArrayContent (\_ f -> f 0#) (const $ \o _ f -> f o) a (\i -> i)
{-# INLINE offsetElems #-}
uniqueOrCumulDims :: a -> Either t CumulDims
uniqueOrCumulDims = withArrayContent Left (\cd _ _ -> Right cd)
{-# INLINE uniqueOrCumulDims #-}
fromElems# :: CumulDims -> Int# -> ByteArray# -> a
{-# WARNING fromElems# "Please, use fromElems instead." #-}
{-# WARNING withArrayContent# "Please, use withArrayContent instead." #-}
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 #)
withArrayContent :: forall (t :: Type) (a :: Type)
(rep :: RuntimeRep) (r :: TYPE rep)
. PrimArray t a
=> (t -> r)
-> (CumulDims -> Int# -> ByteArray# -> r)
-> a -> r
withArrayContent = withArrayContent#
#if __GLASGOW_HASKELL__ == 806
{-# NOINLINE withArrayContent #-}
#else
{-# INLINE[1] withArrayContent #-}
#endif
fromElems :: forall (t :: Type) (a :: Type)
. PrimArray t a => CumulDims -> Int# -> ByteArray# -> a
fromElems = fromElems#
#if __GLASGOW_HASKELL__ == 806
{-# NOINLINE fromElems #-}
#else
{-# INLINE[1] fromElems #-}
#endif
broadcast :: forall (t :: Type) (a :: Type)
. PrimArray t a => t -> a
broadcast = 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
#-}