{-# 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.
-}
module DataFrame.Internal.PackedText (
    PackedTextData (..),
    mkPackedContiguous,
    packedGather,
    packedTake,
    packedRowOffsetVec,
    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.Ord (comparing)
import Data.Text.Internal (Text (Text))
import DataFrame.Internal.Utf8 (isValidUtf8Slice, lenientDecodeSlice)

{- | 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.
-}
data PackedTextData = PackedTextData
    { PackedTextData -> Array
ptBytes :: {-# UNPACK #-} !A.Array
    , PackedTextData -> Vector Int
ptOffsets :: {-# UNPACK #-} !(VU.Vector Int)
    , PackedTextData -> Maybe (Vector Int)
ptSel :: !(Maybe (VU.Vector Int))
    }

-- | 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 -> Vector Int -> Maybe (Vector Int) -> PackedTextData
PackedTextData Array
arr Vector Int
offs Maybe (Vector Int)
forall a. Maybe a
Nothing
{-# INLINE mkPackedContiguous #-}

{- | 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.
-}
packedGather :: VU.Vector Int -> PackedTextData -> PackedTextData
packedGather :: Vector Int -> PackedTextData -> PackedTextData
packedGather Vector Int
indices (PackedTextData Array
arr Vector Int
offs Maybe (Vector Int)
msel) =
    let !base :: Int
base = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
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
        sel' :: Vector Int
sel' = case Maybe (Vector Int)
msel of
            Maybe (Vector Int)
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
            Just Vector Int
s ->
                (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
< Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
s then Int -> Int
clamp (Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
s Int
i) else -Int
1)
                    Vector Int
indices
     in Array -> Vector Int -> Maybe (Vector Int) -> PackedTextData
PackedTextData Array
arr Vector Int
offs (Vector Int -> Maybe (Vector Int)
forall a. a -> Maybe a
Just Vector Int
sel')
{-# 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 Vector Int
offs Maybe (Vector Int)
msel) =
    let !base :: Int
base = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
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
     in case Maybe (Vector Int)
msel of
            Just Vector Int
s -> Array -> Vector Int -> Maybe (Vector Int) -> PackedTextData
PackedTextData Array
arr Vector Int
offs (Vector Int -> Maybe (Vector Int)
forall a. a -> Maybe a
Just (Int -> Vector Int -> Vector Int
forall a. Unbox a => Int -> Vector a -> Vector a
VU.take Int
k' Vector Int
s))
            Maybe (Vector Int)
Nothing -> Array -> Vector Int -> Maybe (Vector Int) -> PackedTextData
PackedTextData Array
arr Vector Int
offs (Vector Int -> Maybe (Vector Int)
forall a. a -> Maybe a
Just (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)))
{-# 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
_ Vector Int
_ Maybe (Vector Int)
Nothing) Int
i = Int
i
baseRow (PackedTextData Array
_ Vector Int
_ (Just Vector Int
sel)) Int
i = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
sel Int
i
{-# INLINE baseRow #-}

-- | Row count: @length sel@ when selected, else @length offsets - 1@.
packedLength :: PackedTextData -> Int
packedLength :: PackedTextData -> Int
packedLength (PackedTextData Array
_ Vector Int
offs Maybe (Vector Int)
Nothing) = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
offs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
packedLength (PackedTextData Array
_ Vector Int
_ (Just Vector Int
sel)) = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
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 Vector Int
offs Maybe (Vector Int)
_) 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 = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
offs Int
r in (Array
arr, Int
o, Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
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 the boxed-Text fallback take the fast contiguous path.
-}
packedRowOffsetVec :: PackedTextData -> Maybe (A.Array, VU.Vector Int)
packedRowOffsetVec :: PackedTextData -> Maybe (Array, Vector Int)
packedRowOffsetVec (PackedTextData Array
arr Vector Int
offs Maybe (Vector Int)
Nothing) = (Array, Vector Int) -> Maybe (Array, Vector Int)
forall a. a -> Maybe a
Just (Array
arr, Vector Int
offs)
packedRowOffsetVec PackedTextData
_ = Maybe (Array, Vector Int)
forall a. Maybe a
Nothing
{-# INLINE packedRowOffsetVec #-}

{- | On-demand single 'Data.Text.Text' for row @i@, using the same
validate-or-lenient decode as 'sliceTextVector' 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 'sliceTextVector' 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 #-}