Skip to main content

webp_rust/decoder/
vp8.rs

1use crate::decoder::quant::{parse_quantization, Quantization};
2use crate::decoder::tree::{
3    parse_intra_mode_row, parse_probability_tables, parse_probability_updates, MacroBlockHeader,
4    ProbabilityTables, ProbabilityUpdateSummary,
5};
6use crate::decoder::vp8i::{
7    B_DC_PRED, MAX_NUM_PARTITIONS, MB_FEATURE_TREE_PROBS, NUM_MB_SEGMENTS, NUM_MODE_LF_DELTAS,
8    NUM_REF_LF_DELTAS, VP8L_FRAME_HEADER_SIZE, VP8_FRAME_HEADER_SIZE,
9};
10use crate::decoder::DecoderError;
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub struct Vp8FrameHeader {
14    pub key_frame: bool,
15    pub profile: u8,
16    pub show: bool,
17    pub partition_length: usize,
18}
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub struct Vp8PictureHeader {
22    pub width: u16,
23    pub height: u16,
24    pub xscale: u8,
25    pub yscale: u8,
26    pub colorspace: u8,
27    pub clamp_type: u8,
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub struct SegmentHeader {
32    pub use_segment: bool,
33    pub update_map: bool,
34    pub absolute_delta: bool,
35    pub quantizer: [i8; NUM_MB_SEGMENTS],
36    pub filter_strength: [i8; NUM_MB_SEGMENTS],
37    pub segment_probs: [u8; MB_FEATURE_TREE_PROBS],
38}
39
40impl Default for SegmentHeader {
41    fn default() -> Self {
42        Self {
43            use_segment: false,
44            update_map: false,
45            absolute_delta: true,
46            quantizer: [0; NUM_MB_SEGMENTS],
47            filter_strength: [0; NUM_MB_SEGMENTS],
48            segment_probs: [255; MB_FEATURE_TREE_PROBS],
49        }
50    }
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub enum FilterType {
55    Off,
56    Simple,
57    Complex,
58}
59
60#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61pub struct FilterHeader {
62    pub simple: bool,
63    pub level: u8,
64    pub sharpness: u8,
65    pub use_lf_delta: bool,
66    pub ref_lf_delta: [i8; NUM_REF_LF_DELTAS],
67    pub mode_lf_delta: [i8; NUM_MODE_LF_DELTAS],
68    pub filter_type: FilterType,
69}
70
71#[derive(Debug, Clone, Copy, PartialEq, Eq)]
72pub struct LosslessInfo {
73    pub width: usize,
74    pub height: usize,
75    pub has_alpha: bool,
76}
77
78#[derive(Debug, Clone, PartialEq, Eq)]
79pub struct LossyHeader {
80    pub frame: Vp8FrameHeader,
81    pub picture: Vp8PictureHeader,
82    pub macroblock_width: usize,
83    pub macroblock_height: usize,
84    pub segment: SegmentHeader,
85    pub filter: FilterHeader,
86    pub token_partition_sizes: Vec<usize>,
87    pub quantization: Quantization,
88    pub probabilities: ProbabilityUpdateSummary,
89}
90
91#[derive(Debug, Clone, PartialEq, Eq)]
92pub struct MacroBlockHeaders {
93    pub frame: LossyHeader,
94    pub macroblocks: Vec<MacroBlockHeader>,
95}
96
97#[derive(Debug, Clone, PartialEq, Eq)]
98pub struct MacroBlockData {
99    pub header: MacroBlockHeader,
100    pub coeffs: [i16; 384],
101    pub non_zero_y: u32,
102    pub non_zero_uv: u32,
103}
104
105#[derive(Debug, Clone, PartialEq, Eq)]
106pub struct MacroBlockDataFrame {
107    pub frame: LossyHeader,
108    pub macroblocks: Vec<MacroBlockData>,
109}
110
111#[derive(Debug, Clone, Copy, Default)]
112struct NonZeroContext {
113    nz: u8,
114    nz_dc: u8,
115}
116
117#[derive(Debug, Clone)]
118pub struct Vp8BoolDecoder<'a> {
119    data: &'a [u8],
120    position: usize,
121    value: u64,
122    range: u32,
123    bits: i32,
124    eof: bool,
125}
126
127impl<'a> Vp8BoolDecoder<'a> {
128    pub fn new(data: &'a [u8]) -> Self {
129        let mut reader = Self {
130            data,
131            position: 0,
132            value: 0,
133            range: 255 - 1,
134            bits: -8,
135            eof: false,
136        };
137        reader.load_new_bytes();
138        reader
139    }
140
141    pub fn eof(&self) -> bool {
142        self.eof
143    }
144
145    fn load_new_bytes(&mut self) {
146        while self.bits < 0 {
147            if self.position < self.data.len() {
148                self.bits += 8;
149                self.value = (self.value << 8) | self.data[self.position] as u64;
150                self.position += 1;
151            } else if !self.eof {
152                self.value <<= 8;
153                self.bits += 8;
154                self.eof = true;
155            } else {
156                self.bits = 0;
157            }
158        }
159    }
160
161    pub fn get(&mut self) -> u32 {
162        self.get_bit(0x80)
163    }
164
165    pub fn get_value(&mut self, num_bits: usize) -> u32 {
166        let mut value = 0u32;
167        for bit_index in (0..num_bits).rev() {
168            value |= self.get() << bit_index;
169        }
170        value
171    }
172
173    pub fn get_signed_value(&mut self, num_bits: usize) -> i32 {
174        let value = self.get_value(num_bits) as i32;
175        if self.get() == 1 {
176            -value
177        } else {
178            value
179        }
180    }
181
182    pub fn get_signed(&mut self, value: i32) -> i32 {
183        if self.get() == 1 {
184            -value
185        } else {
186            value
187        }
188    }
189
190    pub fn get_bit(&mut self, prob: u8) -> u32 {
191        if self.bits < 0 {
192            self.load_new_bytes();
193        }
194
195        let pos = self.bits as u32;
196        let mut range = self.range;
197        let split = (range * prob as u32) >> 8;
198        let value = (self.value >> pos) as u32;
199        let bit = (value > split) as u32;
200        if bit == 1 {
201            range -= split;
202            self.value -= ((split + 1) as u64) << pos;
203        } else {
204            range = split + 1;
205        }
206
207        let shift = 7 ^ (31 - range.leading_zeros()) as i32;
208        range <<= shift as u32;
209        self.bits -= shift;
210        self.range = range - 1;
211        bit
212    }
213}
214
215pub fn check_lossy_signature(data: &[u8]) -> bool {
216    data.len() >= 3 && data[0] == 0x9d && data[1] == 0x01 && data[2] == 0x2a
217}
218
219pub fn get_info(data: &[u8], chunk_size: usize) -> Result<(usize, usize), DecoderError> {
220    if data.len() < VP8_FRAME_HEADER_SIZE {
221        return Err(DecoderError::NotEnoughData("VP8 frame header"));
222    }
223    if !check_lossy_signature(&data[3..]) {
224        return Err(DecoderError::Bitstream("bad VP8 signature"));
225    }
226
227    let bits = data[0] as u32 | ((data[1] as u32) << 8) | ((data[2] as u32) << 16);
228    let key_frame = (bits & 1) == 0;
229    let profile = ((bits >> 1) & 0x07) as u8;
230    let show = ((bits >> 4) & 1) == 1;
231    let partition_length = (bits >> 5) as usize;
232    let width = ((((data[7] as u16) << 8) | data[6] as u16) & 0x3fff) as usize;
233    let height = ((((data[9] as u16) << 8) | data[8] as u16) & 0x3fff) as usize;
234
235    if !key_frame {
236        return Err(DecoderError::Unsupported("interframes are not supported"));
237    }
238    if profile > 3 {
239        return Err(DecoderError::Bitstream("unknown VP8 profile"));
240    }
241    if !show {
242        return Err(DecoderError::Unsupported("invisible VP8 frame"));
243    }
244    if partition_length >= chunk_size {
245        return Err(DecoderError::Bitstream("bad VP8 partition length"));
246    }
247    if width == 0 || height == 0 {
248        return Err(DecoderError::Bitstream("invalid VP8 dimensions"));
249    }
250
251    Ok((width, height))
252}
253
254pub fn check_lossless_signature(data: &[u8]) -> bool {
255    data.len() >= VP8L_FRAME_HEADER_SIZE && data[0] == 0x2f && (data[4] >> 5) == 0
256}
257
258pub fn get_lossless_info(data: &[u8]) -> Result<LosslessInfo, DecoderError> {
259    if data.len() < VP8L_FRAME_HEADER_SIZE {
260        return Err(DecoderError::NotEnoughData("VP8L frame header"));
261    }
262    if !check_lossless_signature(data) {
263        return Err(DecoderError::Bitstream("bad VP8L signature"));
264    }
265
266    let bits = u32::from_le_bytes([data[1], data[2], data[3], data[4]]);
267    let width = ((bits & 0x3fff) + 1) as usize;
268    let height = (((bits >> 14) & 0x3fff) + 1) as usize;
269    let has_alpha = ((bits >> 28) & 1) == 1;
270    let version = (bits >> 29) & 0x07;
271
272    if version != 0 {
273        return Err(DecoderError::Bitstream("unsupported VP8L version"));
274    }
275
276    Ok(LosslessInfo {
277        width,
278        height,
279        has_alpha,
280    })
281}
282
283fn parse_segment_header(br: &mut Vp8BoolDecoder<'_>) -> Result<SegmentHeader, DecoderError> {
284    let mut header = SegmentHeader {
285        use_segment: br.get() == 1,
286        ..SegmentHeader::default()
287    };
288    if header.use_segment {
289        header.update_map = br.get() == 1;
290        if br.get() == 1 {
291            header.absolute_delta = br.get() == 1;
292            for value in &mut header.quantizer {
293                *value = if br.get() == 1 {
294                    br.get_signed_value(7) as i8
295                } else {
296                    0
297                };
298            }
299            for value in &mut header.filter_strength {
300                *value = if br.get() == 1 {
301                    br.get_signed_value(6) as i8
302                } else {
303                    0
304                };
305            }
306        }
307        if header.update_map {
308            for value in &mut header.segment_probs {
309                *value = if br.get() == 1 {
310                    br.get_value(8) as u8
311                } else {
312                    255
313                };
314            }
315        }
316    }
317
318    if br.eof() {
319        return Err(DecoderError::Bitstream("cannot parse segment header"));
320    }
321
322    Ok(header)
323}
324
325fn parse_filter_header(br: &mut Vp8BoolDecoder<'_>) -> Result<FilterHeader, DecoderError> {
326    let simple = br.get() == 1;
327    let level = br.get_value(6) as u8;
328    let sharpness = br.get_value(3) as u8;
329    let use_lf_delta = br.get() == 1;
330    let mut header = FilterHeader {
331        simple,
332        level,
333        sharpness,
334        use_lf_delta,
335        ref_lf_delta: [0; NUM_REF_LF_DELTAS],
336        mode_lf_delta: [0; NUM_MODE_LF_DELTAS],
337        filter_type: FilterType::Off,
338    };
339
340    if use_lf_delta && br.get() == 1 {
341        for value in &mut header.ref_lf_delta {
342            if br.get() == 1 {
343                *value = br.get_signed_value(6) as i8;
344            }
345        }
346        for value in &mut header.mode_lf_delta {
347            if br.get() == 1 {
348                *value = br.get_signed_value(6) as i8;
349            }
350        }
351    }
352
353    header.filter_type = if level == 0 {
354        FilterType::Off
355    } else if simple {
356        FilterType::Simple
357    } else {
358        FilterType::Complex
359    };
360
361    if br.eof() {
362        return Err(DecoderError::Bitstream("cannot parse filter header"));
363    }
364
365    Ok(header)
366}
367
368fn parse_token_partitions(
369    br: &mut Vp8BoolDecoder<'_>,
370    data: &[u8],
371) -> Result<Vec<usize>, DecoderError> {
372    let num_parts_minus_one = (1usize << br.get_value(2)) - 1;
373    if num_parts_minus_one >= MAX_NUM_PARTITIONS {
374        return Err(DecoderError::Bitstream("too many VP8 token partitions"));
375    }
376
377    let size_bytes = num_parts_minus_one * 3;
378    if data.len() < size_bytes {
379        return Err(DecoderError::NotEnoughData("VP8 token partition sizes"));
380    }
381
382    let mut partitions = Vec::with_capacity(num_parts_minus_one + 1);
383    let mut size_left = data.len() - size_bytes;
384    for chunk in data[..size_bytes].chunks_exact(3) {
385        let stored = chunk[0] as usize | ((chunk[1] as usize) << 8) | ((chunk[2] as usize) << 16);
386        if stored > size_left {
387            return Err(DecoderError::NotEnoughData("VP8 token partition"));
388        }
389        partitions.push(stored);
390        size_left -= stored;
391    }
392    partitions.push(size_left);
393
394    if data.len() == size_bytes {
395        return Err(DecoderError::NotEnoughData("VP8 token partitions"));
396    }
397
398    Ok(partitions)
399}
400
401const CAT3: [u8; 4] = [173, 148, 140, 0];
402const CAT4: [u8; 5] = [176, 155, 140, 135, 0];
403const CAT5: [u8; 6] = [180, 157, 141, 134, 130, 0];
404const CAT6: [u8; 12] = [254, 254, 243, 230, 196, 177, 153, 140, 133, 130, 129, 0];
405const ZIGZAG: [usize; 16] = [0, 1, 4, 8, 5, 2, 3, 6, 9, 12, 13, 10, 7, 11, 14, 15];
406
407fn transform_wht(input: &[i16; 16]) -> [i16; 16] {
408    let mut tmp = [0i32; 16];
409    for i in 0..4 {
410        let a0 = input[i] as i32 + input[12 + i] as i32;
411        let a1 = input[4 + i] as i32 + input[8 + i] as i32;
412        let a2 = input[4 + i] as i32 - input[8 + i] as i32;
413        let a3 = input[i] as i32 - input[12 + i] as i32;
414        tmp[i] = a0 + a1;
415        tmp[8 + i] = a0 - a1;
416        tmp[4 + i] = a3 + a2;
417        tmp[12 + i] = a3 - a2;
418    }
419
420    let mut out = [0i16; 16];
421    for i in 0..4 {
422        let base = i * 4;
423        let dc = tmp[base] + 3;
424        let a0 = dc + tmp[base + 3];
425        let a1 = tmp[base + 1] + tmp[base + 2];
426        let a2 = tmp[base + 1] - tmp[base + 2];
427        let a3 = dc - tmp[base + 3];
428        out[base] = ((a0 + a1) >> 3) as i16;
429        out[base + 1] = ((a3 + a2) >> 3) as i16;
430        out[base + 2] = ((a0 - a1) >> 3) as i16;
431        out[base + 3] = ((a3 - a2) >> 3) as i16;
432    }
433    out
434}
435
436fn get_large_value(br: &mut Vp8BoolDecoder<'_>, p: &[u8; 11]) -> i32 {
437    if br.get_bit(p[3]) == 0 {
438        if br.get_bit(p[4]) == 0 {
439            2
440        } else {
441            3 + br.get_bit(p[5]) as i32
442        }
443    } else if br.get_bit(p[6]) == 0 {
444        if br.get_bit(p[7]) == 0 {
445            5 + br.get_bit(159) as i32
446        } else {
447            7 + 2 * br.get_bit(165) as i32 + br.get_bit(145) as i32
448        }
449    } else {
450        let (cat, table): (usize, &[u8]) = if br.get_bit(p[8]) == 0 {
451            if br.get_bit(p[9]) == 0 {
452                (0, &CAT3)
453            } else {
454                (1, &CAT4)
455            }
456        } else if br.get_bit(p[10]) == 0 {
457            (2, &CAT5)
458        } else {
459            (3, &CAT6)
460        };
461        let mut value = 0i32;
462        for &prob in table {
463            if prob == 0 {
464                break;
465            }
466            value = value + value + br.get_bit(prob) as i32;
467        }
468        value + 3 + (8 << cat)
469    }
470}
471
472fn get_coeffs(
473    br: &mut Vp8BoolDecoder<'_>,
474    probabilities: &ProbabilityTables,
475    coeff_type: usize,
476    ctx: usize,
477    dq: [u16; 2],
478    start: usize,
479    out: &mut [i16],
480) -> usize {
481    let mut n = start;
482    let mut p = probabilities.coeff_probs(coeff_type, n, ctx);
483    while n < 16 {
484        if br.get_bit(p[0]) == 0 {
485            return n;
486        }
487        while br.get_bit(p[1]) == 0 {
488            n += 1;
489            if n == 16 {
490                return 16;
491            }
492            p = probabilities.coeff_probs(coeff_type, n, 0);
493        }
494
495        let next_ctx;
496        let value = if br.get_bit(p[2]) == 0 {
497            next_ctx = 1;
498            1
499        } else {
500            next_ctx = 2;
501            get_large_value(br, p)
502        };
503        let dequant = if n > 0 { dq[1] } else { dq[0] } as i32;
504        out[ZIGZAG[n]] = (br.get_signed(value) * dequant) as i16;
505        n += 1;
506        p = probabilities.coeff_probs(coeff_type, n, next_ctx);
507    }
508    16
509}
510
511fn nz_code_bits(nz_coeffs: u32, nz: usize, dc_nz: bool) -> u32 {
512    (nz_coeffs << 2)
513        | if nz > 3 {
514            3
515        } else if nz > 1 {
516            2
517        } else if dc_nz {
518            1
519        } else {
520            0
521        }
522}
523
524fn parse_residuals(
525    header: MacroBlockHeader,
526    top: &mut NonZeroContext,
527    left: &mut NonZeroContext,
528    token_br: &mut Vp8BoolDecoder<'_>,
529    quantization: &Quantization,
530    probabilities: &ProbabilityTables,
531) -> MacroBlockData {
532    let mut coeffs = [0i16; 384];
533    if header.skip {
534        top.nz = 0;
535        left.nz = 0;
536        if !header.is_i4x4 {
537            top.nz_dc = 0;
538            left.nz_dc = 0;
539        }
540        return MacroBlockData {
541            header,
542            coeffs,
543            non_zero_y: 0,
544            non_zero_uv: 0,
545        };
546    }
547
548    let q = &quantization.matrices[header.segment as usize];
549    let mut offset = 0usize;
550    let first;
551    let coeff_type;
552    if !header.is_i4x4 {
553        let mut dc = [0i16; 16];
554        let ctx = (top.nz_dc + left.nz_dc) as usize;
555        let nz = get_coeffs(token_br, probabilities, 1, ctx, q.y2, 0, &mut dc);
556        let has_dc = nz > 0;
557        top.nz_dc = has_dc as u8;
558        left.nz_dc = has_dc as u8;
559        if nz > 1 {
560            let transformed = transform_wht(&dc);
561            for (block, value) in transformed.into_iter().enumerate() {
562                coeffs[block * 16] = value;
563            }
564        } else {
565            let dc0 = ((dc[0] as i32 + 3) >> 3) as i16;
566            for block in 0..16 {
567                coeffs[block * 16] = dc0;
568            }
569        }
570        first = 1;
571        coeff_type = 0;
572    } else {
573        first = 0;
574        coeff_type = 3;
575    }
576
577    let mut non_zero_y = 0u32;
578    let mut tnz = top.nz & 0x0f;
579    let mut lnz = left.nz & 0x0f;
580    for _y in 0..4 {
581        let mut l = lnz & 1;
582        let mut nz_coeffs = 0u32;
583        for _x in 0..4 {
584            let ctx = (l + (tnz & 1)) as usize;
585            let nz = get_coeffs(
586                token_br,
587                probabilities,
588                coeff_type,
589                ctx,
590                q.y1,
591                first,
592                &mut coeffs[offset..offset + 16],
593            );
594            l = (nz > first) as u8;
595            tnz = (tnz >> 1) | (l << 7);
596            nz_coeffs = nz_code_bits(nz_coeffs, nz, coeffs[offset] != 0);
597            offset += 16;
598        }
599        tnz >>= 4;
600        lnz = (lnz >> 1) | (l << 7);
601        non_zero_y = (non_zero_y << 8) | nz_coeffs;
602    }
603
604    let mut out_t_nz = tnz;
605    let mut out_l_nz = lnz >> 4;
606    let mut non_zero_uv = 0u32;
607    for ch in [0usize, 2usize] {
608        let mut nz_coeffs = 0u32;
609        let mut tnz = top.nz >> (4 + ch);
610        let mut lnz = left.nz >> (4 + ch);
611        for _y in 0..2 {
612            let mut l = lnz & 1;
613            for _x in 0..2 {
614                let ctx = (l + (tnz & 1)) as usize;
615                let nz = get_coeffs(
616                    token_br,
617                    probabilities,
618                    2,
619                    ctx,
620                    q.uv,
621                    0,
622                    &mut coeffs[offset..offset + 16],
623                );
624                l = (nz > 0) as u8;
625                tnz = (tnz >> 1) | (l << 3);
626                nz_coeffs = nz_code_bits(nz_coeffs, nz, coeffs[offset] != 0);
627                offset += 16;
628            }
629            tnz >>= 2;
630            lnz = (lnz >> 1) | (l << 5);
631        }
632        non_zero_uv |= nz_coeffs << (4 * ch);
633        out_t_nz |= (tnz << 4) << ch;
634        out_l_nz |= (lnz & 0xf0) << ch;
635    }
636    top.nz = out_t_nz;
637    left.nz = out_l_nz;
638
639    MacroBlockData {
640        header,
641        coeffs,
642        non_zero_y,
643        non_zero_uv,
644    }
645}
646
647struct LossyPrefix<'a> {
648    frame: Vp8FrameHeader,
649    picture: Vp8PictureHeader,
650    macroblock_width: usize,
651    macroblock_height: usize,
652    segment: SegmentHeader,
653    filter: FilterHeader,
654    token_partition_sizes: Vec<usize>,
655    quantization: Quantization,
656    partition0_end: usize,
657    br: Vp8BoolDecoder<'a>,
658}
659
660fn parse_lossy_prefix(data: &[u8]) -> Result<LossyPrefix<'_>, DecoderError> {
661    if data.len() < VP8_FRAME_HEADER_SIZE {
662        return Err(DecoderError::NotEnoughData("VP8 frame header"));
663    }
664
665    let frame_bits = data[0] as u32 | ((data[1] as u32) << 8) | ((data[2] as u32) << 16);
666    let frame = Vp8FrameHeader {
667        key_frame: (frame_bits & 1) == 0,
668        profile: ((frame_bits >> 1) & 0x07) as u8,
669        show: ((frame_bits >> 4) & 1) == 1,
670        partition_length: (frame_bits >> 5) as usize,
671    };
672    if !frame.key_frame {
673        return Err(DecoderError::Unsupported("interframes are not supported"));
674    }
675    if frame.profile > 3 {
676        return Err(DecoderError::Bitstream("unknown VP8 profile"));
677    }
678    if !frame.show {
679        return Err(DecoderError::Unsupported("invisible VP8 frame"));
680    }
681    if !check_lossy_signature(&data[3..]) {
682        return Err(DecoderError::Bitstream("bad VP8 signature"));
683    }
684
685    let picture = Vp8PictureHeader {
686        width: (((data[7] as u16) << 8) | data[6] as u16) & 0x3fff,
687        height: (((data[9] as u16) << 8) | data[8] as u16) & 0x3fff,
688        xscale: data[7] >> 6,
689        yscale: data[9] >> 6,
690        colorspace: 0,
691        clamp_type: 0,
692    };
693    if picture.width == 0 || picture.height == 0 {
694        return Err(DecoderError::Bitstream("invalid VP8 dimensions"));
695    }
696
697    let partition0_offset = VP8_FRAME_HEADER_SIZE;
698    let partition0_end = partition0_offset
699        .checked_add(frame.partition_length)
700        .ok_or(DecoderError::Bitstream("VP8 partition length overflow"))?;
701    if partition0_end > data.len() {
702        return Err(DecoderError::NotEnoughData("VP8 partition 0"));
703    }
704
705    let mut br = Vp8BoolDecoder::new(&data[partition0_offset..partition0_end]);
706    let mut picture = picture;
707    picture.colorspace = br.get() as u8;
708    picture.clamp_type = br.get() as u8;
709
710    let segment = parse_segment_header(&mut br)?;
711    let filter = parse_filter_header(&mut br)?;
712    let token_partition_sizes = parse_token_partitions(&mut br, &data[partition0_end..])?;
713    let quantization = parse_quantization(&mut br, &segment)?;
714    let _ = br.get();
715
716    Ok(LossyPrefix {
717        frame,
718        picture,
719        macroblock_width: (picture.width as usize + 15) >> 4,
720        macroblock_height: (picture.height as usize + 15) >> 4,
721        segment,
722        filter,
723        token_partition_sizes,
724        quantization,
725        partition0_end,
726        br,
727    })
728}
729
730fn finish_lossy_header(
731    prefix: &LossyPrefix<'_>,
732    probabilities: ProbabilityUpdateSummary,
733) -> LossyHeader {
734    LossyHeader {
735        frame: prefix.frame,
736        picture: prefix.picture,
737        macroblock_width: prefix.macroblock_width,
738        macroblock_height: prefix.macroblock_height,
739        segment: prefix.segment,
740        filter: prefix.filter,
741        token_partition_sizes: prefix.token_partition_sizes.clone(),
742        quantization: prefix.quantization.clone(),
743        probabilities,
744    }
745}
746
747pub fn parse_lossy_headers(data: &[u8]) -> Result<LossyHeader, DecoderError> {
748    let mut prefix = parse_lossy_prefix(data)?;
749    let probabilities = parse_probability_updates(&mut prefix.br)?;
750    Ok(finish_lossy_header(&prefix, probabilities))
751}
752
753pub fn parse_macroblock_headers(data: &[u8]) -> Result<MacroBlockHeaders, DecoderError> {
754    let mut prefix = parse_lossy_prefix(data)?;
755    let probabilities = parse_probability_updates(&mut prefix.br)?;
756    let frame = finish_lossy_header(&prefix, probabilities);
757
758    let mut top_modes = vec![B_DC_PRED; frame.macroblock_width * 4];
759    let mut macroblocks = Vec::with_capacity(frame.macroblock_width * frame.macroblock_height);
760    for _mb_y in 0..frame.macroblock_height {
761        let mut left_modes = [B_DC_PRED; 4];
762        let row = parse_intra_mode_row(
763            &mut prefix.br,
764            frame.macroblock_width,
765            frame.segment.update_map,
766            &frame.segment.segment_probs,
767            probabilities.use_skip_probability,
768            probabilities.skip_probability.unwrap_or(0),
769            &mut top_modes,
770            &mut left_modes,
771        )?;
772        macroblocks.extend(row);
773    }
774
775    Ok(MacroBlockHeaders { frame, macroblocks })
776}
777
778pub(crate) struct MacroBlockRows<'a> {
779    frame: LossyHeader,
780    br: Vp8BoolDecoder<'a>,
781    token_readers: Vec<Vp8BoolDecoder<'a>>,
782    probabilities: ProbabilityTables,
783    top_modes: Vec<u8>,
784    top_contexts: Vec<NonZeroContext>,
785    next_mb_y: usize,
786}
787
788impl<'a> MacroBlockRows<'a> {
789    pub(crate) fn new(data: &'a [u8]) -> Result<Self, DecoderError> {
790        let mut prefix = parse_lossy_prefix(data)?;
791        let probabilities = parse_probability_tables(&mut prefix.br)?;
792        let frame = finish_lossy_header(&prefix, probabilities.summary);
793        let partition_size_bytes = (prefix.token_partition_sizes.len() - 1) * 3;
794        let mut token_offset = prefix
795            .partition0_end
796            .checked_add(partition_size_bytes)
797            .ok_or(DecoderError::Bitstream(
798                "VP8 token partition offset overflow",
799            ))?;
800        let mut token_readers = Vec::with_capacity(prefix.token_partition_sizes.len());
801        for size in &prefix.token_partition_sizes {
802            let end = token_offset
803                .checked_add(*size)
804                .ok_or(DecoderError::Bitstream("VP8 token partition size overflow"))?;
805            let partition = data
806                .get(token_offset..end)
807                .ok_or(DecoderError::NotEnoughData("VP8 token partition"))?;
808            token_readers.push(Vp8BoolDecoder::new(partition));
809            token_offset = end;
810        }
811
812        Ok(Self {
813            top_modes: vec![B_DC_PRED; frame.macroblock_width * 4],
814            top_contexts: vec![NonZeroContext::default(); frame.macroblock_width],
815            frame,
816            br: prefix.br,
817            token_readers,
818            probabilities,
819            next_mb_y: 0,
820        })
821    }
822
823    pub(crate) fn frame(&self) -> &LossyHeader {
824        &self.frame
825    }
826
827    pub(crate) fn next_row(&mut self) -> Result<Option<Vec<MacroBlockData>>, DecoderError> {
828        if self.next_mb_y == self.frame.macroblock_height {
829            return Ok(None);
830        }
831
832        let mb_y = self.next_mb_y;
833        let mut left_modes = [B_DC_PRED; 4];
834        let row = parse_intra_mode_row(
835            &mut self.br,
836            self.frame.macroblock_width,
837            self.frame.segment.update_map,
838            &self.frame.segment.segment_probs,
839            self.probabilities.summary.use_skip_probability,
840            self.probabilities.summary.skip_probability.unwrap_or(0),
841            &mut self.top_modes,
842            &mut left_modes,
843        )?;
844
845        let part_mask = self.token_readers.len() - 1;
846        let token_br = &mut self.token_readers[mb_y & part_mask];
847        let mut left_context = NonZeroContext::default();
848        let mut macroblocks = Vec::with_capacity(self.frame.macroblock_width);
849        for (mb_x, header) in row.into_iter().enumerate() {
850            let reads_tokens = !header.skip;
851            let mb = parse_residuals(
852                header,
853                &mut self.top_contexts[mb_x],
854                &mut left_context,
855                token_br,
856                &self.frame.quantization,
857                &self.probabilities,
858            );
859            if reads_tokens && token_br.eof() {
860                return Err(DecoderError::NotEnoughData("VP8 token partition"));
861            }
862            macroblocks.push(mb);
863        }
864        self.next_mb_y += 1;
865        Ok(Some(macroblocks))
866    }
867}
868
869pub fn parse_macroblock_data(data: &[u8]) -> Result<MacroBlockDataFrame, DecoderError> {
870    let mut rows = MacroBlockRows::new(data)?;
871    let frame = rows.frame().clone();
872    let mut macroblocks = Vec::with_capacity(frame.macroblock_width * frame.macroblock_height);
873    while let Some(row) = rows.next_row()? {
874        macroblocks.extend(row);
875    }
876    Ok(MacroBlockDataFrame { frame, macroblocks })
877}
878
879#[cfg(test)]
880mod tests {
881    use super::{parse_token_partitions, Vp8BoolDecoder};
882    use crate::decoder::DecoderError;
883
884    #[test]
885    fn token_partition_sizes_must_fit_available_data() {
886        let mut br = Vp8BoolDecoder::new(&[0x40]);
887        assert_eq!(
888            parse_token_partitions(&mut br, &[3, 0, 0, 0xaa]),
889            Err(DecoderError::NotEnoughData("VP8 token partition"))
890        );
891
892        let mut br = Vp8BoolDecoder::new(&[0x40]);
893        assert_eq!(
894            parse_token_partitions(&mut br, &[3, 0, 0, 1, 2, 3, 4, 5]),
895            Ok(vec![3, 2])
896        );
897    }
898}