{-# LANGUAGE BangPatterns #-}

module DataFrame.IO.Parquet.Levels (
    -- Level readers
    readLevelsV1,
    readLevelsV2,
    -- Stitch functions
    stitchList,
    stitchList2,
    stitchList3,
) where

import Control.Monad.ST (runST)
import qualified Data.ByteString as BS
import Data.Int (Int32)
import qualified Data.Vector as VB
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import Data.Word (Word32)
import DataFrame.IO.Parquet.Encoding (
    bitWidthForMaxLevel,
    decodeRLEBitPackedHybrid,
 )
import DataFrame.Internal.Binary (littleEndianWord32)

-- ---------------------------------------------------------------------------
-- Level readers
-- ---------------------------------------------------------------------------

{- | Convert a 'Word32' level vector to an 'Int' level vector while counting
how many entries equal @maxDef@. Single pass; allocates a single
'VU.Vector Int' of length @VU.length raw@.
-}
convertAndCount :: Int -> VU.Vector Word32 -> (VU.Vector Int, Int)
convertAndCount :: Int -> Vector Word32 -> (Vector Int, Int)
convertAndCount Int
maxDef Vector Word32
raw = (forall s. ST s (Vector Int, Int)) -> (Vector Int, Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int, Int)) -> (Vector Int, Int))
-> (forall s. ST s (Vector Int, Int)) -> (Vector Int, Int)
forall a b. (a -> b) -> a -> b
$ do
    let !n :: Int
n = Vector Word32 -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Word32
raw
    MVector s Int
mv <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
n
    let !maxDefW :: Word32
maxDefW = Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
maxDef :: Word32
        go :: Int -> t -> f t
go !Int
i !t
nPresent
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = t -> f t
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure t
nPresent
            | Bool
otherwise = do
                let !w :: Word32
w = Vector Word32 -> Int -> Word32
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Word32
raw Int
i
                    !d :: Int
d = Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
w :: Int
                MVector (PrimState f) Int -> Int -> Int -> f ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState f) Int
mv Int
i Int
d
                if Word32
w Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
== Word32
maxDefW
                    then Int -> t -> f t
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (t
nPresent t -> t -> t
forall a. Num a => a -> a -> a
+ t
1)
                    else Int -> t -> f t
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) t
nPresent
    !Int
nPresent <- Int -> Int -> ST s Int
forall {f :: * -> *} {t}.
(PrimState f ~ s, PrimMonad f, Num t) =>
Int -> t -> f t
go Int
0 Int
0
    !Vector Int
out <- 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
mv
    (Vector Int, Int) -> ST s (Vector Int, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector Int
out, Int
nPresent)

readLevelsV1 ::
    -- | Total number of values in the page
    Int ->
    -- | maxDefinitionLevel
    Int ->
    -- | maxRepetitionLevel
    Int ->
    BS.ByteString ->
    (VU.Vector Int, VU.Vector Int, Int, BS.ByteString)
readLevelsV1 :: Int
-> Int
-> Int
-> ByteString
-> (Vector Int, Vector Int, Int, ByteString)
readLevelsV1 Int
n Int
maxDef Int
maxRep ByteString
bs =
    let bwRep :: Int
bwRep = Int -> Int
bitWidthForMaxLevel Int
maxRep
        bwDef :: Int
bwDef = Int -> Int
bitWidthForMaxLevel Int
maxDef
        (Vector Int
repVec, Int
_, ByteString
afterRep) = Int -> Int -> ByteString -> (Vector Int, Int, ByteString)
decodeLevelBlock Int
bwRep Int
n ByteString
bs
        (Vector Int
defVec, Int
nPresent, ByteString
afterDef) = Int -> Int -> ByteString -> (Vector Int, Int, ByteString)
decodeLevelBlock Int
bwDef Int
n ByteString
afterRep
     in (Vector Int
defVec, Vector Int
repVec, Int
nPresent, ByteString
afterDef)
  where
    -- For rep block we don't need nPresent; we still get one cheaply.
    decodeLevelBlock :: Int -> Int -> ByteString -> (Vector Int, Int, ByteString)
decodeLevelBlock Int
0 Int
n' ByteString
buf = (Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
n' Int
0, Int
n' Int -> Int -> Int
forall a. Num a => a -> a -> a
* Bool -> Int
forall a. Enum a => a -> Int
fromEnum (Int
maxDef Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0), ByteString
buf)
    decodeLevelBlock Int
bw Int
n' ByteString
buf =
        let blockLen :: Int
blockLen = Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Word32
littleEndianWord32 (Int -> ByteString -> ByteString
BS.take Int
4 ByteString
buf)) :: Int
            blockData :: ByteString
blockData = Int -> ByteString -> ByteString
BS.take Int
blockLen (Int -> ByteString -> ByteString
BS.drop Int
4 ByteString
buf)
            after :: ByteString
after = Int -> ByteString -> ByteString
BS.drop (Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
blockLen) ByteString
buf
            (Vector Word32
raw, ByteString
_) = Int -> Int -> ByteString -> (Vector Word32, ByteString)
decodeRLEBitPackedHybrid Int
bw Int
n' ByteString
blockData
            (Vector Int
out, Int
np) = Int -> Vector Word32 -> (Vector Int, Int)
convertAndCount Int
maxDef Vector Word32
raw
         in (Vector Int
out, Int
np, ByteString
after)

readLevelsV2 ::
    -- | Total number of values
    Int ->
    -- | maxDefinitionLevel
    Int ->
    -- | maxRepetitionLevel
    Int ->
    -- | Repetition-level byte length (from page header)
    Int32 ->
    -- | Definition-level byte length (from page header)
    Int32 ->
    BS.ByteString ->
    (VU.Vector Int, VU.Vector Int, Int, BS.ByteString)
readLevelsV2 :: Int
-> Int
-> Int
-> Int32
-> Int32
-> ByteString
-> (Vector Int, Vector Int, Int, ByteString)
readLevelsV2 Int
n Int
maxDef Int
maxRep Int32
repLen Int32
defLen ByteString
bs =
    let (ByteString
repBytes, ByteString
afterRepBytes) = Int -> ByteString -> (ByteString, ByteString)
BS.splitAt (Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
repLen) ByteString
bs
        (ByteString
defBytes, ByteString
afterDefBytes) = Int -> ByteString -> (ByteString, ByteString)
BS.splitAt (Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
defLen) ByteString
afterRepBytes
        bwRep :: Int
bwRep = Int -> Int
bitWidthForMaxLevel Int
maxRep
        bwDef :: Int
bwDef = Int -> Int
bitWidthForMaxLevel Int
maxDef
        repVec :: Vector Int
repVec
            | Int
bwRep Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
n Int
0
            | Bool
otherwise =
                let (Vector Word32
raw, ByteString
_) = Int -> Int -> ByteString -> (Vector Word32, ByteString)
decodeRLEBitPackedHybrid Int
bwRep Int
n ByteString
repBytes
                    (Vector Int
out, Int
_) = Int -> Vector Word32 -> (Vector Int, Int)
convertAndCount Int
maxDef Vector Word32
raw
                 in Vector Int
out
        (Vector Int
defVec, Int
nPresent)
            | Int
bwDef Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = (Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
VU.replicate Int
n Int
0, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Bool -> Int
forall a. Enum a => a -> Int
fromEnum (Int
maxDef Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0))
            | Bool
otherwise =
                let (Vector Word32
raw, ByteString
_) = Int -> Int -> ByteString -> (Vector Word32, ByteString)
decodeRLEBitPackedHybrid Int
bwDef Int
n ByteString
defBytes
                 in Int -> Vector Word32 -> (Vector Int, Int)
convertAndCount Int
maxDef Vector Word32
raw
     in (Vector Int
defVec, Vector Int
repVec, Int
nPresent, ByteString
afterDefBytes)

{- | Stitch a singly-nested list column (@maxRep == 1@) from vector-format
definition and repetition levels plus a compact present-values vector.
Returns one @Maybe [Maybe a]@ per top-level row.
-}
stitchList ::
    Int ->
    VU.Vector Int ->
    VU.Vector Int ->
    VB.Vector a ->
    [Maybe [Maybe a]]
stitchList :: forall a.
Int -> Vector Int -> Vector Int -> Vector a -> [Maybe [Maybe a]]
stitchList Int
maxDef Vector Int
repVec Vector Int
defVec Vector a
values =
    ([(Int, Int, Maybe a)] -> Maybe [Maybe a])
-> [[(Int, Int, Maybe a)]] -> [Maybe [Maybe a]]
forall a b. (a -> b) -> [a] -> [b]
map [(Int, Int, Maybe a)] -> Maybe [Maybe a]
forall {a} {a} {a}. (Eq a, Num a) => [(a, a, a)] -> Maybe [a]
toRow (Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
forall a. Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
splitAtRepBound Int
0 (Int
-> Vector Int -> Vector Int -> Vector a -> [(Int, Int, Maybe a)]
forall a.
Int
-> Vector Int -> Vector Int -> Vector a -> [(Int, Int, Maybe a)]
pairWithValsV Int
maxDef Vector Int
repVec Vector Int
defVec Vector a
values))
  where
    toRow :: [(a, a, a)] -> Maybe [a]
toRow [] = Maybe [a]
forall a. Maybe a
Nothing
    toRow ((a
_, a
d, a
_) : [(a, a, a)]
_) | a
d a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 = Maybe [a]
forall a. Maybe a
Nothing
    toRow [(a, a, a)]
grp = [a] -> Maybe [a]
forall a. a -> Maybe a
Just [a
v | (a
_, a
_, a
v) <- [(a, a, a)]
grp]

{- | Stitch a doubly-nested list column (@maxRep == 2@).
@defT1@ is the def threshold at which the depth-1 element is present.
-}
stitchList2 ::
    Int ->
    Int ->
    VU.Vector Int ->
    VU.Vector Int ->
    VB.Vector a ->
    [Maybe [Maybe [Maybe a]]]
stitchList2 :: forall a.
Int
-> Int
-> Vector Int
-> Vector Int
-> Vector a
-> [Maybe [Maybe [Maybe a]]]
stitchList2 Int
defT1 Int
maxDef Vector Int
repVec Vector Int
defVec Vector a
values =
    ([(Int, Int, Maybe a)] -> Maybe [Maybe [Maybe a]])
-> [[(Int, Int, Maybe a)]] -> [Maybe [Maybe [Maybe a]]]
forall a b. (a -> b) -> [a] -> [b]
map [(Int, Int, Maybe a)] -> Maybe [Maybe [Maybe a]]
forall {a}. [(Int, Int, Maybe a)] -> Maybe [Maybe [Maybe a]]
toRow (Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
forall a. Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
splitAtRepBound Int
0 [(Int, Int, Maybe a)]
triplets)
  where
    triplets :: [(Int, Int, Maybe a)]
triplets = Int
-> Vector Int -> Vector Int -> Vector a -> [(Int, Int, Maybe a)]
forall a.
Int
-> Vector Int -> Vector Int -> Vector a -> [(Int, Int, Maybe a)]
pairWithValsV Int
maxDef Vector Int
repVec Vector Int
defVec Vector a
values
    toRow :: [(Int, Int, Maybe a)] -> Maybe [Maybe [Maybe a]]
toRow [] = Maybe [Maybe [Maybe a]]
forall a. Maybe a
Nothing
    toRow ((Int
_, Int
d, Maybe a
_) : [(Int, Int, Maybe a)]
_) | Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Maybe [Maybe [Maybe a]]
forall a. Maybe a
Nothing
    toRow [(Int, Int, Maybe a)]
row = [Maybe [Maybe a]] -> Maybe [Maybe [Maybe a]]
forall a. a -> Maybe a
Just (([(Int, Int, Maybe a)] -> Maybe [Maybe a])
-> [[(Int, Int, Maybe a)]] -> [Maybe [Maybe a]]
forall a b. (a -> b) -> [a] -> [b]
map [(Int, Int, Maybe a)] -> Maybe [Maybe a]
forall {a}. [(Int, Int, Maybe a)] -> Maybe [Maybe a]
toOuter (Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
forall a. Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
splitAtRepBound Int
1 [(Int, Int, Maybe a)]
row))
    toOuter :: [(Int, Int, Maybe a)] -> Maybe [Maybe a]
toOuter [] = Maybe [Maybe a]
forall a. Maybe a
Nothing
    toOuter ((Int
_, Int
d, Maybe a
_) : [(Int, Int, Maybe a)]
_) | Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
defT1 = Maybe [Maybe a]
forall a. Maybe a
Nothing
    toOuter [(Int, Int, Maybe a)]
outer = [Maybe a] -> Maybe [Maybe a]
forall a. a -> Maybe a
Just (([(Int, Int, Maybe a)] -> Maybe a)
-> [[(Int, Int, Maybe a)]] -> [Maybe a]
forall a b. (a -> b) -> [a] -> [b]
map [(Int, Int, Maybe a)] -> Maybe a
forall {a} {b} {a}. [(a, b, Maybe a)] -> Maybe a
toLeaf (Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
forall a. Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
splitAtRepBound Int
2 [(Int, Int, Maybe a)]
outer))
    toLeaf :: [(a, b, Maybe a)] -> Maybe a
toLeaf [] = Maybe a
forall a. Maybe a
Nothing
    toLeaf ((a
_, b
_, Maybe a
v) : [(a, b, Maybe a)]
_) = Maybe a
v

{- | Stitch a triply-nested list column (@maxRep == 3@).
@defT1@ and @defT2@ are the def thresholds for depth-1 and depth-2
elements respectively.
-}
stitchList3 ::
    Int ->
    Int ->
    Int ->
    VU.Vector Int ->
    VU.Vector Int ->
    VB.Vector a ->
    [Maybe [Maybe [Maybe [Maybe a]]]]
stitchList3 :: forall a.
Int
-> Int
-> Int
-> Vector Int
-> Vector Int
-> Vector a
-> [Maybe [Maybe [Maybe [Maybe a]]]]
stitchList3 Int
defT1 Int
defT2 Int
maxDef Vector Int
repVec Vector Int
defVec Vector a
values =
    ([(Int, Int, Maybe a)] -> Maybe [Maybe [Maybe [Maybe a]]])
-> [[(Int, Int, Maybe a)]] -> [Maybe [Maybe [Maybe [Maybe a]]]]
forall a b. (a -> b) -> [a] -> [b]
map [(Int, Int, Maybe a)] -> Maybe [Maybe [Maybe [Maybe a]]]
forall {a}.
[(Int, Int, Maybe a)] -> Maybe [Maybe [Maybe [Maybe a]]]
toRow (Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
forall a. Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
splitAtRepBound Int
0 [(Int, Int, Maybe a)]
triplets)
  where
    triplets :: [(Int, Int, Maybe a)]
triplets = Int
-> Vector Int -> Vector Int -> Vector a -> [(Int, Int, Maybe a)]
forall a.
Int
-> Vector Int -> Vector Int -> Vector a -> [(Int, Int, Maybe a)]
pairWithValsV Int
maxDef Vector Int
repVec Vector Int
defVec Vector a
values
    toRow :: [(Int, Int, Maybe a)] -> Maybe [Maybe [Maybe [Maybe a]]]
toRow [] = Maybe [Maybe [Maybe [Maybe a]]]
forall a. Maybe a
Nothing
    toRow ((Int
_, Int
d, Maybe a
_) : [(Int, Int, Maybe a)]
_) | Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Maybe [Maybe [Maybe [Maybe a]]]
forall a. Maybe a
Nothing
    toRow [(Int, Int, Maybe a)]
row = [Maybe [Maybe [Maybe a]]] -> Maybe [Maybe [Maybe [Maybe a]]]
forall a. a -> Maybe a
Just (([(Int, Int, Maybe a)] -> Maybe [Maybe [Maybe a]])
-> [[(Int, Int, Maybe a)]] -> [Maybe [Maybe [Maybe a]]]
forall a b. (a -> b) -> [a] -> [b]
map [(Int, Int, Maybe a)] -> Maybe [Maybe [Maybe a]]
forall {a}. [(Int, Int, Maybe a)] -> Maybe [Maybe [Maybe a]]
toOuter (Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
forall a. Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
splitAtRepBound Int
1 [(Int, Int, Maybe a)]
row))
    toOuter :: [(Int, Int, Maybe a)] -> Maybe [Maybe [Maybe a]]
toOuter [] = Maybe [Maybe [Maybe a]]
forall a. Maybe a
Nothing
    toOuter ((Int
_, Int
d, Maybe a
_) : [(Int, Int, Maybe a)]
_) | Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
defT1 = Maybe [Maybe [Maybe a]]
forall a. Maybe a
Nothing
    toOuter [(Int, Int, Maybe a)]
outer = [Maybe [Maybe a]] -> Maybe [Maybe [Maybe a]]
forall a. a -> Maybe a
Just (([(Int, Int, Maybe a)] -> Maybe [Maybe a])
-> [[(Int, Int, Maybe a)]] -> [Maybe [Maybe a]]
forall a b. (a -> b) -> [a] -> [b]
map [(Int, Int, Maybe a)] -> Maybe [Maybe a]
forall {a}. [(Int, Int, Maybe a)] -> Maybe [Maybe a]
toMiddle (Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
forall a. Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
splitAtRepBound Int
2 [(Int, Int, Maybe a)]
outer))
    toMiddle :: [(Int, Int, Maybe a)] -> Maybe [Maybe a]
toMiddle [] = Maybe [Maybe a]
forall a. Maybe a
Nothing
    toMiddle ((Int
_, Int
d, Maybe a
_) : [(Int, Int, Maybe a)]
_) | Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
defT2 = Maybe [Maybe a]
forall a. Maybe a
Nothing
    toMiddle [(Int, Int, Maybe a)]
middle = [Maybe a] -> Maybe [Maybe a]
forall a. a -> Maybe a
Just (([(Int, Int, Maybe a)] -> Maybe a)
-> [[(Int, Int, Maybe a)]] -> [Maybe a]
forall a b. (a -> b) -> [a] -> [b]
map [(Int, Int, Maybe a)] -> Maybe a
forall {a} {b} {a}. [(a, b, Maybe a)] -> Maybe a
toLeaf (Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
forall a. Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
splitAtRepBound Int
3 [(Int, Int, Maybe a)]
middle))
    toLeaf :: [(a, b, Maybe a)] -> Maybe a
toLeaf [] = Maybe a
forall a. Maybe a
Nothing
    toLeaf ((a
_, b
_, Maybe a
v) : [(a, b, Maybe a)]
_) = Maybe a
v

-- ---------------------------------------------------------------------------
-- Internal helpers
-- ---------------------------------------------------------------------------

{- | Zip rep and def level vectors with a present-values vector, tagging each
position as @Just value@ (when @def == maxDef@) or @Nothing@.
Returns a flat list of @(rep, def, Maybe a)@ triplets for row-splitting.
-}
pairWithValsV ::
    Int ->
    VU.Vector Int ->
    VU.Vector Int ->
    VB.Vector a ->
    [(Int, Int, Maybe a)]
pairWithValsV :: forall a.
Int
-> Vector Int -> Vector Int -> Vector a -> [(Int, Int, Maybe a)]
pairWithValsV Int
maxDef Vector Int
repVec Vector Int
defVec Vector a
values = Int -> Int -> [(Int, Int, Maybe a)]
go Int
0 Int
0
  where
    n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
defVec
    go :: Int -> Int -> [(Int, Int, Maybe a)]
go Int
i Int
j
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = []
        | Bool
otherwise =
            let r :: Int
r = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
repVec Int
i
                d :: Int
d = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
defVec Int
i
             in if Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
maxDef
                    then (Int
r, Int
d, a -> Maybe a
forall a. a -> Maybe a
Just (Vector a -> Int -> a
forall a. Vector a -> Int -> a
VB.unsafeIndex Vector a
values Int
j)) (Int, Int, Maybe a)
-> [(Int, Int, Maybe a)] -> [(Int, Int, Maybe a)]
forall a. a -> [a] -> [a]
: Int -> Int -> [(Int, Int, Maybe a)]
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                    else (Int
r, Int
d, Maybe a
forall a. Maybe a
Nothing) (Int, Int, Maybe a)
-> [(Int, Int, Maybe a)] -> [(Int, Int, Maybe a)]
forall a. a -> [a] -> [a]
: Int -> Int -> [(Int, Int, Maybe a)]
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
j

{- | Group a flat triplet list into rows.
A new group begins whenever @rep <= bound@.
-}
splitAtRepBound :: Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
splitAtRepBound :: forall a. Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
splitAtRepBound Int
_ [] = []
splitAtRepBound Int
bound ((Int, Int, Maybe a)
t : [(Int, Int, Maybe a)]
ts) =
    let ([(Int, Int, Maybe a)]
rest, [(Int, Int, Maybe a)]
remaining) = ((Int, Int, Maybe a) -> Bool)
-> [(Int, Int, Maybe a)]
-> ([(Int, Int, Maybe a)], [(Int, Int, Maybe a)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span (\(Int
r, Int
_, Maybe a
_) -> Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
bound) [(Int, Int, Maybe a)]
ts
     in ((Int, Int, Maybe a)
t (Int, Int, Maybe a)
-> [(Int, Int, Maybe a)] -> [(Int, Int, Maybe a)]
forall a. a -> [a] -> [a]
: [(Int, Int, Maybe a)]
rest) [(Int, Int, Maybe a)]
-> [[(Int, Int, Maybe a)]] -> [[(Int, Int, Maybe a)]]
forall a. a -> [a] -> [a]
: Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
forall a. Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]]
splitAtRepBound Int
bound [(Int, Int, Maybe a)]
remaining