1{-# LANGUAGE BangPatterns #-}
2
3module Network.HPACK.Huffman.Tree (
4  -- * Huffman decoding
5    HTree(..)
6  , eosInfo
7  , toHTree
8  , showTree
9  , printTree
10  , flatten
11  ) where
12
13import Control.Arrow (second)
14
15import Imports
16import Network.HPACK.Huffman.Bit
17import Network.HPACK.Huffman.Params
18
19----------------------------------------------------------------
20
21type EOSInfo = Maybe Int
22
23-- | Type for Huffman decoding.
24data HTree = Tip
25             !EOSInfo            -- EOS info from 1
26             {-# UNPACK #-} !Int -- Decoded value. Essentially Word8
27           | Bin
28             !EOSInfo            -- EOS info from 1
29             {-# UNPACK #-} !Int -- Sequence no from 0
30             !HTree              -- Left
31             !HTree              -- Right
32           deriving Show
33
34eosInfo :: HTree -> EOSInfo
35eosInfo (Tip mx _)     = mx
36eosInfo (Bin mx _ _ _) = mx
37
38----------------------------------------------------------------
39
40showTree :: HTree -> String
41showTree = showTree' ""
42
43showTree' :: String -> HTree -> String
44showTree' _    (Tip _ i)     = show i ++ "\n"
45showTree' pref (Bin _ n l r) = "No " ++ show n ++ "\n"
46                            ++ pref ++ "+ " ++ showTree' pref' l
47                            ++ pref ++ "+ " ++ showTree' pref' r
48  where
49    pref' = "  " ++ pref
50
51printTree :: HTree -> IO ()
52printTree = putStr . showTree
53
54----------------------------------------------------------------
55
56-- | Creating 'HTree'.
57toHTree :: [Bits] -> HTree
58toHTree bs = mark 1 eos $ snd $ build 0 $ zip [0..idxEos] bs
59  where
60    eos = bs !! idxEos
61
62build :: Int -> [(Int,Bits)] -> (Int, HTree)
63build !cnt0 [(v,[])] = (cnt0,Tip Nothing v)
64build !cnt0 xs       = let (cnt1,l) = build (cnt0 + 1) fs
65                           (cnt2,r) = build cnt1 ts
66                       in (cnt2, Bin Nothing cnt0 l r)
67  where
68    (fs',ts') = partition ((==) F . head . snd) xs
69    fs = map (second tail) fs'
70    ts = map (second tail) ts'
71
72-- | Marking the EOS path
73mark :: Int -> Bits -> HTree -> HTree
74mark i []     (Tip Nothing v)     = Tip (Just i) v
75mark i (F:bs) (Bin Nothing n l r) = Bin (Just i) n (mark (i+1) bs l) r
76mark i (T:bs) (Bin Nothing n l r) = Bin (Just i) n l (mark (i+1) bs r)
77mark _ _      _                   = error "mark"
78
79----------------------------------------------------------------
80
81flatten :: HTree -> [HTree]
82flatten (Tip _ _)       = []
83flatten t@(Bin _ _ l r) = t : (flatten l ++ flatten r)
84