{-# LANGUAGE BangPatterns #-}

{- | Packed-text payload + byte-slice primitives. A 'PackedTextData' shares one
UTF-8 byte buffer across all rows of a string column, with @n+1@ row offsets, so
no per-row 'Data.Text.Text' header is materialized until decode is demanded.
Offsets and selection vectors are stored 'Int32' whenever their values fit
(Arrow-style), halving the per-row footprint of large string columns.
-}
module DataFrame.Internal.PackedText (
    PackedTextData (..),
    PackedOffsets (..),
    PackedSel (..),
    offAt,
    offCount,
    selAt,
    selLength,
    mkPackedContiguous,
    mkPackedContiguous32,
    mkOffsets,
    mkSel,
    packedGather,
    packedTake,
    packedRowOffsets,
    packedLength,
    packedSlice,
    packedIndexText,
    sliceEqBytes,
    sliceCmpBytes,
) where

import qualified Data.Text as T
import qualified Data.Text.Array as A
import qualified Data.Vector.Unboxed as VU

import Data.Int (Int32)
import Data.Ord (comparing)
import Data.Text.Internal (Text (Text))
import DataFrame.Internal.Utf8 (isValidUtf8Slice, lenientDecodeSlice)

{- | Row byte-offsets, physically 'Int32' when every value fits (total buffer
bytes < 2^31) and 'Int' otherwise. Values are non-negative byte positions.
-}
data PackedOffsets
    = Offs32 {-# UNPACK #-} !(VU.Vector Int32)
    | Offs64 {-# UNPACK #-} !(VU.Vector Int)

-- | Offset at index @i@, widened to 'Int'.
offAt :: PackedOffsets -> Int -> Int
offAt :: PackedOffsets -> Int -> Int
offAt (Offs32 Vector Int32
v) Int
i = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector Int32 -> Int -> Int32
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int32
v Int
i)
offAt (Offs64 Vector Int
v) Int
i = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
v Int
i
{-# INLINE offAt #-}

-- | Number of offset entries (row count + 1).
offCount :: PackedOffsets -> Int
offCount :: PackedOffsets -> Int
offCount (Offs32 Vector Int32
v) = Vector Int32 -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int32
v
offCount (Offs64 Vector Int
v) = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
v
{-# INLINE offCount #-}

{- | A selection layer mapping logical rows to base rows; @-1@ marks an
invalid/null row. 'Int32' when the base row count fits.
-}
data PackedSel
    = Sel32 {-# UNPACK #-} !(VU.Vector Int32)
    | Sel64 {-# UNPACK #-} !(VU.Vector Int)

-- | Base row for logical row @i@ (may be @-1@).
selAt :: PackedSel -> Int -> Int
selAt :: PackedSel -> Int -> Int
selAt (Sel32 Vector Int32
v) Int
i = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector Int32 -> Int -> Int32
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int32
v Int
i)
selAt (Sel64 Vector Int
v) Int
i = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
v Int
i
{-# INLINE selAt #-}

selLength :: PackedSel -> Int
selLength :: PackedSel -> Int
selLength (Sel32 Vector Int32
v) = Vector Int32 -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int32
v
selLength (Sel64 Vector Int
v) = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
v
{-# INLINE selLength #-}

{- | A shared UTF-8 byte buffer plus @n+1@ row offsets (base row @r@ spans bytes
@[offsets!r, offsets!(r+1))@); validity lives in the column's bitmap. @ptSel@ is
an optional selection layer letting a gather/join/sort result share the buffer.

@ptCanonicalSel@ marks a selection that is a canonical dictionary encoding:
equal byte slices always map to the same base row (codes). Set by dictionary
compaction; preserved by gather/take over an already-canonical selection (a
row keeps its code); 'False' for a gather over an unselected base, where two
logical rows can select different but equal-byted base rows. Grouping keys on
codes directly when it holds.
-}
data PackedTextData = PackedTextData
    { PackedTextData -> Array
ptBytes :: {-# UNPACK #-} !A.Array
    , PackedTextData -> PackedOffsets
ptOffsets :: !PackedOffsets
    , PackedTextData -> Maybe PackedSel
ptSel :: !(Maybe PackedSel)
    , PackedTextData -> Bool
ptCanonicalSel :: !Bool
    }

int32Max :: Int
int32Max :: Int
int32Max = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32
forall a. Bounded a => a
maxBound :: Int32)

-- | Narrow an 'Int' offset vector when the final offset (total bytes) fits.
mkOffsets :: VU.Vector Int -> PackedOffsets
mkOffsets :: Vector Int -> PackedOffsets
mkOffsets Vector Int
offs
    | Bool -> Bool
not (Vector Int -> Bool
forall a. Unbox a => Vector a -> Bool
VU.null Vector Int
offs) Bool -> Bool -> Bool
&& Vector Int -> Int
forall a. Unbox a => Vector a -> a
VU.last Vector Int
offs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
int32Max =
        Vector Int32 -> PackedOffsets
Offs32 ((Int -> Int32) -> Vector Int -> Vector Int32
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Vector Int
offs)
    | Bool
otherwise = Vector Int -> PackedOffsets
Offs64 Vector Int
offs
{-# INLINE mkOffsets #-}

{- | Narrow an 'Int' base-row vector (@-1@ sentinels allowed) when the base
row count fits in 'Int32'.
-}
mkSel :: Int -> VU.Vector Int -> PackedSel
mkSel :: Int -> Vector Int -> PackedSel
mkSel Int
base Vector Int
rows
    | Int
base Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
int32Max = Vector Int32 -> PackedSel
Sel32 ((Int -> Int32) -> Vector Int -> Vector Int32
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Vector Int
rows)
    | Bool
otherwise = Vector Int -> PackedSel
Sel64 Vector Int
rows
{-# INLINE mkSel #-}

-- | Build a contiguous packed payload (no selection): the freeze-path shape.
mkPackedContiguous :: A.Array -> VU.Vector Int -> PackedTextData
mkPackedContiguous :: Array -> Vector Int -> PackedTextData
mkPackedContiguous Array
arr Vector Int
offs = Array -> PackedOffsets -> Maybe PackedSel -> Bool -> PackedTextData
PackedTextData Array
arr (Vector Int -> PackedOffsets
mkOffsets Vector Int
offs) Maybe PackedSel
forall a. Maybe a
Nothing Bool
False
{-# INLINE mkPackedContiguous #-}

-- | 'mkPackedContiguous' from offsets already produced at 'Int32' width.
mkPackedContiguous32 :: A.Array -> VU.Vector Int32 -> PackedTextData
mkPackedContiguous32 :: Array -> Vector Int32 -> PackedTextData
mkPackedContiguous32 Array
arr Vector Int32
offs = Array -> PackedOffsets -> Maybe PackedSel -> Bool -> PackedTextData
PackedTextData Array
arr (Vector Int32 -> PackedOffsets
Offs32 Vector Int32
offs) Maybe PackedSel
forall a. Maybe a
Nothing Bool
False
{-# INLINE mkPackedContiguous32 #-}

{- | Reindex a packed payload by a selection vector, sharing the byte buffer;
logical row @i@ becomes base row @indices!i@. A negative or out-of-range index
decodes to the empty slice. Composes with an existing selection; canonicality
survives composition (a kept row keeps its code) but not a first selection
over the unselected base.
-}
packedGather :: VU.Vector Int -> PackedTextData -> PackedTextData
packedGather :: Vector Int -> PackedTextData -> PackedTextData
packedGather Vector Int
indices (PackedTextData Array
arr PackedOffsets
offs Maybe PackedSel
msel Bool
canon) =
    let !base :: Int
base = PackedOffsets -> Int
offCount PackedOffsets
offs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        clamp :: Int -> Int
clamp Int
r = if Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
base then Int
r else -Int
1
        (Vector Int
sel', Bool
canon') = case Maybe PackedSel
msel of
            Maybe PackedSel
Nothing -> ((Int -> Int) -> Vector Int -> Vector Int
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map Int -> Int
clamp Vector Int
indices, Bool
False)
            Just PackedSel
s ->
                let !sn :: Int
sn = PackedSel -> Int
selLength PackedSel
s
                 in ( (Int -> Int) -> Vector Int -> Vector Int
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map
                        (\Int
i -> if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sn then Int -> Int
clamp (PackedSel -> Int -> Int
selAt PackedSel
s Int
i) else -Int
1)
                        Vector Int
indices
                    , Bool
canon
                    )
     in Array -> PackedOffsets -> Maybe PackedSel -> Bool -> PackedTextData
PackedTextData Array
arr PackedOffsets
offs (PackedSel -> Maybe PackedSel
forall a. a -> Maybe a
Just (Int -> Vector Int -> PackedSel
mkSel Int
base Vector Int
sel')) Bool
canon'
{-# INLINE packedGather #-}

{- | Take the first @k@ logical rows, sharing the byte buffer via a capped
selection layer. O(k), no byte copy or decode — cheap @take@/display on a
large packed column.
-}
packedTake :: Int -> PackedTextData -> PackedTextData
packedTake :: Int -> PackedTextData -> PackedTextData
packedTake Int
k (PackedTextData Array
arr PackedOffsets
offs Maybe PackedSel
msel Bool
canon) =
    let !base :: Int
base = PackedOffsets -> Int
offCount PackedOffsets
offs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        !k' :: Int
k' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 Int
k
        (PackedSel
sel', Bool
canon') = case Maybe PackedSel
msel of
            Just (Sel32 Vector Int32
s) -> (Vector Int32 -> PackedSel
Sel32 (Int -> Vector Int32 -> Vector Int32
forall a. Unbox a => Int -> Vector a -> Vector a
VU.take Int
k' Vector Int32
s), Bool
canon)
            Just (Sel64 Vector Int
s) -> (Vector Int -> PackedSel
Sel64 (Int -> Vector Int -> Vector Int
forall a. Unbox a => Int -> Vector a -> Vector a
VU.take Int
k' Vector Int
s), Bool
canon)
            Maybe PackedSel
Nothing -> (Int -> Vector Int -> PackedSel
mkSel Int
base (Int -> Int -> Vector Int
forall a. (Unbox a, Num a) => a -> Int -> Vector a
VU.enumFromN Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
k' Int
base)), Bool
False)
     in Array -> PackedOffsets -> Maybe PackedSel -> Bool -> PackedTextData
PackedTextData Array
arr PackedOffsets
offs (PackedSel -> Maybe PackedSel
forall a. a -> Maybe a
Just PackedSel
sel') Bool
canon'
{-# INLINE packedTake #-}

-- | Map a logical row index to its base row, honoring any selection layer.
baseRow :: PackedTextData -> Int -> Int
baseRow :: PackedTextData -> Int -> Int
baseRow (PackedTextData Array
_ PackedOffsets
_ Maybe PackedSel
Nothing Bool
_) Int
i = Int
i
baseRow (PackedTextData Array
_ PackedOffsets
_ (Just PackedSel
sel) Bool
_) Int
i = PackedSel -> Int -> Int
selAt PackedSel
sel Int
i
{-# INLINE baseRow #-}

-- | Row count: @length sel@ when selected, else @length offsets - 1@.
packedLength :: PackedTextData -> Int
packedLength :: PackedTextData -> Int
packedLength (PackedTextData Array
_ PackedOffsets
offs Maybe PackedSel
Nothing Bool
_) = PackedOffsets -> Int
offCount PackedOffsets
offs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
packedLength (PackedTextData Array
_ PackedOffsets
_ (Just PackedSel
sel) Bool
_) = PackedSel -> Int
selLength PackedSel
sel
{-# INLINE packedLength #-}

-- | Raw byte slice for logical row @i@: @(buffer, offset, length)@. The hot accessor.
packedSlice :: PackedTextData -> Int -> (A.Array, Int, Int)
packedSlice :: PackedTextData -> Int -> (Array, Int, Int)
packedSlice p :: PackedTextData
p@(PackedTextData Array
arr PackedOffsets
offs Maybe PackedSel
_ Bool
_) Int
i =
    let !r :: Int
r = PackedTextData -> Int -> Int
baseRow PackedTextData
p Int
i
     in if Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
            then (Array
arr, Int
0, Int
0)
            else
                let o :: Int
o = PackedOffsets -> Int -> Int
offAt PackedOffsets
offs Int
r in (Array
arr, Int
o, PackedOffsets -> Int -> Int
offAt PackedOffsets
offs (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
o)
{-# INLINE packedSlice #-}

{- | The shared buffer + contiguous @n+1@ offsets when the payload is the
unselected base; a selected (gathered) payload returns 'Nothing' (its rows are
non-contiguous). Lets contiguous consumers skip the selection indirection.
-}
packedRowOffsets :: PackedTextData -> Maybe (A.Array, PackedOffsets)
packedRowOffsets :: PackedTextData -> Maybe (Array, PackedOffsets)
packedRowOffsets (PackedTextData Array
arr PackedOffsets
offs Maybe PackedSel
Nothing Bool
_) = (Array, PackedOffsets) -> Maybe (Array, PackedOffsets)
forall a. a -> Maybe a
Just (Array
arr, PackedOffsets
offs)
packedRowOffsets PackedTextData
_ = Maybe (Array, PackedOffsets)
forall a. Maybe a
Nothing
{-# INLINE packedRowOffsets #-}

{- | On-demand single 'Data.Text.Text' for row @i@, using the same
validate-or-lenient decode as the freeze path so output is bit-identical.
-}
packedIndexText :: PackedTextData -> Int -> T.Text
packedIndexText :: PackedTextData -> Int -> Text
packedIndexText PackedTextData
p Int
i =
    let (Array
arr, Int
o, Int
l) = PackedTextData -> Int -> (Array, Int, Int)
packedSlice PackedTextData
p Int
i
     in Array -> Int -> Int -> Text
decodeField Array
arr Int
o Int
l
{-# INLINE packedIndexText #-}

-- Decode one field exactly as the boxed freeze path does per row.
decodeField :: A.Array -> Int -> Int -> T.Text
decodeField :: Array -> Int -> Int -> Text
decodeField Array
arr Int
o Int
l
    | Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Text
T.empty
    | Array -> Int -> Int -> Bool
isValidUtf8Slice Array
arr Int
o Int
l = Array -> Int -> Int -> Text
Text Array
arr Int
o Int
l
    | Bool
otherwise = Array -> Int -> Int -> Text
lenientDecodeSlice Array
arr Int
o Int
l
{-# INLINE decodeField #-}

{- | Byte-wise equality of two slices. UTF-8 is injective on valid scalar
sequences and lenient decode is deterministic, so this agrees with
@Text@'s '==' on the decoded values.
-}
sliceEqBytes :: A.Array -> Int -> Int -> A.Array -> Int -> Int -> Bool
sliceEqBytes :: Array -> Int -> Int -> Array -> Int -> Int -> Bool
sliceEqBytes Array
a Int
ao Int
al Array
b Int
bo Int
bl
    | Int
al Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
bl = Bool
False
    | Bool
otherwise = Int -> Bool
go Int
0
  where
    go :: Int -> Bool
go !Int
k
        | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
al = Bool
True
        | Array -> Int -> Word8
A.unsafeIndex Array
a (Int
ao Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Array -> Int -> Word8
A.unsafeIndex Array
b (Int
bo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) = Int -> Bool
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise = Bool
False
{-# INLINE sliceEqBytes #-}

{- | Unsigned byte-lexicographic comparison (memcmp semantics). For
well-formed UTF-8 this matches 'Data.Text.compare' exactly, since UTF-8
byte order equals codepoint order for all valid scalars.
-}
sliceCmpBytes :: A.Array -> Int -> Int -> A.Array -> Int -> Int -> Ordering
sliceCmpBytes :: Array -> Int -> Int -> Array -> Int -> Int -> Ordering
sliceCmpBytes Array
a Int
ao Int
al Array
b Int
bo Int
bl = Int -> Ordering
go Int
0
  where
    !m :: Int
m = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
al Int
bl
    go :: Int -> Ordering
go !Int
k
        | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
m = Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
al Int
bl
        | Bool
otherwise = case (Word8 -> Word8) -> Word8 -> Word8 -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing Word8 -> Word8
forall a. a -> a
id (Array -> Int -> Word8
A.unsafeIndex Array
a (Int
ao Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k)) (Array -> Int -> Word8
A.unsafeIndex Array
b (Int
bo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k)) of
            Ordering
EQ -> Int -> Ordering
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            Ordering
r -> Ordering
r
{-# INLINE sliceCmpBytes #-}