module DataFrame.IO.Parquet.Levels where

import qualified Data.ByteString as BS
import Data.Int
import Data.List
import qualified Data.Text as T

import DataFrame.IO.Parquet.Binary
import DataFrame.IO.Parquet.Encoding
import DataFrame.IO.Parquet.Thrift
import DataFrame.IO.Parquet.Types

readLevelsV1 ::
    Int -> Int -> Int -> BS.ByteString -> ([Int], [Int], BS.ByteString)
readLevelsV1 :: Int -> Int -> Int -> ByteString -> ([Int], [Int], ByteString)
readLevelsV1 Int
n Int
maxDef Int
maxRep ByteString
bs =
    let bwDef :: Int
bwDef = Int -> Int
bitWidthForMaxLevel Int
maxDef
        bwRep :: Int
bwRep = Int -> Int
bitWidthForMaxLevel Int
maxRep

        ([Int]
repLvls, ByteString
afterRep) =
            if Int
bwRep Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                then (Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate Int
n Int
0, ByteString
bs)
                else
                    let repLength :: Word32
repLength = ByteString -> Word32
littleEndianWord32 (Int -> ByteString -> ByteString
BS.take Int
4 ByteString
bs)
                        repData :: ByteString
repData = Int -> ByteString -> ByteString
BS.take (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
repLength) (Int -> ByteString -> ByteString
BS.drop Int
4 ByteString
bs)
                        afterRepData :: ByteString
afterRepData = Int -> ByteString -> ByteString
BS.drop (Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
repLength) ByteString
bs
                        ([Word32]
repVals, ByteString
_) = Int -> Int -> ByteString -> ([Word32], ByteString)
decodeRLEBitPackedHybrid Int
bwRep Int
n ByteString
repData
                     in ((Word32 -> Int) -> [Word32] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral [Word32]
repVals, ByteString
afterRepData)

        ([Int]
defLvls, ByteString
afterDef) =
            if Int
bwDef Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                then (Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate Int
n Int
0, ByteString
afterRep)
                else
                    let defLength :: Word32
defLength = ByteString -> Word32
littleEndianWord32 (Int -> ByteString -> ByteString
BS.take Int
4 ByteString
afterRep)
                        defData :: ByteString
defData = Int -> ByteString -> ByteString
BS.take (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
defLength) (Int -> ByteString -> ByteString
BS.drop Int
4 ByteString
afterRep)
                        afterDefData :: ByteString
afterDefData = Int -> ByteString -> ByteString
BS.drop (Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
defLength) ByteString
afterRep
                        ([Word32]
defVals, ByteString
_) = Int -> Int -> ByteString -> ([Word32], ByteString)
decodeRLEBitPackedHybrid Int
bwDef Int
n ByteString
defData
                     in ((Word32 -> Int) -> [Word32] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral [Word32]
defVals, ByteString
afterDefData)
     in ([Int]
defLvls, [Int]
repLvls, ByteString
afterDef)

readLevelsV2 ::
    Int ->
    Int ->
    Int ->
    Int32 ->
    Int32 ->
    BS.ByteString ->
    ([Int], [Int], BS.ByteString)
readLevelsV2 :: Int
-> Int
-> Int
-> Int32
-> Int32
-> ByteString
-> ([Int], [Int], ByteString)
readLevelsV2 Int
n Int
maxDef Int
maxRep Int32
defLen Int32
repLen 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
        bwDef :: Int
bwDef = Int -> Int
bitWidthForMaxLevel Int
maxDef
        bwRep :: Int
bwRep = Int -> Int
bitWidthForMaxLevel Int
maxRep
        ([Word32]
repLvlsRaw, ByteString
_) =
            if Int
bwRep Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                then (Int -> Word32 -> [Word32]
forall a. Int -> a -> [a]
replicate Int
n Word32
0, ByteString
repBytes)
                else Int -> Int -> ByteString -> ([Word32], ByteString)
decodeRLEBitPackedHybrid Int
bwRep Int
n ByteString
repBytes
        ([Word32]
defLvlsRaw, ByteString
_) =
            if Int
bwDef Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                then (Int -> Word32 -> [Word32]
forall a. Int -> a -> [a]
replicate Int
n Word32
0, ByteString
defBytes)
                else Int -> Int -> ByteString -> ([Word32], ByteString)
decodeRLEBitPackedHybrid Int
bwDef Int
n ByteString
defBytes
     in ((Word32 -> Int) -> [Word32] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral [Word32]
defLvlsRaw, (Word32 -> Int) -> [Word32] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral [Word32]
repLvlsRaw, ByteString
afterDefBytes)

stitchNullable :: Int -> [Int] -> [a] -> [Maybe a]
stitchNullable :: forall a. Int -> [Int] -> [a] -> [Maybe a]
stitchNullable Int
maxDef = [Int] -> [a] -> [Maybe a]
forall {a}. [Int] -> [a] -> [Maybe a]
go
  where
    go :: [Int] -> [a] -> [Maybe a]
go [] [a]
_ = []
    go (Int
d : [Int]
ds) [a]
vs
        | Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
maxDef = case [a]
vs of
            (a
v : [a]
vs') -> a -> Maybe a
forall a. a -> Maybe a
Just a
v Maybe a -> [Maybe a] -> [Maybe a]
forall a. a -> [a] -> [a]
: [Int] -> [a] -> [Maybe a]
go [Int]
ds [a]
vs'
            [] -> [Char] -> [Maybe a]
forall a. HasCallStack => [Char] -> a
error [Char]
"value stream exhausted"
        | Bool
otherwise = Maybe a
forall a. Maybe a
Nothing Maybe a -> [Maybe a] -> [Maybe a]
forall a. a -> [a] -> [a]
: [Int] -> [a] -> [Maybe a]
go [Int]
ds [a]
vs

data SNode = SNode
    { SNode -> [Char]
sName :: String
    , SNode -> RepetitionType
sRep :: RepetitionType
    , SNode -> [SNode]
sChildren :: [SNode]
    }
    deriving (Int -> SNode -> ShowS
[SNode] -> ShowS
SNode -> [Char]
(Int -> SNode -> ShowS)
-> (SNode -> [Char]) -> ([SNode] -> ShowS) -> Show SNode
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SNode -> ShowS
showsPrec :: Int -> SNode -> ShowS
$cshow :: SNode -> [Char]
show :: SNode -> [Char]
$cshowList :: [SNode] -> ShowS
showList :: [SNode] -> ShowS
Show, SNode -> SNode -> Bool
(SNode -> SNode -> Bool) -> (SNode -> SNode -> Bool) -> Eq SNode
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SNode -> SNode -> Bool
== :: SNode -> SNode -> Bool
$c/= :: SNode -> SNode -> Bool
/= :: SNode -> SNode -> Bool
Eq)

parseOne :: [SchemaElement] -> (SNode, [SchemaElement])
parseOne :: [SchemaElement] -> (SNode, [SchemaElement])
parseOne [] = [Char] -> (SNode, [SchemaElement])
forall a. HasCallStack => [Char] -> a
error [Char]
"parseOne: empty schema list"
parseOne (SchemaElement
se : [SchemaElement]
rest) =
    let childCount :: Int
childCount = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (SchemaElement -> Int32
numChildren SchemaElement
se)
        ([SNode]
kids, [SchemaElement]
rest') = Int -> [SchemaElement] -> ([SNode], [SchemaElement])
parseMany Int
childCount [SchemaElement]
rest
     in ( SNode
            { sName :: [Char]
sName = Text -> [Char]
T.unpack (SchemaElement -> Text
elementName SchemaElement
se)
            , sRep :: RepetitionType
sRep = SchemaElement -> RepetitionType
repetitionType SchemaElement
se
            , sChildren :: [SNode]
sChildren = [SNode]
kids
            }
        , [SchemaElement]
rest'
        )

parseMany :: Int -> [SchemaElement] -> ([SNode], [SchemaElement])
parseMany :: Int -> [SchemaElement] -> ([SNode], [SchemaElement])
parseMany Int
0 [SchemaElement]
xs = ([], [SchemaElement]
xs)
parseMany Int
n [SchemaElement]
xs =
    let (SNode
node, [SchemaElement]
xs') = [SchemaElement] -> (SNode, [SchemaElement])
parseOne [SchemaElement]
xs
        ([SNode]
nodes, [SchemaElement]
xs'') = Int -> [SchemaElement] -> ([SNode], [SchemaElement])
parseMany (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [SchemaElement]
xs'
     in (SNode
node SNode -> [SNode] -> [SNode]
forall a. a -> [a] -> [a]
: [SNode]
nodes, [SchemaElement]
xs'')

parseAll :: [SchemaElement] -> [SNode]
parseAll :: [SchemaElement] -> [SNode]
parseAll [] = []
parseAll [SchemaElement]
xs = let (SNode
n, [SchemaElement]
xs') = [SchemaElement] -> (SNode, [SchemaElement])
parseOne [SchemaElement]
xs in SNode
n SNode -> [SNode] -> [SNode]
forall a. a -> [a] -> [a]
: [SchemaElement] -> [SNode]
parseAll [SchemaElement]
xs'

-- | Tag leaf values as Just/Nothing according to maxDef.
pairWithVals :: Int -> [(Int, Int)] -> [a] -> [(Int, Int, Maybe a)]
pairWithVals :: forall a. Int -> [(Int, Int)] -> [a] -> [(Int, Int, Maybe a)]
pairWithVals Int
_ [] [a]
_ = []
pairWithVals Int
maxDef ((Int
r, Int
d) : [(Int, Int)]
rds) [a]
vs
    | Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
maxDef = case [a]
vs of
        (a
v : [a]
vs') -> (Int
r, Int
d, a -> Maybe a
forall a. a -> Maybe a
Just a
v) (Int, Int, Maybe a)
-> [(Int, Int, Maybe a)] -> [(Int, Int, Maybe a)]
forall a. a -> [a] -> [a]
: Int -> [(Int, Int)] -> [a] -> [(Int, Int, Maybe a)]
forall a. Int -> [(Int, Int)] -> [a] -> [(Int, Int, Maybe a)]
pairWithVals Int
maxDef [(Int, Int)]
rds [a]
vs'
        [] -> [Char] -> [(Int, Int, Maybe a)]
forall a. HasCallStack => [Char] -> a
error [Char]
"pairWithVals: value stream exhausted"
    | Bool
otherwise = (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)] -> [a] -> [(Int, Int, Maybe a)]
forall a. Int -> [(Int, Int)] -> [a] -> [(Int, Int, Maybe a)]
pairWithVals Int
maxDef [(Int, Int)]
rds [a]
vs

-- | Split triplets into groups; 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

{- | Reconstruct a list column from Dremel encoding levels.
rep=0 starts a new top-level row; def=0 means the entire list slot is null.
Returns one Maybe [Maybe a] per row.
-}
stitchList :: Int -> [Int] -> [Int] -> [a] -> [Maybe [Maybe a]]
stitchList :: forall a. Int -> [Int] -> [Int] -> [a] -> [Maybe [Maybe a]]
stitchList Int
maxDef [Int]
repLvls [Int]
defLvls [a]
vals =
    let triplets :: [(Int, Int, Maybe a)]
triplets = Int -> [(Int, Int)] -> [a] -> [(Int, Int, Maybe a)]
forall a. Int -> [(Int, Int)] -> [a] -> [(Int, Int, Maybe a)]
pairWithVals Int
maxDef ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
repLvls [Int]
defLvls) [a]
vals
        rows :: [[(Int, Int, Maybe a)]]
rows = 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
     in ([(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, Maybe a)]]
rows
  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]

{- | Reconstruct a 2-level nested list (maxRep=2) from Dremel triplets.
defT1: def threshold at which the depth-1 element is present (not null).
maxDef: def threshold at which the leaf is present.
-}
stitchList2 :: Int -> Int -> [Int] -> [Int] -> [a] -> [Maybe [Maybe [Maybe a]]]
stitchList2 :: forall a.
Int -> Int -> [Int] -> [Int] -> [a] -> [Maybe [Maybe [Maybe a]]]
stitchList2 Int
defT1 Int
maxDef [Int]
repLvls [Int]
defLvls [a]
vals =
    let triplets :: [(Int, Int, Maybe a)]
triplets = Int -> [(Int, Int)] -> [a] -> [(Int, Int, Maybe a)]
forall a. Int -> [(Int, Int)] -> [a] -> [(Int, Int, Maybe a)]
pairWithVals Int
maxDef ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
repLvls [Int]
defLvls) [a]
vals
     in ([(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
    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

{- | Reconstruct a 3-level nested list (maxRep=3) from Dremel triplets.
defT1, defT2: def thresholds at which depth-1 and depth-2 elements are present.
maxDef: def threshold at which the leaf is present.
-}
stitchList3 ::
    Int -> Int -> Int -> [Int] -> [Int] -> [a] -> [Maybe [Maybe [Maybe [Maybe a]]]]
stitchList3 :: forall a.
Int
-> Int
-> Int
-> [Int]
-> [Int]
-> [a]
-> [Maybe [Maybe [Maybe [Maybe a]]]]
stitchList3 Int
defT1 Int
defT2 Int
maxDef [Int]
repLvls [Int]
defLvls [a]
vals =
    let triplets :: [(Int, Int, Maybe a)]
triplets = Int -> [(Int, Int)] -> [a] -> [(Int, Int, Maybe a)]
forall a. Int -> [(Int, Int)] -> [a] -> [(Int, Int, Maybe a)]
pairWithVals Int
maxDef ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
repLvls [Int]
defLvls) [a]
vals
     in ([(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
    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

levelsForPath :: [SchemaElement] -> [String] -> (Int, Int)
levelsForPath :: [SchemaElement] -> [[Char]] -> (Int, Int)
levelsForPath [SchemaElement]
schemaTail = Int -> Int -> [SNode] -> [[Char]] -> (Int, Int)
forall {t} {t}.
(Num t, Num t) =>
t -> t -> [SNode] -> [[Char]] -> (t, t)
go Int
0 Int
0 ([SchemaElement] -> [SNode]
parseAll [SchemaElement]
schemaTail)
  where
    go :: t -> t -> [SNode] -> [[Char]] -> (t, t)
go t
defC t
repC [SNode]
_ [] = (t
defC, t
repC)
    go t
defC t
repC [SNode]
nodes ([Char]
p : [[Char]]
ps) =
        case (SNode -> Bool) -> [SNode] -> Maybe SNode
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (\SNode
n -> SNode -> [Char]
sName SNode
n [Char] -> [Char] -> Bool
forall a. Eq a => a -> a -> Bool
== [Char]
p) [SNode]
nodes of
            Maybe SNode
Nothing -> (t
defC, t
repC)
            Just SNode
n ->
                let defC' :: t
defC' = t
defC t -> t -> t
forall a. Num a => a -> a -> a
+ (if SNode -> RepetitionType
sRep SNode
n RepetitionType -> RepetitionType -> Bool
forall a. Eq a => a -> a -> Bool
== RepetitionType
OPTIONAL Bool -> Bool -> Bool
|| SNode -> RepetitionType
sRep SNode
n RepetitionType -> RepetitionType -> Bool
forall a. Eq a => a -> a -> Bool
== RepetitionType
REPEATED then t
1 else t
0)
                    repC' :: t
repC' = t
repC t -> t -> t
forall a. Num a => a -> a -> a
+ (if SNode -> RepetitionType
sRep SNode
n RepetitionType -> RepetitionType -> Bool
forall a. Eq a => a -> a -> Bool
== RepetitionType
REPEATED then t
1 else t
0)
                 in t -> t -> [SNode] -> [[Char]] -> (t, t)
go t
defC' t
repC' (SNode -> [SNode]
sChildren SNode
n) [[Char]]
ps