module Network.HPACK.Huffman.Tree (
    -- * Huffman decoding
    HTree (..),
    eosInfo,
    toHTree,
    showTree,
    printTree,
    flatten,
) where

import Control.Arrow (second)

import Imports
import Network.HPACK.Huffman.Bit
import Network.HPACK.Huffman.Params

----------------------------------------------------------------

type EOSInfo = Maybe Int

-- | Type for Huffman decoding.
data HTree
    = Tip
        EOSInfo -- EOS info from 1
        {-# UNPACK #-} Int -- Decoded value. Essentially Word8
    | Bin
        EOSInfo -- EOS info from 1
        {-# UNPACK #-} Int -- Sequence no from 0
        HTree -- Left
        HTree -- Right
    deriving (Int -> HTree -> ShowS
[HTree] -> ShowS
HTree -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [HTree] -> ShowS
$cshowList :: [HTree] -> ShowS
show :: HTree -> String
$cshow :: HTree -> String
showsPrec :: Int -> HTree -> ShowS
$cshowsPrec :: Int -> HTree -> ShowS
Show)

eosInfo :: HTree -> EOSInfo
eosInfo :: HTree -> EOSInfo
eosInfo (Tip EOSInfo
mx Int
_) = EOSInfo
mx
eosInfo (Bin EOSInfo
mx Int
_ HTree
_ HTree
_) = EOSInfo
mx

----------------------------------------------------------------

showTree :: HTree -> String
showTree :: HTree -> String
showTree = String -> HTree -> String
showTree' String
""

showTree' :: String -> HTree -> String
showTree' :: String -> HTree -> String
showTree' String
_ (Tip EOSInfo
_ Int
i) = forall a. Show a => a -> String
show Int
i forall a. [a] -> [a] -> [a]
++ String
"\n"
showTree' String
pref (Bin EOSInfo
_ Int
n HTree
l HTree
r) =
    String
"No "
        forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
n
        forall a. [a] -> [a] -> [a]
++ String
"\n"
        forall a. [a] -> [a] -> [a]
++ String
pref
        forall a. [a] -> [a] -> [a]
++ String
"+ "
        forall a. [a] -> [a] -> [a]
++ String -> HTree -> String
showTree' String
pref' HTree
l
        forall a. [a] -> [a] -> [a]
++ String
pref
        forall a. [a] -> [a] -> [a]
++ String
"+ "
        forall a. [a] -> [a] -> [a]
++ String -> HTree -> String
showTree' String
pref' HTree
r
  where
    pref' :: String
pref' = String
"  " forall a. [a] -> [a] -> [a]
++ String
pref

printTree :: HTree -> IO ()
printTree :: HTree -> IO ()
printTree = String -> IO ()
putStr forall b c a. (b -> c) -> (a -> b) -> a -> c
. HTree -> String
showTree

----------------------------------------------------------------

-- | Creating 'HTree'.
toHTree :: [Bits] -> HTree
toHTree :: [Bits] -> HTree
toHTree [Bits]
bs = Int -> Bits -> HTree -> HTree
mark Int
1 Bits
eos forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Int -> [(Int, Bits)] -> (Int, HTree)
build Int
0 forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 .. Int
idxEos] [Bits]
bs
  where
    eos :: Bits
eos = [Bits]
bs forall a. [a] -> Int -> a
!! Int
idxEos

build :: Int -> [(Int, Bits)] -> (Int, HTree)
build :: Int -> [(Int, Bits)] -> (Int, HTree)
build Int
cnt0 [(Int
v, [])] = (Int
cnt0, EOSInfo -> Int -> HTree
Tip forall a. Maybe a
Nothing Int
v)
build Int
cnt0 [(Int, Bits)]
xs =
    let (Int
cnt1, HTree
l) = Int -> [(Int, Bits)] -> (Int, HTree)
build (Int
cnt0 forall a. Num a => a -> a -> a
+ Int
1) [(Int, Bits)]
fs
        (Int
cnt2, HTree
r) = Int -> [(Int, Bits)] -> (Int, HTree)
build Int
cnt1 [(Int, Bits)]
ts
     in (Int
cnt2, EOSInfo -> Int -> HTree -> HTree -> HTree
Bin forall a. Maybe a
Nothing Int
cnt0 HTree
l HTree
r)
  where
    ([(Int, Bits)]
fs', [(Int, Bits)]
ts') = forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (forall a. Eq a => a -> a -> Bool
(==) B
F forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> a
head forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Int, Bits)]
xs
    fs :: [(Int, Bits)]
fs = forall a b. (a -> b) -> [a] -> [b]
map (forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second forall a. [a] -> [a]
tail) [(Int, Bits)]
fs'
    ts :: [(Int, Bits)]
ts = forall a b. (a -> b) -> [a] -> [b]
map (forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second forall a. [a] -> [a]
tail) [(Int, Bits)]
ts'

-- | Marking the EOS path
mark :: Int -> Bits -> HTree -> HTree
mark :: Int -> Bits -> HTree -> HTree
mark Int
i [] (Tip EOSInfo
Nothing Int
v) = EOSInfo -> Int -> HTree
Tip (forall a. a -> Maybe a
Just Int
i) Int
v
mark Int
i (B
F : Bits
bs) (Bin EOSInfo
Nothing Int
n HTree
l HTree
r) = EOSInfo -> Int -> HTree -> HTree -> HTree
Bin (forall a. a -> Maybe a
Just Int
i) Int
n (Int -> Bits -> HTree -> HTree
mark (Int
i forall a. Num a => a -> a -> a
+ Int
1) Bits
bs HTree
l) HTree
r
mark Int
i (B
T : Bits
bs) (Bin EOSInfo
Nothing Int
n HTree
l HTree
r) = EOSInfo -> Int -> HTree -> HTree -> HTree
Bin (forall a. a -> Maybe a
Just Int
i) Int
n HTree
l (Int -> Bits -> HTree -> HTree
mark (Int
i forall a. Num a => a -> a -> a
+ Int
1) Bits
bs HTree
r)
mark Int
_ Bits
_ HTree
_ = forall a. HasCallStack => String -> a
error String
"mark"

----------------------------------------------------------------

flatten :: HTree -> [HTree]
flatten :: HTree -> [HTree]
flatten (Tip EOSInfo
_ Int
_) = []
flatten t :: HTree
t@(Bin EOSInfo
_ Int
_ HTree
l HTree
r) = HTree
t forall a. a -> [a] -> [a]
: (HTree -> [HTree]
flatten HTree
l forall a. [a] -> [a] -> [a]
++ HTree -> [HTree]
flatten HTree
r)