1{-# LANGUAGE TupleSections #-}
2{-# LANGUAGE FlexibleContexts #-}
3-- | Module used by the jpeg decoder internally, shouldn't be used
4-- in user code.
5module Codec.Picture.Jpg.Internal.DefaultTable( DctComponent( .. )
6                                     , HuffmanTree( .. )
7                                     , HuffmanTable
8                                     , HuffmanPackedTree
9                                     , MacroBlock
10                                     , QuantificationTable
11                                     , HuffmanWriterCode
12                                     , scaleQuantisationMatrix
13                                     , makeMacroBlock
14                                     , makeInverseTable
15                                     , buildHuffmanTree
16                                     , packHuffmanTree
17                                     , huffmanPackedDecode
18
19                                     , defaultChromaQuantizationTable
20
21                                     , defaultLumaQuantizationTable
22
23                                     , defaultAcChromaHuffmanTree
24                                     , defaultAcChromaHuffmanTable
25
26                                     , defaultAcLumaHuffmanTree
27                                     , defaultAcLumaHuffmanTable
28
29                                     , defaultDcChromaHuffmanTree
30                                     , defaultDcChromaHuffmanTable
31
32                                     , defaultDcLumaHuffmanTree
33                                     , defaultDcLumaHuffmanTable
34                                     ) where
35
36import Data.Int( Int16 )
37import Foreign.Storable ( Storable )
38import Control.Monad.ST( runST )
39import qualified Data.Vector.Storable as SV
40import qualified Data.Vector as V
41import Data.Bits( unsafeShiftL, (.|.), (.&.) )
42import Data.Word( Word8, Word16 )
43import Data.List( foldl' )
44import qualified Data.Vector.Storable.Mutable as M
45
46import Codec.Picture.BitWriter
47
48-- | Tree storing the code used for huffman encoding.
49data HuffmanTree = Branch HuffmanTree HuffmanTree -- ^ If bit is 0 take the first subtree, if 1, the right.
50                 | Leaf Word8       -- ^ We should output the value
51                 | Empty            -- ^ no value present
52                 deriving (Eq, Show)
53
54type HuffmanPackedTree = SV.Vector Word16
55
56type HuffmanWriterCode = V.Vector (Word8, Word16)
57
58packHuffmanTree :: HuffmanTree -> HuffmanPackedTree
59packHuffmanTree tree = runST $ do
60    table <- M.replicate 512 0x8000
61    let aux (Empty) idx = return $ idx + 1
62        aux (Leaf v) idx = do
63            (table `M.unsafeWrite` idx) $ fromIntegral v .|. 0x4000
64            return $ idx + 1
65
66        aux (Branch i1@(Leaf _) i2@(Leaf _)) idx =
67            aux i1 idx >>= aux i2
68
69        aux (Branch i1@(Leaf _) i2) idx = do
70            _ <- aux i1 idx
71            ix2 <- aux i2 $ idx + 2
72            (table `M.unsafeWrite` (idx + 1)) $ fromIntegral $ idx + 2
73            return ix2
74
75        aux (Branch i1 i2@(Leaf _)) idx = do
76            ix1 <- aux i1 (idx + 2)
77            _ <- aux i2 (idx + 1)
78            (table `M.unsafeWrite` idx) . fromIntegral $ idx + 2
79            return ix1
80
81        aux (Branch i1 i2) idx = do
82            ix1 <- aux i1 (idx + 2)
83            ix2 <- aux i2 ix1
84            (table `M.unsafeWrite` idx) (fromIntegral $ idx + 2)
85            (table `M.unsafeWrite` (idx + 1)) (fromIntegral ix1)
86            return ix2
87    _ <- aux tree 0
88    SV.unsafeFreeze table
89
90makeInverseTable :: HuffmanTree -> HuffmanWriterCode
91makeInverseTable t = V.replicate 255 (0,0) V.// inner 0 0 t
92  where inner _     _     Empty   = []
93        inner depth code (Leaf v) = [(fromIntegral v, (depth, code))]
94        inner depth code (Branch l r) =
95          inner (depth + 1) shifted l ++ inner (depth + 1) (shifted .|. 1) r
96            where shifted = code `unsafeShiftL` 1
97
98-- | Represent a compact array of 8 * 8 values. The size
99-- is not guarenteed by type system, but if makeMacroBlock is
100-- used, everything should be fine size-wise
101type MacroBlock a = SV.Vector a
102
103type QuantificationTable = MacroBlock Int16
104
105-- | Helper function to create pure macro block of the good size.
106makeMacroBlock :: (Storable a) => [a] -> MacroBlock a
107makeMacroBlock = SV.fromListN 64
108
109-- | Enumeration used to search in the tables for different components.
110data DctComponent = DcComponent | AcComponent
111    deriving (Eq, Show)
112
113-- | Transform parsed coefficients from the jpeg header to a
114-- tree which can be used to decode data.
115buildHuffmanTree :: [[Word8]] -> HuffmanTree
116buildHuffmanTree table = foldl' insertHuffmanVal Empty
117                       . concatMap (\(i, t) -> map (i + 1,) t)
118                       $ zip ([0..] :: [Int]) table
119  where isTreeFullyDefined Empty = False
120        isTreeFullyDefined (Leaf _) = True
121        isTreeFullyDefined (Branch l r) = isTreeFullyDefined l && isTreeFullyDefined r
122
123        insertHuffmanVal Empty (0, val) = Leaf val
124        insertHuffmanVal Empty (d, val) = Branch (insertHuffmanVal Empty (d - 1, val)) Empty
125        insertHuffmanVal (Branch l r) (d, val)
126            | isTreeFullyDefined l = Branch l (insertHuffmanVal r (d - 1, val))
127            | otherwise            = Branch (insertHuffmanVal l (d - 1, val)) r
128        insertHuffmanVal (Leaf _) _ = error "Inserting in value, shouldn't happen"
129
130scaleQuantisationMatrix :: Int -> QuantificationTable -> QuantificationTable
131scaleQuantisationMatrix quality
132    | quality < 0 = scaleQuantisationMatrix 0
133        -- shouldn't show much difference than with 1,
134        -- but hey, at least we're complete
135    | quality == 0 = SV.map (scale (10000 :: Int))
136    | quality < 50 = let qq = 5000 `div` quality
137                     in SV.map (scale qq)
138    | otherwise    = SV.map (scale q)
139          where q = 200 - quality * 2
140                scale coeff i = fromIntegral . min 255
141                                             . max 1
142                                             $ fromIntegral i * coeff `div` 100
143
144huffmanPackedDecode :: HuffmanPackedTree -> BoolReader s Word8
145huffmanPackedDecode table = getNextBitJpg >>= aux 0
146  where aux idx b
147            | (v .&. 0x8000) /= 0 = return  0
148            | (v .&. 0x4000) /= 0 = return . fromIntegral $ v .&. 0xFF
149            | otherwise = getNextBitJpg >>= aux v
150          where tableIndex | b = idx + 1
151                           | otherwise = idx
152                v = table `SV.unsafeIndex` fromIntegral tableIndex
153
154defaultLumaQuantizationTable :: QuantificationTable
155defaultLumaQuantizationTable = makeMacroBlock
156    [16, 11, 10, 16,  24,  40,  51,  61
157    ,12, 12, 14, 19,  26,  58,  60,  55
158    ,14, 13, 16, 24,  40,  57,  69,  56
159    ,14, 17, 22, 29,  51,  87,  80,  62
160    ,18, 22, 37, 56,  68, 109, 103,  77
161    ,24, 35, 55, 64,  81, 104, 113,  92
162    ,49, 64, 78, 87, 103, 121, 120, 101
163    ,72, 92, 95, 98, 112, 100, 103,  99
164    ]
165
166defaultChromaQuantizationTable :: QuantificationTable
167defaultChromaQuantizationTable = makeMacroBlock
168    [17, 18, 24, 47, 99, 99, 99, 99
169    ,18, 21, 26, 66, 99, 99, 99, 99
170    ,24, 26, 56, 99, 99, 99, 99, 99
171    ,47, 66, 99, 99, 99, 99, 99, 99
172    ,99, 99, 99, 99, 99, 99, 99, 99
173    ,99, 99, 99, 99, 99, 99, 99, 99
174    ,99, 99, 99, 99, 99, 99, 99, 99
175    ,99, 99, 99, 99, 99, 99, 99, 99
176    ]
177
178defaultDcLumaHuffmanTree :: HuffmanTree
179defaultDcLumaHuffmanTree = buildHuffmanTree defaultDcLumaHuffmanTable
180
181-- | From the Table K.3 of ITU-81 (p153)
182defaultDcLumaHuffmanTable :: HuffmanTable
183defaultDcLumaHuffmanTable =
184    [ []
185    , [0]
186    , [1, 2, 3, 4, 5]
187    , [6]
188    , [7]
189    , [8]
190    , [9]
191    , [10]
192    , [11]
193    , []
194    , []
195    , []
196    , []
197    , []
198    , []
199    , []
200    ]
201
202defaultDcChromaHuffmanTree :: HuffmanTree
203defaultDcChromaHuffmanTree = buildHuffmanTree defaultDcChromaHuffmanTable
204
205-- | From the Table K.4 of ITU-81 (p153)
206defaultDcChromaHuffmanTable :: HuffmanTable
207defaultDcChromaHuffmanTable =
208    [ []
209    , [0, 1, 2]
210    , [3]
211    , [4]
212    , [5]
213    , [6]
214    , [7]
215    , [8]
216    , [9]
217    , [10]
218    , [11]
219    , []
220    , []
221    , []
222    , []
223    , []
224    ]
225
226defaultAcLumaHuffmanTree :: HuffmanTree
227defaultAcLumaHuffmanTree = buildHuffmanTree defaultAcLumaHuffmanTable
228
229-- | From the Table K.5 of ITU-81 (p154)
230defaultAcLumaHuffmanTable :: HuffmanTable
231defaultAcLumaHuffmanTable =
232    [ []
233    , [0x01, 0x02]
234    , [0x03]
235    , [0x00, 0x04, 0x11]
236    , [0x05, 0x12, 0x21]
237    , [0x31, 0x41]
238    , [0x06, 0x13, 0x51, 0x61]
239    , [0x07, 0x22, 0x71]
240    , [0x14, 0x32, 0x81, 0x91, 0xA1]
241    , [0x08, 0x23, 0x42, 0xB1, 0xC1]
242    , [0x15, 0x52, 0xD1, 0xF0]
243    , [0x24, 0x33, 0x62, 0x72]
244    , []
245    , []
246    , [0x82]
247    , [0x09, 0x0A, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2A, 0x34, 0x35
248      ,0x36, 0x37, 0x38, 0x39, 0x3A, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48, 0x49, 0x4A, 0x53, 0x54
249      ,0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x6A, 0x73
250      ,0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7A, 0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89, 0x8A
251      ,0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6, 0xA7
252      ,0xA8, 0xA9, 0xAA, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xC2, 0xC3, 0xC4
253      ,0xC5, 0xC6, 0xC7, 0xC8, 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8, 0xD9, 0xDA
254      ,0xE1, 0xE2, 0xE3, 0xE4, 0xE5, 0xE6, 0xE7, 0xE8, 0xE9, 0xEA, 0xF1, 0xF2, 0xF3, 0xF4, 0xF5
255      ,0xF6, 0xF7, 0xF8, 0xF9, 0xFA]
256    ]
257
258type HuffmanTable = [[Word8]]
259
260defaultAcChromaHuffmanTree :: HuffmanTree
261defaultAcChromaHuffmanTree = buildHuffmanTree defaultAcChromaHuffmanTable
262
263defaultAcChromaHuffmanTable :: HuffmanTable
264defaultAcChromaHuffmanTable =
265    [ []
266    , [0x00, 0x01]
267    , [0x02]
268    , [0x03, 0x11]
269    , [0x04, 0x05, 0x21, 0x31]
270    , [0x06, 0x12, 0x41, 0x51]
271    , [0x07, 0x61, 0x71]
272    , [0x13, 0x22, 0x32, 0x81]
273    , [0x08, 0x14, 0x42, 0x91, 0xA1, 0xB1, 0xC1]
274    , [0x09, 0x23, 0x33, 0x52, 0xF0]
275    , [0x15, 0x62, 0x72, 0xD1]
276    , [0x0A, 0x16, 0x24, 0x34]
277    , []
278    , [0xE1]
279    , [0x25, 0xF1]
280    , [ 0x17, 0x18, 0x19, 0x1A, 0x26, 0x27, 0x28, 0x29, 0x2A, 0x35
281      , 0x36, 0x37, 0x38, 0x39, 0x3A, 0x43, 0x44, 0x45, 0x46, 0x47
282      , 0x48, 0x49, 0x4A, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59
283      , 0x5A, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x6A, 0x73
284      , 0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7A, 0x82, 0x83, 0x84
285      , 0x85, 0x86, 0x87, 0x88, 0x89, 0x8A, 0x92, 0x93, 0x94, 0x95
286      , 0x96, 0x97, 0x98, 0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6
287      , 0xA7, 0xA8, 0xA9, 0xAA, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7
288      , 0xB8, 0xB9, 0xBA, 0xC2, 0xC3, 0xC4, 0xC5, 0xC6, 0xC7, 0xC8
289      , 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8, 0xD9
290      , 0xDA, 0xE2, 0xE3, 0xE4, 0xE5, 0xE6, 0xE7, 0xE8, 0xE9, 0xEA
291      , 0xF2, 0xF3, 0xF4, 0xF5, 0xF6, 0xF7, 0xF8, 0xF9, 0xFA
292      ]
293    ]
294
295