{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UnboxedTuples #-}
{-# OPTIONS_GHC -fno-warn-redundant-constraints #-}
module Numeric.DataFrame.Internal.Mutable
( MDataFrame ()
, newDataFrame#, newPinnedDataFrame#
, copyDataFrame#, copyMDataFrame#
, copyDataFrame'#, copyMDataFrame'#
, freezeDataFrame#, unsafeFreezeDataFrame#
, thawDataFrame#, thawPinDataFrame#, unsafeThawDataFrame#
, writeDataFrame#, writeDataFrameOff#
, readDataFrame#, readDataFrameOff#
, withDataFramePtr#, isDataFramePinned#
) where
import GHC.Base
import Numeric.DataFrame.Internal.PrimArray
import Numeric.DataFrame.Type
import Numeric.Dimensions
import Numeric.PrimBytes
data MDataFrame s t (ns :: [Nat])
= MDataFrame# Int# CumulDims (MutableByteArray# s)
newDataFrame# :: forall t (ns :: [Nat]) s
. ( PrimBytes t, Dimensions ns)
=> State# s -> (# State# s, MDataFrame s t ns #)
newDataFrame# s0
| steps <- cumulDims $ dims @ns
, n <- cdTotalDim# steps
, (# s1, mba #) <- newByteArray# (n *# byteSize @t undefined) s0
= (# s1, MDataFrame# 0# steps mba #)
{-# INLINE newDataFrame# #-}
newPinnedDataFrame# :: forall t (ns :: [Nat]) s
. ( PrimBytes t, Dimensions ns)
=> State# s -> (# State# s, MDataFrame s t ns #)
newPinnedDataFrame# s0
| steps <- cumulDims $ dims @ns
, n <- cdTotalDim# steps
, (# s1, mba #) <- newAlignedPinnedByteArray#
(n *# byteSize @t undefined)
(byteAlign @t undefined) s0
= (# s1, MDataFrame# 0# steps mba #)
{-# INLINE newPinnedDataFrame# #-}
copyDataFrame# :: forall (t :: Type)
(b :: Nat) (bi :: Nat) (bd :: Nat)
(as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) s
. ( b ~ (bi + bd - 1)
, PrimBytes t
, PrimBytes (DataFrame t (bd :+ bs))
, ConcatList as (b :+ bs) asbs
)
=> Idxs (as +: bi) -> DataFrame t (bd :+ bs) -> MDataFrame s t asbs
-> State# s -> (# State# s, () #)
copyDataFrame# ei df (MDataFrame# off steps mba) s
| I# i <- cdIx steps ei
= (# writeBytes mba ((off +# i) *# byteSize @t undefined) df s, () #)
{-# INLINE copyDataFrame# #-}
{-# ANN copyMDataFrame# "HLint: ignore Use camelCase" #-}
copyMDataFrame# :: forall (t :: Type)
(b :: Nat) (bi :: Nat) (bd :: Nat)
(as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) s
. ( b ~ (bi + bd - 1)
, PrimBytes t
, ConcatList as (b :+ bs) asbs
)
=> Idxs (as +: bi) -> MDataFrame s t (bd :+ bs) -> MDataFrame s t asbs
-> State# s -> (# State# s, () #)
copyMDataFrame# ei df =
copyMDataFrame'#
(unsafeCoerce# ei :: Idxs as)
(unsafeCoerce# df :: MDataFrame s t (b :+ bs))
{-# INLINE copyMDataFrame# #-}
{-# ANN copyDataFrame'# "HLint: ignore Use camelCase" #-}
copyDataFrame'# :: forall (t :: Type) (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) s
. ( PrimBytes t
, PrimBytes (DataFrame t bs)
, ConcatList as bs asbs
)
=> Idxs as -> DataFrame t bs -> MDataFrame s t asbs
-> State# s -> (# State# s, () #)
copyDataFrame'# ei df (MDataFrame# off steps mba) s
| I# i <- cdIx steps ei
= (# writeBytes mba ((off +# i) *# byteSize @t undefined) df s, () #)
{-# INLINE copyDataFrame'# #-}
{-# ANN copyMDataFrame'# "HLint: ignore Use camelCase" #-}
copyMDataFrame'# :: forall (t :: Type) (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) s
. ( PrimBytes t
, ConcatList as bs asbs
)
=> Idxs as -> MDataFrame s t bs -> MDataFrame s t asbs
-> State# s -> (# State# s, () #)
copyMDataFrame'# ei (MDataFrame# offA stepsA arrA) (MDataFrame# offM stepsM arrM) s
| elS <- byteSize @t undefined
, lenA <- cdTotalDim# stepsA
, I# i <- cdIx stepsM ei
= (# copyMutableByteArray# arrA (offA *# elS)
arrM ((offM +# i) *# elS) (lenA *# elS) s
, () #)
{-# INLINE copyMDataFrame'# #-}
unsafeFreezeDataFrame# :: forall (t :: Type) (ns :: [Nat]) s
. PrimArray t (DataFrame t ns)
=> MDataFrame s t ns
-> State# s -> (# State# s, DataFrame t ns #)
unsafeFreezeDataFrame# (MDataFrame# offM steps arrM) s1
| (# s2, arrA #) <- unsafeFreezeByteArray# arrM s1
= (# s2, fromElems steps offM arrA #)
{-# INLINE unsafeFreezeDataFrame# #-}
freezeDataFrame# :: forall (t :: Type) (ns :: [Nat]) s
. PrimArray t (DataFrame t ns)
=> MDataFrame s t ns -> State# s -> (# State# s, DataFrame t ns #)
freezeDataFrame# (MDataFrame# offM steps arrM) s0
| elS <- byteSize @t undefined
, n <- cdTotalDim# steps
, (# s1, mba #) <- newByteArray# (n *# elS) s0
, s2 <- copyMutableByteArray# arrM (offM *# elS) mba 0# (n *# elS) s1
, (# s3, arrA #) <- unsafeFreezeByteArray# mba s2
= (# s3, fromElems steps 0# arrA #)
{-# INLINE freezeDataFrame# #-}
thawDataFrame# :: forall (t :: Type) (ns :: [Nat]) s
. ( Dimensions ns
, PrimBytes (DataFrame t ns))
=> DataFrame t ns -> State# s -> (# State# s, MDataFrame s t ns #)
thawDataFrame# df s0
| arrA <- getBytes df
, boff <- byteOffset df
, bsize <- byteSize df
, steps <- cumulDims $ dims @ns
, (# s1, arrM #) <- newByteArray# bsize s0
, s2 <- copyByteArray# arrA boff arrM 0# bsize s1
= (# s2, MDataFrame# 0# steps arrM #)
{-# INLINE thawDataFrame# #-}
thawPinDataFrame# :: forall (t :: Type) (ns :: [Nat]) s
. ( Dimensions ns
, PrimBytes (DataFrame t ns))
=> DataFrame t ns -> State# s -> (# State# s, MDataFrame s t ns #)
thawPinDataFrame# df s0
| arrA <- getBytes df
, boff <- byteOffset df
, bsize <- byteSize df
, steps <- cumulDims $ dims @ns
, (# s1, arrM #) <- newAlignedPinnedByteArray# bsize (byteAlign df) s0
, s2 <- copyByteArray# arrA boff arrM 0# bsize s1
= (# s2, MDataFrame# 0# steps arrM #)
{-# INLINE thawPinDataFrame# #-}
unsafeThawDataFrame# :: forall (t :: Type) (ns :: [Nat]) s
. ( Dimensions ns
, PrimBytes (DataFrame t ns), PrimBytes t)
=> DataFrame t ns
-> State# s -> (# State# s, MDataFrame s t ns #)
unsafeThawDataFrame# df s0
| elS <- byteSize @t undefined
, arrA <- getBytes df
, boff <- byteOffset df
, steps <- cumulDims $ dims @ns
= (# s0
, MDataFrame# (quotInt# boff elS) steps (unsafeCoerce# arrA)
#)
{-# INLINE unsafeThawDataFrame# #-}
writeDataFrameOff# :: forall (t :: Type) (ns :: [Nat]) s
. PrimBytes t
=> MDataFrame s t ns -> Int# -> t -> State# s -> (# State# s, () #)
writeDataFrameOff# (MDataFrame# off _ mba) i x s
= (# writeArray mba (off +# i) x s, () #)
{-# INLINE writeDataFrameOff# #-}
writeDataFrame# :: forall (t :: Type) (ns :: [Nat]) s
. PrimBytes t
=> MDataFrame s t ns -> Idxs ns -> t -> State# s -> (# State# s, () #)
writeDataFrame# mdf@(MDataFrame# _ st _) ei
| I# i <- cdIx st ei = writeDataFrameOff# mdf i
{-# INLINE writeDataFrame# #-}
readDataFrameOff# :: forall (t :: Type) (ns :: [Nat]) s
. PrimBytes t
=> MDataFrame s t ns -> Int# -> State# s -> (# State# s, t #)
readDataFrameOff# (MDataFrame# off _ mba) i = readArray mba (off +# i)
{-# INLINE readDataFrameOff# #-}
readDataFrame# :: forall (t :: Type) (ns :: [Nat]) s
. PrimBytes t
=> MDataFrame s t ns -> Idxs ns -> State# s -> (# State# s, t #)
readDataFrame# mdf@(MDataFrame# _ st _) ei
| I# i <- cdIx st ei = readDataFrameOff# mdf i
{-# INLINE readDataFrame# #-}
withDataFramePtr# :: forall (t :: Type) (ns :: [Nat]) (r :: Type)
. PrimBytes t
=> MDataFrame RealWorld t ns
-> ( Addr# -> State# RealWorld -> (# State# RealWorld, r #) )
-> State# RealWorld -> (# State# RealWorld, r #)
withDataFramePtr# (MDataFrame# off _ mba) k s0
| (# s1, a #) <- unsafeFreezeByteArray# mba s0
, (# s2, r #) <- k ( byteArrayContents# a
`plusAddr#` (off *# byteSize @t undefined)
) s1
= (# touch# mba s2, r #)
isDataFramePinned# :: forall (t :: Type) (ns :: [Nat]) s
. MDataFrame s t ns -> Bool
isDataFramePinned# (MDataFrame# _ _ mba)
= isTrue# (isMutableByteArrayPinned# mba)