{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

{- | Concatenation of per-chunk 'Column's (e.g. from parallel CSV chunks). Text
columns merge at the byte level via 'TextChunk' \/ 'mergeTextChunks', so no
per-chunk 'Data.Text.Text' values are ever materialized.
-}
module DataFrame.Internal.ColumnMerge (
    TextChunk (..),
    mergeColumns,
    mergeTextChunks,
    packedFromTextChunk,
    packValidity,
    spliceBitmaps,
    tcRows,
) where

import qualified Data.Text.Array as A
import qualified Data.Vector as VB
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM

import Control.Monad (foldM_, forM_, when)
import Control.Monad.ST (ST, runST)
import Data.Bits (shiftL, shiftR, (.&.), (.|.))
import Data.Maybe (fromMaybe, isNothing)
import Data.Type.Equality (testEquality, (:~:) (Refl))
import Data.Word (Word8)
import DataFrame.Internal.Column (
    Bitmap,
    Column (..),
    Columnable,
    allValidBitmap,
    isMergedColumn,
    isPackedText,
    materializeMerged,
    materializePacked,
 )
import DataFrame.Internal.PackedText (mkPackedContiguous)
import Type.Reflection (typeRep)

{- | A frozen text-builder chunk: raw UTF-8 bytes plus row offsets (row @i@
spans bytes @[offsets!i, offsets!(i+1))@) and an optional validity bitmap.
'Data.Text.Text' values are only created when chunks merge into a 'Column'.
-}
data TextChunk = TextChunk
    { TextChunk -> Array
tcBytes :: !A.Array
    , TextChunk -> Int
tcUsed :: !Int
    , TextChunk -> Vector Int
tcOffsets :: !(VU.Vector Int)
    , TextChunk -> Maybe Bitmap
tcBitmap :: !(Maybe Bitmap)
    }

tcRows :: TextChunk -> Int
tcRows :: TextChunk -> Int
tcRows TextChunk
c = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length (TextChunk -> Vector Int
tcOffsets TextChunk
c) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

{- | Freeze a builder chunk directly into a packed-text column: no
'Data.Text.Text' materialization, no UTF-8 validation pass (deferred to decode).
Not yet called by any reader.
-}
packedFromTextChunk :: TextChunk -> Column
packedFromTextChunk :: TextChunk -> Column
packedFromTextChunk (TextChunk Array
arr Int
_used Vector Int
offs Maybe Bitmap
bm) =
    Maybe Bitmap -> PackedTextData -> Column
PackedText Maybe Bitmap
bm (Array -> Vector Int -> PackedTextData
mkPackedContiguous Array
arr Vector Int
offs)

{- | Merge text chunks into one packed-text 'Column': one byte-array copy per
chunk, one offset rebase, then wrap the shared buffer + offsets as 'PackedText'
(no per-row header, decode deferred).
-}
mergeTextChunks :: [TextChunk] -> Column
mergeTextChunks :: [TextChunk] -> Column
mergeTextChunks [] = [Char] -> Column
forall a. HasCallStack => [Char] -> a
error [Char]
"DataFrame.Internal.ColumnMerge.mergeTextChunks: empty list"
mergeTextChunks [TextChunk
c] = TextChunk -> Column
packedFromTextChunk TextChunk
c
mergeTextChunks [TextChunk]
cs = (forall s. ST s Column) -> Column
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s Column) -> Column)
-> (forall s. ST s Column) -> Column
forall a b. (a -> b) -> a -> b
$ do
    let totalBytes :: Int
totalBytes = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((TextChunk -> Int) -> [TextChunk] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map TextChunk -> Int
tcUsed [TextChunk]
cs)
        totalRows :: Int
totalRows = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((TextChunk -> Int) -> [TextChunk] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map TextChunk -> Int
tcRows [TextChunk]
cs)
    MArray s
arr <- Int -> ST s (MArray s)
forall s. Int -> ST s (MArray s)
A.new (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
totalBytes)
    MVector s Int
offs <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Int
totalRows Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
offs Int
0 Int
0
    let splice :: Int -> Int -> TextChunk -> ST s (Int, Int)
splice !Int
byteBase !Int
rowBase TextChunk
c = do
            let n :: Int
n = TextChunk -> Int
tcRows TextChunk
c
                co :: Vector Int
co = TextChunk -> Vector Int
tcOffsets TextChunk
c
            Int -> MArray s -> Int -> Array -> Int -> ST s ()
forall s. Int -> MArray s -> Int -> Array -> Int -> ST s ()
A.copyI (TextChunk -> Int
tcUsed TextChunk
c) MArray s
arr Int
byteBase (TextChunk -> Array
tcBytes TextChunk
c) Int
0
            [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
1 .. Int
n] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
offs (Int
rowBase Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) (Int
byteBase Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
co Int
i)
            (Int, Int) -> ST s (Int, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
byteBase Int -> Int -> Int
forall a. Num a => a -> a -> a
+ TextChunk -> Int
tcUsed TextChunk
c, Int
rowBase Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)
    ((Int, Int) -> TextChunk -> ST s (Int, Int))
-> (Int, Int) -> [TextChunk] -> ST s ()
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
foldM_ (\(Int
b, Int
r) TextChunk
c -> Int -> Int -> TextChunk -> ST s (Int, Int)
splice Int
b Int
r TextChunk
c) (Int
0, Int
0) [TextChunk]
cs
    Array
farr <- MArray s -> ST s Array
forall s. MArray s -> ST s Array
A.unsafeFreeze MArray s
arr
    Vector Int
foffs <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
offs
    let !bm :: Maybe Bitmap
bm = [(Maybe Bitmap, Int)] -> Maybe Bitmap
spliceBitmaps [(TextChunk -> Maybe Bitmap
tcBitmap TextChunk
c, TextChunk -> Int
tcRows TextChunk
c) | TextChunk
c <- [TextChunk]
cs]
    Column -> ST s Column
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Bitmap -> PackedTextData -> Column
PackedText Maybe Bitmap
bm (Array -> Vector Int -> PackedTextData
mkPackedContiguous Array
farr Vector Int
foffs))

{- | Merge per-chunk columns into one column: one allocation + memcpy per
payload, with bitmaps spliced across non-byte-aligned chunk boundaries.
All chunks must have the same element type.
-}
mergeColumns :: [Column] -> Column
mergeColumns :: [Column] -> Column
mergeColumns [] = [Char] -> Column
forall a. HasCallStack => [Char] -> a
error [Char]
"DataFrame.Internal.ColumnBuilder.mergeColumns: empty list"
mergeColumns [Column
c] = Column
c
-- Normalize on the whole list, not the head: a packed or merged chunk in any
-- position must demote every chunk to the common boxed form.
mergeColumns cols :: [Column]
cols@(Column
c0 : [Column]
_)
    | (Column -> Bool) -> [Column] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Column -> Bool
isMergedColumn [Column]
cols = [Column] -> Column
mergeColumns ((Column -> Column) -> [Column] -> [Column]
forall a b. (a -> b) -> [a] -> [b]
map Column -> Column
materializeMerged [Column]
cols)
    | (Column -> Bool) -> [Column] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Column -> Bool
isPackedText [Column]
cols = [Column] -> Column
mergeColumns ((Column -> Column) -> [Column] -> [Column]
forall a b. (a -> b) -> [a] -> [b]
map Column -> Column
materializePacked [Column]
cols)
mergeColumns cols :: [Column]
cols@(Column
c0 : [Column]
_) = case Column
c0 of
    PackedText Maybe Bitmap
_ PackedTextData
_ -> [Column] -> Column
mergeColumns ((Column -> Column) -> [Column] -> [Column]
forall a b. (a -> b) -> [a] -> [b]
map Column -> Column
materializePacked [Column]
cols)
    MergedColumn Column
_ Column
_ -> [Column] -> Column
mergeColumns ((Column -> Column) -> [Column] -> [Column]
forall a b. (a -> b) -> [a] -> [b]
map Column -> Column
materializeMerged [Column]
cols)
    UnboxedColumn Maybe Bitmap
_ (Vector a
_ :: VU.Vector a) ->
        let parts :: [(Maybe Bitmap, Vector a)]
parts = (Column -> (Maybe Bitmap, Vector a))
-> [Column] -> [(Maybe Bitmap, Vector a)]
forall a b. (a -> b) -> [a] -> [b]
map (forall a.
(Columnable a, Unbox a) =>
Column -> (Maybe Bitmap, Vector a)
unboxedPart @a) [Column]
cols
            !merged :: Vector a
merged = [Vector a] -> Vector a
forall a. Unbox a => [Vector a] -> Vector a
VU.concat (((Maybe Bitmap, Vector a) -> Vector a)
-> [(Maybe Bitmap, Vector a)] -> [Vector a]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe Bitmap, Vector a) -> Vector a
forall a b. (a, b) -> b
snd [(Maybe Bitmap, Vector a)]
parts)
            !bm :: Maybe Bitmap
bm = [(Maybe Bitmap, Int)] -> Maybe Bitmap
spliceBitmaps [(Maybe Bitmap
mb, Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
v) | (Maybe Bitmap
mb, Vector a
v) <- [(Maybe Bitmap, Vector a)]
parts]
         in Maybe Bitmap -> Vector a -> Column
forall a.
(Columnable a, Unbox a) =>
Maybe Bitmap -> Vector a -> Column
UnboxedColumn Maybe Bitmap
bm Vector a
merged
    BoxedColumn Maybe Bitmap
_ (Vector a
_ :: VB.Vector a) ->
        let parts :: [(Maybe Bitmap, Vector a)]
parts = (Column -> (Maybe Bitmap, Vector a))
-> [Column] -> [(Maybe Bitmap, Vector a)]
forall a b. (a -> b) -> [a] -> [b]
map (forall a. Columnable a => Column -> (Maybe Bitmap, Vector a)
boxedPart @a) [Column]
cols
            !merged :: Vector a
merged = [Vector a] -> Vector a
forall a. [Vector a] -> Vector a
VB.concat (((Maybe Bitmap, Vector a) -> Vector a)
-> [(Maybe Bitmap, Vector a)] -> [Vector a]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe Bitmap, Vector a) -> Vector a
forall a b. (a, b) -> b
snd [(Maybe Bitmap, Vector a)]
parts)
            !bm :: Maybe Bitmap
bm = [(Maybe Bitmap, Int)] -> Maybe Bitmap
spliceBitmaps [(Maybe Bitmap
mb, Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
v) | (Maybe Bitmap
mb, Vector a
v) <- [(Maybe Bitmap, Vector a)]
parts]
         in Maybe Bitmap -> Vector a -> Column
forall a. Columnable a => Maybe Bitmap -> Vector a -> Column
BoxedColumn Maybe Bitmap
bm Vector a
merged

unboxedPart ::
    forall a. (Columnable a, VU.Unbox a) => Column -> (Maybe Bitmap, VU.Vector a)
unboxedPart :: forall a.
(Columnable a, Unbox a) =>
Column -> (Maybe Bitmap, Vector a)
unboxedPart (UnboxedColumn Maybe Bitmap
mb (Vector a
v :: VU.Vector b)) =
    case TypeRep a -> TypeRep a -> Maybe (a :~: a)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b) of
        Just a :~: a
Refl -> (Maybe Bitmap
mb, Vector a
Vector a
v)
        Maybe (a :~: a)
Nothing -> (Maybe Bitmap, Vector a)
forall a. a
mergeMismatch
unboxedPart Column
_ = (Maybe Bitmap, Vector a)
forall a. a
mergeMismatch

boxedPart ::
    forall a. (Columnable a) => Column -> (Maybe Bitmap, VB.Vector a)
boxedPart :: forall a. Columnable a => Column -> (Maybe Bitmap, Vector a)
boxedPart (BoxedColumn Maybe Bitmap
mb (Vector a
v :: VB.Vector b)) =
    case TypeRep a -> TypeRep a -> Maybe (a :~: a)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b) of
        Just a :~: a
Refl -> (Maybe Bitmap
mb, Vector a
Vector a
v)
        Maybe (a :~: a)
Nothing -> (Maybe Bitmap, Vector a)
forall a. a
mergeMismatch
boxedPart Column
_ = (Maybe Bitmap, Vector a)
forall a. a
mergeMismatch

mergeMismatch :: a
mergeMismatch :: forall a. a
mergeMismatch =
    [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"DataFrame.Internal.ColumnBuilder.mergeColumns: chunk column types differ"

{- | Splice chunk bitmaps end to end at the bit level. 'Nothing' if no chunk
carries a bitmap; chunks without one count as all-valid otherwise.
-}
spliceBitmaps :: [(Maybe Bitmap, Int)] -> Maybe Bitmap
spliceBitmaps :: [(Maybe Bitmap, Int)] -> Maybe Bitmap
spliceBitmaps [(Maybe Bitmap, Int)]
parts
    | ((Maybe Bitmap, Int) -> Bool) -> [(Maybe Bitmap, Int)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Maybe Bitmap -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe Bitmap -> Bool)
-> ((Maybe Bitmap, Int) -> Maybe Bitmap)
-> (Maybe Bitmap, Int)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe Bitmap, Int) -> Maybe Bitmap
forall a b. (a, b) -> a
fst) [(Maybe Bitmap, Int)]
parts = Maybe Bitmap
forall a. Maybe a
Nothing
    | Bool
otherwise = Bitmap -> Maybe Bitmap
forall a. a -> Maybe a
Just (Bitmap -> Maybe Bitmap) -> Bitmap -> Maybe Bitmap
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (MVector s Word8)) -> Bitmap
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Word8)) -> Bitmap)
-> (forall s. ST s (MVector s Word8)) -> Bitmap
forall a b. (a -> b) -> a -> b
$ do
        let total :: Int
total = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (((Maybe Bitmap, Int) -> Int) -> [(Maybe Bitmap, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe Bitmap, Int) -> Int
forall a b. (a, b) -> b
snd [(Maybe Bitmap, Int)]
parts)
            outBytes :: Int
outBytes = (Int
total Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3
        MVector s Word8
mv <- Int -> Word8 -> ST s (MVector (PrimState (ST s)) Word8)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
outBytes Word8
0
        let orInto :: Int -> Word8 -> ST s ()
orInto Int
i Word8
w =
                Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
outBytes Bool -> Bool -> Bool
&& Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
                    Word8
old <- MVector (PrimState (ST s)) Word8 -> Int -> ST s Word8
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector s Word8
MVector (PrimState (ST s)) Word8
mv Int
i
                    MVector (PrimState (ST s)) Word8 -> Int -> Word8 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Word8
MVector (PrimState (ST s)) Word8
mv Int
i (Word8
old Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8
w)
            splice :: Int -> (Maybe Bitmap, Int) -> ST s Int
splice !Int
bitPos (Maybe Bitmap
mb, Int
len) = do
                let bm :: Bitmap
bm = Bitmap -> Maybe Bitmap -> Bitmap
forall a. a -> Maybe a -> a
fromMaybe (Int -> Bitmap
allValidBitmap Int
len) Maybe Bitmap
mb
                    sh :: Int
sh = Int
bitPos Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
7
                    byte0 :: Int
byte0 = Int
bitPos Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3
                    lastIdx :: Int
lastIdx = ((Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
                    tailBits :: Int
tailBits = Int
len Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
7
                    lastMask :: Word8
lastMask =
                        if Int
tailBits Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Word8
0xFF else (Word8
1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
tailBits) Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
- Word8
1
                [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
lastIdx] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
k -> do
                    let raw :: Word8
raw = Bitmap -> Int -> Word8
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Bitmap
bm Int
k
                        masked :: Word8
masked = if Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
lastIdx then Word8
raw Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
lastMask else Word8
raw
                        w :: Word
w = Word8 -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
masked :: Word
                    Int -> Word8 -> ST s ()
orInto (Int
byte0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) (Word -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word
w Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`shiftL` Int
sh))
                    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
sh Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
                        Int -> Word8 -> ST s ()
orInto (Int
byte0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Word -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word
w Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`shiftR` (Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
sh)))
                Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
bitPos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len)
        (Int -> (Maybe Bitmap, Int) -> ST s Int)
-> Int -> [(Maybe Bitmap, Int)] -> ST s ()
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
foldM_ Int -> (Maybe Bitmap, Int) -> ST s Int
splice Int
0 [(Maybe Bitmap, Int)]
parts
        MVector s Word8 -> ST s (MVector s Word8)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Word8
mv

-- | Pack a 0\/1 byte-per-row validity prefix into a bit-packed 'Bitmap'.
packValidity :: Int -> VUM.MVector s Word8 -> ST s Bitmap
packValidity :: forall s. Int -> MVector s Word8 -> ST s Bitmap
packValidity Int
n MVector s Word8
val = do
    Bitmap
bytes <- MVector (PrimState (ST s)) Word8 -> ST s Bitmap
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze (Int -> Int -> MVector s Word8 -> MVector s Word8
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
VUM.slice Int
0 Int
n MVector s Word8
val)
    let assemble :: Int -> Word8
assemble Int
b =
            let base :: Int
base = Int
b Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
3
                m :: Int
m = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
8 (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
base)
                go :: Word8 -> Int -> Word8
go !Word8
acc !Int
k
                    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
m = Word8
acc
                    | Bitmap -> Int -> Word8
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Bitmap
bytes (Int
base Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0 =
                        Word8 -> Int -> Word8
go (Word8
acc Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. (Word8
1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
k)) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                    | Bool
otherwise = Word8 -> Int -> Word8
go Word8
acc (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
             in Word8 -> Int -> Word8
go (Word8
0 :: Word8) Int
0
    Bitmap -> ST s Bitmap
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bitmap -> ST s Bitmap) -> Bitmap -> ST s Bitmap
forall a b. (a -> b) -> a -> b
$! Int -> (Int -> Word8) -> Bitmap
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate ((Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3) Int -> Word8
assemble