{-# LANGUAGE DataKinds                 #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE MagicHash                 #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE PolyKinds                 #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeFamilies              #-}
{-# LANGUAGE TypeInType                #-}
{-# LANGUAGE TypeOperators             #-}
{-# LANGUAGE UnboxedTuples             #-}
{-# OPTIONS_GHC -fno-warn-redundant-constraints #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.DataFrame.Internal.Mutable
-- Copyright   :  (c) Artem Chirkin
-- License     :  BSD3
--
--
-- Interface to perform primitive stateful operations on mutable frames.
--
-----------------------------------------------------------------------------

module Numeric.DataFrame.Internal.Mutable
    ( MDataFrame ()
    , castDataFrame#
    , newDataFrame#, newPinnedDataFrame#
    , oneMoreDataFrame#
    , subDataFrameView#, subDataFrameView'#
    , copyDataFrame#, copyMDataFrame#
    , copyDataFrame'#, copyMDataFrame'#
    , copyDataFrameOff#, copyMDataFrameOff#
    , freezeDataFrame#, unsafeFreezeDataFrame#
    , thawDataFrame#, thawPinDataFrame#, unsafeThawDataFrame#, withThawDataFrame#
    , writeDataFrame#, writeDataFrameOff#
    , readDataFrame#, readDataFrameOff#
    , withDataFramePtr#, isDataFramePinned#
    , getDataFrameSteps#
    ) where

import GHC.Base (Type)
import GHC.Exts
import Numeric.DataFrame.Internal.PrimArray
import Numeric.DataFrame.Type
import Numeric.Dimensions
import Numeric.PrimBytes

-- | Mutable DataFrame type.
--   Keeps element offset, number of elements, and a mutable byte storage
data MDataFrame s t (ns :: [k])
  = MDataFrame# Int# CumulDims (MutableByteArray# s)

-- | Allow coercing between @XNat@-indexed and @Nat@-indexed Mutable DataFrames.
castDataFrame# ::
       forall (t :: Type) (xns :: [XNat]) (ns :: [Nat]) s
     . FixedDims xns ns
    => MDataFrame s t ns -> MDataFrame s t xns
castDataFrame# :: MDataFrame s t ns -> MDataFrame s t xns
castDataFrame# (MDataFrame# Int#
o CumulDims
c MutableByteArray# s
a) = Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t xns
forall k k s (t :: k) (ns :: [k]).
Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
MDataFrame# Int#
o CumulDims
c MutableByteArray# s
a
{-# INLINE castDataFrame# #-}

-- | Create a new mutable DataFrame.
newDataFrame# ::
       forall (t :: Type) (k :: Type) (ns :: [k]) s
     . (PrimBytes t, Dimensions ns)
    => State# s -> (# State# s, MDataFrame s t ns #)
newDataFrame# :: State# s -> (# State# s, MDataFrame s t ns #)
newDataFrame#
    | CumulDims
steps <- Dims ns -> CumulDims
forall k (ns :: [k]). Dims ns -> CumulDims
cumulDims (Dims ns -> CumulDims) -> Dims ns -> CumulDims
forall a b. (a -> b) -> a -> b
$ Dimensions ns => Dims ns
forall k (ds :: [k]). Dimensions ds => Dims ds
dims @ns
    , Int#
n <- CumulDims -> Int#
cdTotalDim# CumulDims
steps
      = if Int# -> Bool
isTrue# (Int#
n Int# -> Int# -> Int#
==# Int#
0#)
        then \State# s
s0 -> (# State# s
s0, [Char] -> MDataFrame s t ns
forall a. HasCallStack => [Char] -> a
error [Char]
"Empty DataFrame (DF0)" #)
        else \State# s
s0 -> case Int# -> State# s -> (# State# s, MutableByteArray# s #)
forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newByteArray# (Int#
n Int# -> Int# -> Int#
*# t -> Int#
forall a. PrimBytes a => a -> Int#
byteSize @t t
forall a. HasCallStack => a
undefined) State# s
s0 of
                          (# State# s
s1, MutableByteArray# s
mba #) -> (# State# s
s1,  Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
forall k k s (t :: k) (ns :: [k]).
Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
MDataFrame# Int#
0# CumulDims
steps MutableByteArray# s
mba #)
{-# INLINE newDataFrame# #-}

-- | Create a new mutable DataFrame.
newPinnedDataFrame# ::
       forall (t :: Type) (k :: Type) (ns :: [k]) s
     . (PrimBytes t, Dimensions ns)
    => State# s -> (# State# s, MDataFrame s t ns #)
newPinnedDataFrame# :: State# s -> (# State# s, MDataFrame s t ns #)
newPinnedDataFrame#
    | CumulDims
steps <- Dims ns -> CumulDims
forall k (ns :: [k]). Dims ns -> CumulDims
cumulDims (Dims ns -> CumulDims) -> Dims ns -> CumulDims
forall a b. (a -> b) -> a -> b
$ Dimensions ns => Dims ns
forall k (ds :: [k]). Dimensions ds => Dims ds
dims @ns
    , Int#
n <- CumulDims -> Int#
cdTotalDim# CumulDims
steps
      = if Int# -> Bool
isTrue# (Int#
n Int# -> Int# -> Int#
==# Int#
0#)
        then \State# s
s0 -> (# State# s
s0, [Char] -> MDataFrame s t ns
forall a. HasCallStack => [Char] -> a
error [Char]
"Empty DataFrame (DF0)" #)
        else \State# s
s0 -> case Int# -> Int# -> State# s -> (# State# s, MutableByteArray# s #)
forall d.
Int# -> Int# -> State# d -> (# State# d, MutableByteArray# d #)
newAlignedPinnedByteArray#
                          (Int#
n Int# -> Int# -> Int#
*# t -> Int#
forall a. PrimBytes a => a -> Int#
byteSize @t t
forall a. HasCallStack => a
undefined)
                          (t -> Int#
forall a. PrimBytes a => a -> Int#
byteAlign @t t
forall a. HasCallStack => a
undefined) State# s
s0 of
                          (# State# s
s1, MutableByteArray# s
mba #) -> (# State# s
s1,  Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
forall k k s (t :: k) (ns :: [k]).
Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
MDataFrame# Int#
0# CumulDims
steps MutableByteArray# s
mba #)
{-# INLINE newPinnedDataFrame# #-}

-- | Create a new mutable DataFrame of the same size.
oneMoreDataFrame# ::
       forall (t :: Type) (k :: Type) (ns :: [k]) s
     . MDataFrame s t ns -> State# s -> (# State# s, MDataFrame s t ns #)
oneMoreDataFrame# :: MDataFrame s t ns -> State# s -> (# State# s, MDataFrame s t ns #)
oneMoreDataFrame# mdf :: MDataFrame s t ns
mdf@(MDataFrame# Int#
off CumulDims
steps MutableByteArray# s
mba) State# s
s0
    | Int#
0# <- CumulDims -> Int#
cdTotalDim# CumulDims
steps = (# State# s
s0, MDataFrame s t ns
mdf #)
    | (# State# s
s1, Int#
bs #) <- MutableByteArray# s -> State# s -> (# State# s, Int# #)
forall d. MutableByteArray# d -> State# d -> (# State# d, Int# #)
getSizeofMutableByteArray# MutableByteArray# s
mba State# s
s0
    , (# State# s
s2, MutableByteArray# s
mba' #) <- Int# -> State# s -> (# State# s, MutableByteArray# s #)
forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newByteArray# (Int#
bs Int# -> Int# -> Int#
-# Int#
off) State# s
s1
      = (# State# s
s2,  Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
forall k k s (t :: k) (ns :: [k]).
Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
MDataFrame# Int#
0# CumulDims
steps MutableByteArray# s
mba' #)
{-# INLINE oneMoreDataFrame# #-}

-- | View a part of a DataFrame.
--
--   This function does not perform a copy.
--   All changes to a new DataFrame will be reflected in the original DataFrame as well.
--
--   If any of the dims in @as@ or @b@ is unknown (@a ~ XN m@),
--   then this function is unsafe and can throw an `OutOfDimBounds` exception.
--   Otherwise, its safety is guaranteed by the type system.
subDataFrameView# ::
       forall (t :: Type) (k :: Type)
              (b :: k) (bi :: k) (bd :: k)
              (as :: [k]) (bs :: [k]) (asbs :: [k]) s
     . (SubFrameIndexCtx b bi bd, KnownDim bd, ConcatList as (b :+ bs) asbs)
    => Idxs (as +: bi) -> MDataFrame s t asbs -> MDataFrame s t (bd :+ bs)
subDataFrameView# :: Idxs (as +: bi) -> MDataFrame s t asbs -> MDataFrame s t (bd :+ bs)
subDataFrameView# Idxs (as +: bi)
ei (MDataFrame# Int#
offM CumulDims
stepsM MutableByteArray# s
arr)
    = Int#
-> CumulDims -> MutableByteArray# s -> MDataFrame s t (bd :+ bs)
forall k k s (t :: k) (ns :: [k]).
Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
MDataFrame# (case Int
offA of I# Int#
i -> Int#
i) CumulDims
stepsA MutableByteArray# s
arr
  where
    (Int
offA, CumulDims
stepsA) = Int -> CumulDims -> Idxs (as +: bi) -> Dim bd -> (Int, CumulDims)
forall k1 k2 (ns :: [k1]) (idxN :: k1) (subN :: k2).
Int
-> CumulDims -> Idxs (ns +: idxN) -> Dim subN -> (Int, CumulDims)
getOffAndStepsSub (Int# -> Int
I# Int#
offM) CumulDims
stepsM Idxs (as +: bi)
ei (KnownDim bd => Dim bd
forall k (n :: k). KnownDim n => Dim n
dim @bd)

-- | View a part of a DataFrame.
--
--   This function does not perform a copy.
--   All changes to a new DataFrame will be reflected in the original DataFrame as well.
--
--   This is a simpler version of @subDataFrameView@ that allows
--    to view over one index at a time.
--
--   If any of the dims in @as@ is unknown (@a ~ XN m@),
--   then this function is unsafe and can throw an `OutOfDimBounds` exception.
--   Otherwise, its safety is guaranteed by the type system.
subDataFrameView'# ::
       forall (t :: Type) (k :: Type) (as :: [k]) (bs :: [k]) (asbs :: [k]) s
     . ConcatList as bs asbs
    => Idxs as -> MDataFrame s t asbs -> MDataFrame s t bs
subDataFrameView'# :: Idxs as -> MDataFrame s t asbs -> MDataFrame s t bs
subDataFrameView'# Idxs as
ei (MDataFrame# Int#
offM CumulDims
stepsM MutableByteArray# s
arr)
    = Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t bs
forall k k s (t :: k) (ns :: [k]).
Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
MDataFrame# (case Int
offA of I# Int#
i -> Int#
i) CumulDims
stepsA MutableByteArray# s
arr
  where
    (Int
offA, CumulDims
stepsA) = Int -> CumulDims -> Idxs as -> (Int, CumulDims)
forall k (ns :: [k]).
Int -> CumulDims -> Idxs ns -> (Int, CumulDims)
getOffAndSteps (Int# -> Int
I# Int#
offM) CumulDims
stepsM Idxs as
ei

-- | Copy one DataFrame into another mutable DataFrame at specified position.
--
--   In contrast to @copyDataFrame'@, this function allows to copy over a range
--    of contiguous indices over a single dimension.
--   For example, you can write a 3x4 matrix into a 7x4 matrix, starting at indices 0..3.
--
--   This function is safe (no `OutOfDimBounds` exception possible).
--   If any of the dims in @as@ is unknown (@a ~ XN m@),
--   you may happen to write data beyond dataframe bounds.
--   In this case, this function does nothing.
--   If (@b ~ XN m@) and (@Idx bi + Dim bd > Dim b@), this function copies only as
--   many elements as fits into the dataframe along this dimension (possibly none).
copyDataFrame# ::
       forall (t :: Type) (k :: Type)
              (b :: k) (bi :: k) (bd :: k)
              (as :: [k]) (bs :: [k]) (asbs :: [k]) s
     . ( SubFrameIndexCtx b bi bd, KnownDim bd
       , ExactDims bs
       , PrimArray t (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# :: Idxs (as +: bi)
-> DataFrame t (bd :+ bs)
-> MDataFrame s t asbs
-> State# s
-> (# State# s, () #)
copyDataFrame# Idxs (as +: bi)
ei DataFrame t (bd :+ bs)
df (MDataFrame# Int#
offM CumulDims
stepsM MutableByteArray# s
arrDest)
    | Int#
elS <- t -> Int#
forall a. PrimBytes a => a -> Int#
byteSize @t t
forall a. HasCallStack => a
undefined
    , Just (I# Int#
offDest, CumulDims
stepsB)
         <- Int
-> CumulDims -> Idxs (as +: bi) -> Dim bd -> Maybe (Int, CumulDims)
forall k1 k2 (ns :: [k1]) (idxN :: k1) (subN :: k2).
Int
-> CumulDims
-> Idxs (ns +: idxN)
-> Dim subN
-> Maybe (Int, CumulDims)
getOffAndStepsSubM (Int# -> Int
I# Int#
offM) CumulDims
stepsM Idxs (as +: bi)
ei (KnownDim bd => Dim bd
forall k (n :: k). KnownDim n => Dim n
dim @bd)
    , Int#
n <- CumulDims -> Int#
cdTotalDim# CumulDims
stepsB
    , Int# -> Bool
isTrue# (Int#
n Int# -> Int# -> Int#
># Int#
0#) -- is there enough space to write anything?
    = (t -> State# s -> (# State# s, () #))
-> (CumulDims
    -> Int# -> ByteArray# -> State# s -> (# State# s, () #))
-> DataFrame t (bd :+ bs)
-> State# s
-> (# State# s, () #)
forall t a r.
PrimArray t a =>
(t -> r) -> (CumulDims -> Int# -> ByteArray# -> r) -> a -> r
withArrayContent
      (\t
e State# s
s -> (# MutableByteArray# s -> Int# -> Int# -> t -> State# s -> State# s
forall t s.
PrimBytes t =>
MutableByteArray# s -> Int# -> Int# -> t -> State# s -> State# s
fillArray MutableByteArray# s
arrDest Int#
offDest Int#
n t
e State# s
s, () #))
      (\CumulDims
_ Int#
offSrc ByteArray#
arrSrc State# s
s ->
        (# ByteArray#
-> Int#
-> MutableByteArray# s
-> Int#
-> Int#
-> State# s
-> State# s
forall d.
ByteArray#
-> Int#
-> MutableByteArray# d
-> Int#
-> Int#
-> State# d
-> State# d
copyByteArray# ByteArray#
arrSrc  (Int#
offSrc Int# -> Int# -> Int#
*# Int#
elS)
                          MutableByteArray# s
arrDest (Int#
offDest Int# -> Int# -> Int#
*# Int#
elS) (Int#
n Int# -> Int# -> Int#
*# Int#
elS) State# s
s, () #)) DataFrame t (bd :+ bs)
df
    | Bool
otherwise = \State# s
s -> (# State# s
s, () #)
{-# INLINE copyDataFrame# #-}

{-# ANN copyMDataFrame# "HLint: ignore Use camelCase" #-}
-- | Copy one mutable DataFrame into another mutable DataFrame at specified position.
--
--   In contrast to @copyMutableDataFrame'@, this function allows to copy over a range
--    of contiguous indices over a single dimension.
--   For example, you can write a 3x4 matrix into a 7x4 matrix, starting at indices 0..3.
--
--   This function is safe (no `OutOfDimBounds` exception possible).
--   If any of the dims in @as@ is unknown (@a ~ XN m@),
--   you may happen to write data beyond dataframe bounds.
--   In this case, this function does nothing.
--   If (@b ~ XN m@) and (@Idx bi + Dim bd > Dim b@), this function copies only as
--   many elements as fits into the dataframe along this dimension (possibly none).
copyMDataFrame# ::
       forall (t :: Type) (k :: Type)
              (b :: k) (bi :: k) (bd :: k)
              (as :: [k]) (bs :: [k]) (asbs :: [k]) s
     . ( SubFrameIndexCtx b bi bd
       , ExactDims bs
       , PrimBytes t
       , ConcatList as (b :+ bs) asbs )
    => Idxs (as +: bi) -> MDataFrame s t (bd :+ bs) -> MDataFrame s t asbs
    -> State# s -> (# State# s, () #)
copyMDataFrame# :: Idxs (as +: bi)
-> MDataFrame s t (bd :+ bs)
-> MDataFrame s t asbs
-> State# s
-> (# State# s, () #)
copyMDataFrame# Idxs (as +: bi)
ei (MDataFrame# Int#
offA (CumulDims ~(Word
bb:Word
b:[Word]
_)) MutableByteArray# s
arrA)
                   (MDataFrame# Int#
offM CumulDims
stepsM MutableByteArray# s
arrM)
    | Int#
elS <- t -> Int#
forall a. PrimBytes a => a -> Int#
byteSize @t t
forall a. HasCallStack => a
undefined
    , Just (I# Int#
offDest, CumulDims
stepsB)
         <- Int
-> CumulDims
-> Idxs (as +: bi)
-> Dim Any
-> Maybe (Int, CumulDims)
forall k1 k2 (ns :: [k1]) (idxN :: k1) (subN :: k2).
Int
-> CumulDims
-> Idxs (ns +: idxN)
-> Dim subN
-> Maybe (Int, CumulDims)
getOffAndStepsSubM (Int# -> Int
I# Int#
offM) CumulDims
stepsM Idxs (as +: bi)
ei (Word -> Dim Any
unsafeCoerce# (Word -> Word -> Word
forall a. Integral a => a -> a -> a
quot Word
bb Word
b))
    , Int#
n <- CumulDims -> Int#
cdTotalDim# CumulDims
stepsB
    , Int# -> Bool
isTrue# (Int#
n Int# -> Int# -> Int#
># Int#
0#) -- is there enough space to write anything?
    = \State# s
s -> (# MutableByteArray# s
-> Int#
-> MutableByteArray# s
-> Int#
-> Int#
-> State# s
-> State# s
forall d.
MutableByteArray# d
-> Int#
-> MutableByteArray# d
-> Int#
-> Int#
-> State# d
-> State# d
copyMutableByteArray# MutableByteArray# s
arrA (Int#
offA Int# -> Int# -> Int#
*# Int#
elS)
                                     MutableByteArray# s
arrM (Int#
offDest Int# -> Int# -> Int#
*# Int#
elS) (Int#
n Int# -> Int# -> Int#
*# Int#
elS) State# s
s
             , () #)
    | Bool
otherwise = \State# s
s -> (# State# s
s, () #)
{-# INLINE copyMDataFrame# #-}

{-# ANN copyDataFrame'# "HLint: ignore Use camelCase" #-}
-- | Copy one DataFrame into another mutable DataFrame at specified position.
--
--   This is a simpler version of @copyDataFrame@ that allows
--     to copy over one index at a time.
--
--   This function is safe (no `OutOfDimBounds` exception possible).
--   If any of the dims in @as@ is unknown (@a ~ XN m@),
--   you may happen to write data beyond dataframe bounds.
--   In this case, this function does nothing.
copyDataFrame'# ::
       forall (t :: Type) (k :: Type) (as :: [k]) (bs :: [k]) (asbs :: [k]) s
     . ( ExactDims bs
       , PrimArray t (DataFrame t bs)
       , ConcatList as bs asbs )
    => Idxs as -> DataFrame t bs -> MDataFrame s t asbs
    -> State# s -> (# State# s, () #)
copyDataFrame'# :: Idxs as
-> DataFrame t bs
-> MDataFrame s t asbs
-> State# s
-> (# State# s, () #)
copyDataFrame'# Idxs as
ei DataFrame t bs
df (MDataFrame# Int#
offM CumulDims
stepsM MutableByteArray# s
arrDest)
    | Int#
elS <- t -> Int#
forall a. PrimBytes a => a -> Int#
byteSize @t t
forall a. HasCallStack => a
undefined
    , Just (I# Int#
offDest, CumulDims
stepsA) <- Int -> CumulDims -> Idxs as -> Maybe (Int, CumulDims)
forall k (ns :: [k]).
Int -> CumulDims -> Idxs ns -> Maybe (Int, CumulDims)
getOffAndStepsM (Int# -> Int
I# Int#
offM) CumulDims
stepsM Idxs as
ei
    , Int#
n <- CumulDims -> Int#
cdTotalDim# CumulDims
stepsA
    = (t -> State# s -> (# State# s, () #))
-> (CumulDims
    -> Int# -> ByteArray# -> State# s -> (# State# s, () #))
-> DataFrame t bs
-> State# s
-> (# State# s, () #)
forall t a r.
PrimArray t a =>
(t -> r) -> (CumulDims -> Int# -> ByteArray# -> r) -> a -> r
withArrayContent
      (\t
e State# s
s -> (# MutableByteArray# s -> Int# -> Int# -> t -> State# s -> State# s
forall t s.
PrimBytes t =>
MutableByteArray# s -> Int# -> Int# -> t -> State# s -> State# s
fillArray MutableByteArray# s
arrDest Int#
offDest Int#
n t
e State# s
s, () #))
      (\CumulDims
_ Int#
offSrc ByteArray#
arrSrc State# s
s ->
        (# ByteArray#
-> Int#
-> MutableByteArray# s
-> Int#
-> Int#
-> State# s
-> State# s
forall d.
ByteArray#
-> Int#
-> MutableByteArray# d
-> Int#
-> Int#
-> State# d
-> State# d
copyByteArray# ByteArray#
arrSrc  (Int#
offSrc Int# -> Int# -> Int#
*# Int#
elS)
                          MutableByteArray# s
arrDest (Int#
offDest Int# -> Int# -> Int#
*# Int#
elS) (Int#
n Int# -> Int# -> Int#
*# Int#
elS) State# s
s, () #)) DataFrame t bs
df
    | Bool
otherwise = \State# s
s -> (# State# s
s, () #)
{-# INLINE copyDataFrame'# #-}

{-# ANN copyMDataFrame'# "HLint: ignore Use camelCase" #-}
-- | Copy one mutable DataFrame into another mutable DataFrame at specified position.
--
--   This is a simpler version of @copyMutableDataFrame@ that allows
--     to copy over one index at a time.
--
--   This function is safe (no `OutOfDimBounds` exception possible).
--   If any of the dims in @as@ is unknown (@a ~ XN m@),
--   you may happen to write data beyond dataframe bounds.
--   In this case, this function does nothing.
copyMDataFrame'# ::
       forall (t :: Type) (k :: Type) (as :: [k]) (bs :: [k]) (asbs :: [k]) s
     . (ExactDims bs, PrimBytes t, ConcatList as bs asbs)
    => Idxs as -> MDataFrame s t bs -> MDataFrame s t asbs
    -> State# s -> (# State# s, () #)
copyMDataFrame'# :: Idxs as
-> MDataFrame s t bs
-> MDataFrame s t asbs
-> State# s
-> (# State# s, () #)
copyMDataFrame'# Idxs as
ei (MDataFrame# Int#
offA CumulDims
stepsA MutableByteArray# s
arrA) (MDataFrame# Int#
offM CumulDims
stepsM MutableByteArray# s
arrM)
    | Int#
elS <- t -> Int#
forall a. PrimBytes a => a -> Int#
byteSize @t t
forall a. HasCallStack => a
undefined
    , Int#
lenA <- CumulDims -> Int#
cdTotalDim# CumulDims
stepsA
    , Just (I# Int#
i) <- CumulDims -> Idxs as -> Maybe Int
forall k (ns :: [k]). CumulDims -> Idxs ns -> Maybe Int
cdIxM CumulDims
stepsM Idxs as
ei
    = \State# s
s -> (# MutableByteArray# s
-> Int#
-> MutableByteArray# s
-> Int#
-> Int#
-> State# s
-> State# s
forall d.
MutableByteArray# d
-> Int#
-> MutableByteArray# d
-> Int#
-> Int#
-> State# d
-> State# d
copyMutableByteArray#
                 MutableByteArray# s
arrA (Int#
offA Int# -> Int# -> Int#
*# Int#
elS)
                 MutableByteArray# s
arrM ((Int#
offM Int# -> Int# -> Int#
+# Int#
i) Int# -> Int# -> Int#
*# Int#
elS) (Int#
lenA Int# -> Int# -> Int#
*# Int#
elS) State# s
s, () #)
    | Bool
otherwise = \State# s
s -> (# State# s
s, () #)
{-# INLINE copyMDataFrame'# #-}

-- | Copy one DataFrame into another mutable DataFrame by offset in
--   primitive elements.
--
--   This is a low-level copy function; you have to keep in mind the row-major
--   layout of Mutable DataFrames. Offset bounds are not checked.
--   You will get an undefined behavior if you write beyond the DataFrame bounds.
copyDataFrameOff# ::
       forall (t :: Type) (k :: Type) (as :: [k]) (bs :: [k]) (asbs :: [k]) s
     . ( Dimensions bs
       , PrimArray t (DataFrame t bs)
       , ConcatList as bs asbs )
    => Int -> DataFrame t bs -> MDataFrame s t asbs
    -> State# s -> (# State# s, () #)
copyDataFrameOff# :: Int
-> DataFrame t bs
-> MDataFrame s t asbs
-> State# s
-> (# State# s, () #)
copyDataFrameOff# (I# Int#
off) DataFrame t bs
df (MDataFrame# Int#
offM CumulDims
_ MutableByteArray# s
arrDest)
    | Int#
elS <- t -> Int#
forall a. PrimBytes a => a -> Int#
byteSize @t t
forall a. HasCallStack => a
undefined
    , Int#
offDest <- Int#
offM Int# -> Int# -> Int#
+# Int#
off
    = (t -> State# s -> (# State# s, () #))
-> (CumulDims
    -> Int# -> ByteArray# -> State# s -> (# State# s, () #))
-> DataFrame t bs
-> State# s
-> (# State# s, () #)
forall t a r.
PrimArray t a =>
(t -> r) -> (CumulDims -> Int# -> ByteArray# -> r) -> a -> r
withArrayContent
      (\t
e State# s
s ->
        (# MutableByteArray# s -> Int# -> Int# -> t -> State# s -> State# s
forall t s.
PrimBytes t =>
MutableByteArray# s -> Int# -> Int# -> t -> State# s -> State# s
fillArray MutableByteArray# s
arrDest Int#
offDest
              (case Dims bs -> Word
forall k (xs :: [k]). Dims xs -> Word
totalDim (Dimensions bs => Dims bs
forall k (ds :: [k]). Dimensions ds => Dims ds
dims @bs) of W# Word#
n -> Word# -> Int#
word2Int# Word#
n) t
e State# s
s, () #))
      (\CumulDims
steps Int#
offSrc ByteArray#
arrSrc State# s
s ->
        (# ByteArray#
-> Int#
-> MutableByteArray# s
-> Int#
-> Int#
-> State# s
-> State# s
forall d.
ByteArray#
-> Int#
-> MutableByteArray# d
-> Int#
-> Int#
-> State# d
-> State# d
copyByteArray# ByteArray#
arrSrc  (Int#
offSrc Int# -> Int# -> Int#
*# Int#
elS)
                          MutableByteArray# s
arrDest (Int#
offDest Int# -> Int# -> Int#
*# Int#
elS)
                                  (CumulDims -> Int#
cdTotalDim# CumulDims
steps Int# -> Int# -> Int#
*# Int#
elS) State# s
s, () #)) DataFrame t bs
df
{-# INLINE copyDataFrameOff# #-}

-- | Copy one mutable DataFrame into another mutable DataFrame by offset in
--   primitive elements.
--
--   This is a low-level copy function; you have to keep in mind the row-major
--   layout of Mutable DataFrames. Offset bounds are not checked.
--   You will get an undefined behavior if you write beyond the DataFrame bounds.
copyMDataFrameOff# ::
       forall (t :: Type) (k :: Type) (as :: [k]) (bs :: [k]) (asbs :: [k]) s
     . (ExactDims bs, PrimBytes t, ConcatList as bs asbs)
    => Int -> MDataFrame s t bs -> MDataFrame s t asbs
    -> State# s -> (# State# s, () #)
copyMDataFrameOff# :: Int
-> MDataFrame s t bs
-> MDataFrame s t asbs
-> State# s
-> (# State# s, () #)
copyMDataFrameOff# (I# Int#
off) (MDataFrame# Int#
offA CumulDims
stepsA MutableByteArray# s
arrA)
                            (MDataFrame# Int#
offM CumulDims
_      MutableByteArray# s
arrM)
    | Int#
elS <- t -> Int#
forall a. PrimBytes a => a -> Int#
byteSize @t t
forall a. HasCallStack => a
undefined
    , Int#
lenA <- CumulDims -> Int#
cdTotalDim# CumulDims
stepsA
    = \State# s
s -> (# MutableByteArray# s
-> Int#
-> MutableByteArray# s
-> Int#
-> Int#
-> State# s
-> State# s
forall d.
MutableByteArray# d
-> Int#
-> MutableByteArray# d
-> Int#
-> Int#
-> State# d
-> State# d
copyMutableByteArray#
                 MutableByteArray# s
arrA (Int#
offA Int# -> Int# -> Int#
*# Int#
elS)
                 MutableByteArray# s
arrM ((Int#
offM Int# -> Int# -> Int#
+# Int#
off) Int# -> Int# -> Int#
*# Int#
elS) (Int#
lenA Int# -> Int# -> Int#
*# Int#
elS) State# s
s, () #)
{-# INLINE copyMDataFrameOff# #-}

-- | Make a mutable DataFrame immutable, without copying.
unsafeFreezeDataFrame# ::
       forall (t :: Type) (k :: Type) (ns :: [k]) s
     . PrimArray t (DataFrame t ns)
    => MDataFrame s t ns
    -> State# s -> (# State# s, DataFrame t ns #)
unsafeFreezeDataFrame# :: MDataFrame s t ns -> State# s -> (# State# s, DataFrame t ns #)
unsafeFreezeDataFrame# (MDataFrame# Int#
offM CumulDims
steps MutableByteArray# s
arrM) State# s
s0
    | Int#
0# <- CumulDims -> Int#
cdTotalDim# CumulDims
steps
      = (# State# s
s0, [Char] -> DataFrame t ns
forall a. HasCallStack => [Char] -> a
error [Char]
"Empty DataFrame (DF0)" #)
    | (# State# s
s1, ByteArray#
arrA #) <- MutableByteArray# s -> State# s -> (# State# s, ByteArray# #)
forall d.
MutableByteArray# d -> State# d -> (# State# d, ByteArray# #)
unsafeFreezeByteArray# MutableByteArray# s
arrM State# s
s0
      = (# State# s
s1, CumulDims -> Int# -> ByteArray# -> DataFrame t ns
forall t a. PrimArray t a => CumulDims -> Int# -> ByteArray# -> a
fromElems CumulDims
steps Int#
offM ByteArray#
arrA #)
{-# INLINE unsafeFreezeDataFrame# #-}

-- | Copy content of a mutable DataFrame into a new immutable DataFrame.
freezeDataFrame# ::
       forall (t :: Type) (k :: Type) (ns :: [k]) s
     . PrimArray t (DataFrame t ns)
    => MDataFrame s t ns -> State# s -> (# State# s, DataFrame t ns #)
freezeDataFrame# :: MDataFrame s t ns -> State# s -> (# State# s, DataFrame t ns #)
freezeDataFrame# (MDataFrame# Int#
offM CumulDims
steps MutableByteArray# s
arrM) State# s
s0
    | Int#
0# <- CumulDims -> Int#
cdTotalDim# CumulDims
steps
      = (# State# s
s0, [Char] -> DataFrame t ns
forall a. HasCallStack => [Char] -> a
error [Char]
"Empty DataFrame (DF0)" #)
    | Int#
elS  <- t -> Int#
forall a. PrimBytes a => a -> Int#
byteSize @t t
forall a. HasCallStack => a
undefined
    , Int#
n <- CumulDims -> Int#
cdTotalDim# CumulDims
steps
    , (# State# s
s1, MutableByteArray# s
mba #) <- Int# -> State# s -> (# State# s, MutableByteArray# s #)
forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newByteArray# (Int#
n Int# -> Int# -> Int#
*# Int#
elS) State# s
s0
    , State# s
s2 <- MutableByteArray# s
-> Int#
-> MutableByteArray# s
-> Int#
-> Int#
-> State# s
-> State# s
forall d.
MutableByteArray# d
-> Int#
-> MutableByteArray# d
-> Int#
-> Int#
-> State# d
-> State# d
copyMutableByteArray# MutableByteArray# s
arrM (Int#
offM Int# -> Int# -> Int#
*# Int#
elS) MutableByteArray# s
mba Int#
0# (Int#
n Int# -> Int# -> Int#
*# Int#
elS) State# s
s1
    , (# State# s
s3, ByteArray#
arrA #) <- MutableByteArray# s -> State# s -> (# State# s, ByteArray# #)
forall d.
MutableByteArray# d -> State# d -> (# State# d, ByteArray# #)
unsafeFreezeByteArray# MutableByteArray# s
mba State# s
s2
      = (# State# s
s3, CumulDims -> Int# -> ByteArray# -> DataFrame t ns
forall t a. PrimArray t a => CumulDims -> Int# -> ByteArray# -> a
fromElems CumulDims
steps Int#
0# ByteArray#
arrA #)
{-# INLINE freezeDataFrame# #-}

-- | Create a new mutable DataFrame and copy content of immutable one in there.
thawDataFrame# ::
       forall (t :: Type) (k :: Type) (ns :: [k]) s
     . (Dimensions ns, PrimArray t (DataFrame t ns))
    => DataFrame t ns -> State# s -> (# State# s, MDataFrame s t ns #)
thawDataFrame# :: DataFrame t ns -> State# s -> (# State# s, MDataFrame s t ns #)
thawDataFrame# DataFrame t ns
df State# s
s0
    | Word
nw Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
0
      = (# State# s
s0, [Char] -> MDataFrame s t ns
forall a. HasCallStack => [Char] -> a
error [Char]
"Empty DataFrame (DF0)" #)
    | Int#
bSize <- case Word
nw of W# Word#
w -> t -> Int#
forall a. PrimBytes a => a -> Int#
byteSize @t t
forall a. HasCallStack => a
undefined Int# -> Int# -> Int#
*# Word# -> Int#
word2Int# Word#
w
    , (# State# s
s1, MutableByteArray# s
arrM #) <- Int# -> State# s -> (# State# s, MutableByteArray# s #)
forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newByteArray# Int#
bSize State# s
s0
    , MDataFrame s t ns
r <- Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
forall k k s (t :: k) (ns :: [k]).
Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
MDataFrame# Int#
0# CumulDims
steps MutableByteArray# s
arrM
    , (# State# s
s2, ()
_ #) <- Int
-> DataFrame t ns
-> MDataFrame s t ns
-> State# s
-> (# State# s, () #)
forall t k (as :: [k]) (bs :: [k]) (asbs :: [k]) s.
(Dimensions bs, PrimArray t (DataFrame t bs),
 ConcatList as bs asbs) =>
Int
-> DataFrame t bs
-> MDataFrame s t asbs
-> State# s
-> (# State# s, () #)
copyDataFrameOff# Int
0 DataFrame t ns
df MDataFrame s t ns
r State# s
s1
      = (# State# s
s2, MDataFrame s t ns
r #)
  where
    nw :: Word
nw = CumulDims -> Word
cdTotalDim CumulDims
steps
    steps :: CumulDims
steps = Dims ns -> DataFrame t ns -> CumulDims
forall t a k (ns :: [k]).
PrimArray t a =>
Dims ns -> a -> CumulDims
getSteps (Dimensions ns => Dims ns
forall k (ds :: [k]). Dimensions ds => Dims ds
dims @ns) DataFrame t ns
df
{-# INLINE thawDataFrame# #-}

-- | Create a new mutable DataFrame and copy content of immutable one in there.
--   The result array is pinned and aligned.
thawPinDataFrame# ::
       forall (t :: Type) (k :: Type) (ns :: [k]) s
     . (Dimensions ns, PrimArray t (DataFrame t ns))
    => DataFrame t ns -> State# s -> (# State# s, MDataFrame s t ns #)
thawPinDataFrame# :: DataFrame t ns -> State# s -> (# State# s, MDataFrame s t ns #)
thawPinDataFrame# DataFrame t ns
df State# s
s0
    | Word
nw Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
0
      = (# State# s
s0, [Char] -> MDataFrame s t ns
forall a. HasCallStack => [Char] -> a
error [Char]
"Empty DataFrame (DF0)" #)
    | Int#
bSize <- case Word
nw of W# Word#
w -> t -> Int#
forall a. PrimBytes a => a -> Int#
byteSize @t t
forall a. HasCallStack => a
undefined Int# -> Int# -> Int#
*# Word# -> Int#
word2Int# Word#
w
    , (# State# s
s1, MutableByteArray# s
arrM #) <- Int# -> Int# -> State# s -> (# State# s, MutableByteArray# s #)
forall d.
Int# -> Int# -> State# d -> (# State# d, MutableByteArray# d #)
newAlignedPinnedByteArray# Int#
bSize (t -> Int#
forall a. PrimBytes a => a -> Int#
byteAlign @t t
forall a. HasCallStack => a
undefined) State# s
s0
    , MDataFrame s t ns
r <- Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
forall k k s (t :: k) (ns :: [k]).
Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
MDataFrame# Int#
0# CumulDims
steps MutableByteArray# s
arrM
    , (# State# s
s2, ()
_ #) <- Int
-> DataFrame t ns
-> MDataFrame s t ns
-> State# s
-> (# State# s, () #)
forall t k (as :: [k]) (bs :: [k]) (asbs :: [k]) s.
(Dimensions bs, PrimArray t (DataFrame t bs),
 ConcatList as bs asbs) =>
Int
-> DataFrame t bs
-> MDataFrame s t asbs
-> State# s
-> (# State# s, () #)
copyDataFrameOff# Int
0 DataFrame t ns
df MDataFrame s t ns
r State# s
s1
      = (# State# s
s2, MDataFrame s t ns
r #)
  where
    nw :: Word
nw = CumulDims -> Word
cdTotalDim CumulDims
steps
    steps :: CumulDims
steps = Dims ns -> DataFrame t ns -> CumulDims
forall t a k (ns :: [k]).
PrimArray t a =>
Dims ns -> a -> CumulDims
getSteps (Dimensions ns => Dims ns
forall k (ds :: [k]). Dimensions ds => Dims ds
dims @ns) DataFrame t ns
df
{-# INLINE thawPinDataFrame# #-}

-- | UnsafeCoerces an underlying byte array.
unsafeThawDataFrame# ::
       forall (t :: Type) (k :: Type) (ns :: [k]) s
     . (Dimensions ns, PrimArray t (DataFrame t ns))
    => DataFrame t ns
    -> State# s -> (# State# s, MDataFrame s t ns #)
unsafeThawDataFrame# :: DataFrame t ns -> State# s -> (# State# s, MDataFrame s t ns #)
unsafeThawDataFrame# = (t -> State# s -> (# State# s, MDataFrame s t ns #))
-> (CumulDims
    -> Int#
    -> ByteArray#
    -> State# s
    -> (# State# s, MDataFrame s t ns #))
-> DataFrame t ns
-> State# s
-> (# State# s, MDataFrame s t ns #)
forall t a r.
PrimArray t a =>
(t -> r) -> (CumulDims -> Int# -> ByteArray# -> r) -> a -> r
withArrayContent t -> State# s -> (# State# s, MDataFrame s t ns #)
f CumulDims
-> Int#
-> ByteArray#
-> State# s
-> (# State# s, MDataFrame s t ns #)
g
  where
    f :: t -> State# s -> (# State# s, MDataFrame s t ns #)
    f :: t -> State# s -> (# State# s, MDataFrame s t ns #)
f t
e State# s
s0
      | CumulDims
steps <- Dims ns -> CumulDims
forall k (ns :: [k]). Dims ns -> CumulDims
cumulDims (Dimensions ns => Dims ns
forall k (ds :: [k]). Dimensions ds => Dims ds
dims @ns)
      , Int#
n <- CumulDims -> Int#
cdTotalDim# CumulDims
steps
      , (# State# s
s1, MutableByteArray# s
arrM #) <- Int# -> State# s -> (# State# s, MutableByteArray# s #)
forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newByteArray# (Int#
n Int# -> Int# -> Int#
*# t -> Int#
forall a. PrimBytes a => a -> Int#
byteSize @t t
forall a. HasCallStack => a
undefined) State# s
s0
      = (# MutableByteArray# s -> Int# -> Int# -> t -> State# s -> State# s
forall t s.
PrimBytes t =>
MutableByteArray# s -> Int# -> Int# -> t -> State# s -> State# s
fillArray MutableByteArray# s
arrM Int#
0# Int#
n t
e State# s
s1, Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
forall k k s (t :: k) (ns :: [k]).
Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
MDataFrame# Int#
0# CumulDims
steps MutableByteArray# s
arrM #)
    g :: CumulDims -> Int# -> ByteArray# -> State# s -> (# State# s, MDataFrame s t ns #)
    g :: CumulDims
-> Int#
-> ByteArray#
-> State# s
-> (# State# s, MDataFrame s t ns #)
g CumulDims
steps Int#
off ByteArray#
ba State# s
s0
      = (# State# s
s0, Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
forall k k s (t :: k) (ns :: [k]).
Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
MDataFrame# Int#
off CumulDims
steps (ByteArray# -> MutableByteArray# s
unsafeCoerce# ByteArray#
ba) #)
{-# INLINE unsafeThawDataFrame# #-}

-- | Given two continuations @f@ and @g@.
--   If the input DataFrame is a single broadcast value, use it in @f@.
--   Otherwise, create a new mutable DataFrame and copy content of immutable one
--   in there; then use it in @g@.
--
--   This function is useful when @thawDataFrame@ cannot be used due to
--   @Dimensions ns@ constraint being not available.
withThawDataFrame# ::
       forall (t :: Type) (k :: Type) (ns :: [k]) (r :: Type) s
     . PrimArray t (DataFrame t ns)
    => (t -> State# s -> (# State# s, r #)) -- ^ f
    -> (MDataFrame s t ns -> State# s -> (# State# s, r #)) -- ^ g
    -> DataFrame t ns
    -> State# s -> (# State# s, r #)
withThawDataFrame# :: (t -> State# s -> (# State# s, r #))
-> (MDataFrame s t ns -> State# s -> (# State# s, r #))
-> DataFrame t ns
-> State# s
-> (# State# s, r #)
withThawDataFrame# t -> State# s -> (# State# s, r #)
f MDataFrame s t ns -> State# s -> (# State# s, r #)
g = (t -> State# s -> (# State# s, r #))
-> (CumulDims
    -> Int# -> ByteArray# -> State# s -> (# State# s, r #))
-> DataFrame t ns
-> State# s
-> (# State# s, r #)
forall t a r.
PrimArray t a =>
(t -> r) -> (CumulDims -> Int# -> ByteArray# -> r) -> a -> r
withArrayContent t -> State# s -> (# State# s, r #)
f CumulDims -> Int# -> ByteArray# -> State# s -> (# State# s, r #)
g'
  where
    g' :: CumulDims -> Int# -> ByteArray# -> State# s -> (# State# s, r #)
    g' :: CumulDims -> Int# -> ByteArray# -> State# s -> (# State# s, r #)
g' CumulDims
steps Int#
eOff ByteArray#
arrA State# s
s0 = case CumulDims -> Int#
cdTotalDim# CumulDims
steps of
      Int#
0# -> MDataFrame s t ns -> State# s -> (# State# s, r #)
g ([Char] -> MDataFrame s t ns
forall a. HasCallStack => [Char] -> a
error [Char]
"Empty DataFrame (DF0)") State# s
s0
      Int#
elems
       | Int#
elS <- t -> Int#
forall a. PrimBytes a => a -> Int#
byteSize @t t
forall a. HasCallStack => a
undefined
       , Int#
bsize <- Int#
elS Int# -> Int# -> Int#
*# Int#
elems
       , (# State# s
s1, MutableByteArray# s
arrM #) <- Int# -> State# s -> (# State# s, MutableByteArray# s #)
forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newByteArray# Int#
bsize State# s
s0
       , State# s
s2 <- ByteArray#
-> Int#
-> MutableByteArray# s
-> Int#
-> Int#
-> State# s
-> State# s
forall d.
ByteArray#
-> Int#
-> MutableByteArray# d
-> Int#
-> Int#
-> State# d
-> State# d
copyByteArray# ByteArray#
arrA (Int#
eOff Int# -> Int# -> Int#
*# Int#
elS) MutableByteArray# s
arrM Int#
0# Int#
bsize State# s
s1
         -> MDataFrame s t ns -> State# s -> (# State# s, r #)
g (Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
forall k k s (t :: k) (ns :: [k]).
Int# -> CumulDims -> MutableByteArray# s -> MDataFrame s t ns
MDataFrame# Int#
0# CumulDims
steps MutableByteArray# s
arrM) State# s
s2
{-# INLINE withThawDataFrame# #-}

-- | Write a single element at the specified element offset.
--
--   This is a low-level write function; you have to keep in mind the row-major
--   layout of Mutable DataFrames. Offset bounds are not checked.
--   You will get an undefined behavior if you write beyond the DataFrame bounds.
writeDataFrameOff# ::
       forall (t :: Type) (k :: Type) (ns :: [k]) s
     . PrimBytes (DataFrame t ('[] :: [k]))
    => MDataFrame s t ns -> Int -> DataFrame t ('[] :: [k])
    -> State# s -> (# State# s, () #)
writeDataFrameOff# :: MDataFrame s t ns
-> Int -> DataFrame t '[] -> State# s -> (# State# s, () #)
writeDataFrameOff# (MDataFrame# Int#
off CumulDims
_ MutableByteArray# s
mba) (I# Int#
i) DataFrame t '[]
x State# s
s
  = (# MutableByteArray# s
-> Int# -> DataFrame t '[] -> State# s -> State# s
forall a s.
PrimBytes a =>
MutableByteArray# s -> Int# -> a -> State# s -> State# s
writeArray MutableByteArray# s
mba (Int#
off Int# -> Int# -> Int#
+# Int#
i) DataFrame t '[]
x State# s
s, () #)
{-# INLINE writeDataFrameOff# #-}

-- | Write a single element at the specified index.
--
--   This function is safe (no `OutOfDimBounds` exception possible).
--   If any of the dims in @ns@ is unknown (@n ~ XN m@),
--   you may happen to write data beyond dataframe bounds.
--   In this case, this function does nothing.
writeDataFrame# ::
       forall (t :: Type) (k :: Type) (ns :: [k]) s
     . PrimBytes (DataFrame t ('[] :: [k]))
    => MDataFrame s t ns -> Idxs ns -> DataFrame t ('[] :: [k])
    -> State# s -> (# State# s, () #)
writeDataFrame# :: MDataFrame s t ns
-> Idxs ns -> DataFrame t '[] -> State# s -> (# State# s, () #)
writeDataFrame# mdf :: MDataFrame s t ns
mdf@(MDataFrame# Int#
_ CumulDims
st MutableByteArray# s
_) Idxs ns
ei
  | Just Int
off <- (CumulDims -> Idxs ns -> Maybe Int
forall k (ns :: [k]). CumulDims -> Idxs ns -> Maybe Int
cdIxM CumulDims
st Idxs ns
ei)
  = MDataFrame s t ns
-> Int -> DataFrame t '[] -> State# s -> (# State# s, () #)
forall t k (ns :: [k]) s.
PrimBytes (DataFrame t '[]) =>
MDataFrame s t ns
-> Int -> DataFrame t '[] -> State# s -> (# State# s, () #)
writeDataFrameOff# MDataFrame s t ns
mdf Int
off
  | Bool
otherwise = (State# s -> (# State# s, () #))
-> DataFrame t '[] -> State# s -> (# State# s, () #)
forall a b. a -> b -> a
const (\State# s
s -> (# State# s
s, () #))
{-# INLINE writeDataFrame# #-}

-- | Read a single element at the specified element offset.
--
--   This is a low-level read function; you have to keep in mind the row-major
--   layout of Mutable DataFrames. Offset bounds are not checked.
--   You will get an undefined behavior if you read beyond the DataFrame bounds.
readDataFrameOff# ::
       forall (t :: Type) (k :: Type) (ns :: [k]) s
     . PrimBytes (DataFrame t ('[] :: [k]))
    => MDataFrame s t ns -> Int
    -> State# s -> (# State# s, DataFrame t ('[] :: [k]) #)
readDataFrameOff# :: MDataFrame s t ns
-> Int -> State# s -> (# State# s, DataFrame t '[] #)
readDataFrameOff# (MDataFrame# Int#
off CumulDims
_ MutableByteArray# s
mba) (I# Int#
i)
  = MutableByteArray# s
-> Int# -> State# s -> (# State# s, DataFrame t '[] #)
forall a s.
PrimBytes a =>
MutableByteArray# s -> Int# -> State# s -> (# State# s, a #)
readArray @(DataFrame t ('[] :: [k])) MutableByteArray# s
mba (Int#
off Int# -> Int# -> Int#
+# Int#
i)
{-# INLINE readDataFrameOff# #-}

-- | Read a single element at the specified index.
--
--   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.
readDataFrame# ::
       forall (t :: Type) (k :: Type) (ns :: [k]) s
     . PrimBytes (DataFrame t ('[] :: [k]))
    => MDataFrame s t ns -> Idxs ns
    -> State# s -> (# State# s, DataFrame t ('[] :: [k]) #)
readDataFrame# :: MDataFrame s t ns
-> Idxs ns -> State# s -> (# State# s, DataFrame t '[] #)
readDataFrame# mdf :: MDataFrame s t ns
mdf@(MDataFrame# Int#
_ CumulDims
st MutableByteArray# s
_) Idxs ns
ei
  = MDataFrame s t ns
-> Int -> State# s -> (# State# s, DataFrame t '[] #)
forall t k (ns :: [k]) s.
PrimBytes (DataFrame t '[]) =>
MDataFrame s t ns
-> Int -> State# s -> (# State# s, DataFrame t '[] #)
readDataFrameOff# MDataFrame s t ns
mdf (CumulDims -> Idxs ns -> Int
forall k (ns :: [k]). CumulDims -> Idxs ns -> Int
cdIx CumulDims
st Idxs ns
ei)
{-# INLINE readDataFrame# #-}

-- | Allow arbitrary operations on a pointer to the beginning of the data.
--   Only possible with @RealWord@ state (thus, in @IO@) due to semantics of
--   @touch#@ operation that keeps the data from being garbage collected.
withDataFramePtr# ::
       forall (t :: Type) (k :: Type) (ns :: [k]) (r :: Type)
     . PrimBytes t
    => MDataFrame RealWorld t ns
    -> (Addr# -> State# RealWorld -> (# State# RealWorld, r #))
    -> State# RealWorld -> (# State# RealWorld, r #)
withDataFramePtr# :: MDataFrame RealWorld t ns
-> (Addr# -> State# RealWorld -> (# State# RealWorld, r #))
-> State# RealWorld
-> (# State# RealWorld, r #)
withDataFramePtr# (MDataFrame# Int#
off CumulDims
_ MutableByteArray# RealWorld
mba) Addr# -> State# RealWorld -> (# State# RealWorld, r #)
k State# RealWorld
s0
  | (# State# RealWorld
s1, ByteArray#
a #) <- MutableByteArray# RealWorld
-> State# RealWorld -> (# State# RealWorld, ByteArray# #)
forall d.
MutableByteArray# d -> State# d -> (# State# d, ByteArray# #)
unsafeFreezeByteArray# MutableByteArray# RealWorld
mba State# RealWorld
s0
  , (# State# RealWorld
s2, r
r #) <- Addr# -> State# RealWorld -> (# State# RealWorld, r #)
k ( ByteArray# -> Addr#
byteArrayContents# ByteArray#
a
                       Addr# -> Int# -> Addr#
`plusAddr#` (Int#
off Int# -> Int# -> Int#
*# t -> Int#
forall a. PrimBytes a => a -> Int#
byteSize @t t
forall a. HasCallStack => a
undefined)
                     ) State# RealWorld
s1
    = (# MutableByteArray# RealWorld -> State# RealWorld -> State# RealWorld
forall k1. k1 -> State# RealWorld -> State# RealWorld
touch# MutableByteArray# RealWorld
mba State# RealWorld
s2, r
r #)
{-# INLINE withDataFramePtr# #-}

-- | Check if the byte array wrapped by this DataFrame is pinned,
--   which means cannot be relocated by GC.
isDataFramePinned# ::
       forall (t :: Type) (k :: Type) (ns :: [k]) s . MDataFrame s t ns -> Bool
isDataFramePinned# :: MDataFrame s t ns -> Bool
isDataFramePinned# (MDataFrame# Int#
_ CumulDims
_ MutableByteArray# s
mba)
  = Int# -> Bool
isTrue# (MutableByteArray# s -> Int#
forall d. MutableByteArray# d -> Int#
isMutableByteArrayPinned# MutableByteArray# s
mba)

-- | Get cumulative dimensions @ns@ of a @MDataFrame s t ns@
getDataFrameSteps# ::
       forall (t :: Type) (k :: Type) (ns :: [k]) s
     . MDataFrame s t ns -> CumulDims
getDataFrameSteps# :: MDataFrame s t ns -> CumulDims
getDataFrameSteps# (MDataFrame# Int#
_ CumulDims
c MutableByteArray# s
_) = CumulDims
c
{-# INLINE getDataFrameSteps# #-}

-- | Fill a mutable byte array with the same single element
fillArray :: PrimBytes t
          => MutableByteArray# s
          -> Int# -- ^ Offset in elements
          -> Int# -- ^ Number of elements
          -> t
          -> State# s -> State# s
fillArray :: MutableByteArray# s -> Int# -> Int# -> t -> State# s -> State# s
fillArray MutableByteArray# s
mba Int#
off Int#
n t
e
  = let lim :: Int#
lim = Int#
off Int# -> Int# -> Int#
+# Int#
n
        go :: Int# -> State# s -> State# s
go Int#
i State# s
s | Int# -> Bool
isTrue# (Int#
i Int# -> Int# -> Int#
>=# Int#
lim) = State# s
s
               | Bool
otherwise           = Int# -> State# s -> State# s
go (Int#
i Int# -> Int# -> Int#
+# Int#
1#) (MutableByteArray# s -> Int# -> t -> State# s -> State# s
forall a s.
PrimBytes a =>
MutableByteArray# s -> Int# -> a -> State# s -> State# s
writeArray MutableByteArray# s
mba Int#
i t
e State# s
s)
    in  Int# -> State# s -> State# s
go Int#
off
{-# INLINE fillArray #-}