1use rayon::prelude::*;
8
9use crate::error::{CodecError, CodecResult};
10use crate::frame::{FrameType, Plane, VideoFrame};
11use crate::traits::VideoDecoder;
12use oximedia_core::{CodecId, PixelFormat, Rational, Timestamp};
13
14use super::crc32::crc32_mpeg2;
15use super::prediction::predict_median;
16use super::range_coder::SimpleRangeDecoder;
17use super::types::{
18 Ffv1ChromaType, Ffv1Colorspace, Ffv1Config, Ffv1Version, SliceHeader, CONTEXT_COUNT,
19 INITIAL_STATE,
20};
21
22fn decode_all_planes_in_slice(
33 data: &[u8],
34 plane_headers: &[SliceHeader],
35 plane_states: &mut Vec<Vec<u8>>,
36) -> CodecResult<Vec<Vec<Vec<i32>>>> {
37 let plane_count = plane_headers.len();
38
39 if data.len() < 2 {
40 let mut planes = Vec::with_capacity(plane_count);
42 for header in plane_headers {
43 let w = header.slice_width as usize;
44 let h = header.slice_height as usize;
45 let mut lines = Vec::with_capacity(h);
46 for _ in 0..h {
47 lines.push(vec![0i32; w]);
48 }
49 planes.push(lines);
50 }
51 return Ok(planes);
52 }
53
54 let mut decoder = SimpleRangeDecoder::new(data)?;
55 let mut planes_out: Vec<Vec<Vec<i32>>> = Vec::with_capacity(plane_count);
56
57 for (plane_idx, header) in plane_headers.iter().enumerate() {
58 let w = header.slice_width as usize;
59 let h = header.slice_height as usize;
60
61 let states = plane_states
62 .get_mut(plane_idx)
63 .ok_or_else(|| CodecError::Internal("invalid plane index".to_string()))?;
64
65 if w == 0 || h == 0 {
66 planes_out.push(Vec::new());
67 continue;
68 }
69
70 let mut lines: Vec<Vec<i32>> = Vec::with_capacity(h);
71 let mut prev_line = vec![0i32; w];
72
73 for _y in 0..h {
74 let mut line = Vec::with_capacity(w);
75 for x in 0..w {
76 let residual = decoder.get_symbol(states)?;
77 let left = if x > 0 { line[x - 1] } else { 0 };
78 let top = prev_line[x];
79 let top_left = if x > 0 { prev_line[x - 1] } else { 0 };
80 let pred = predict_median(left, top, top_left);
81 line.push(pred.saturating_add(residual));
84 }
85 prev_line.clone_from(&line);
86 lines.push(line);
87 }
88 planes_out.push(lines);
89 }
90
91 Ok(planes_out)
92}
93
94fn pixel_format_for_config(config: &Ffv1Config) -> PixelFormat {
96 match (
97 config.colorspace,
98 config.chroma_type,
99 config.bits_per_raw_sample,
100 ) {
101 (Ffv1Colorspace::YCbCr, Ffv1ChromaType::Chroma420, 8) => PixelFormat::Yuv420p,
102 (Ffv1Colorspace::YCbCr, Ffv1ChromaType::Chroma420, 10) => PixelFormat::Yuv420p10le,
103 (Ffv1Colorspace::YCbCr, Ffv1ChromaType::Chroma420, 12) => PixelFormat::Yuv420p12le,
104 (Ffv1Colorspace::YCbCr, Ffv1ChromaType::Chroma422, 8) => PixelFormat::Yuv422p,
105 (Ffv1Colorspace::YCbCr, Ffv1ChromaType::Chroma422, 10) => PixelFormat::Yuv422p10le,
106 (Ffv1Colorspace::YCbCr, Ffv1ChromaType::Chroma422, 12) => PixelFormat::Yuv422p12le,
107 (Ffv1Colorspace::YCbCr, Ffv1ChromaType::Chroma444, 8) => PixelFormat::Yuv444p,
108 (Ffv1Colorspace::YCbCr, Ffv1ChromaType::Chroma444, 10) => PixelFormat::Yuv444p10le,
109 (Ffv1Colorspace::YCbCr, Ffv1ChromaType::Chroma444, 12) => PixelFormat::Yuv444p12le,
110 (Ffv1Colorspace::YCbCr, Ffv1ChromaType::Chroma420, 16) => PixelFormat::Yuv420p16le,
111 (Ffv1Colorspace::YCbCr, Ffv1ChromaType::Chroma422, 16) => PixelFormat::Yuv422p16le,
112 (Ffv1Colorspace::YCbCr, Ffv1ChromaType::Chroma444, 16) => PixelFormat::Yuv444p16le,
113 _ => PixelFormat::Yuv420p, }
115}
116
117pub struct Ffv1Decoder {
135 config: Option<Ffv1Config>,
137 output_queue: Vec<VideoFrame>,
139 flushing: bool,
141 frame_count: u64,
143 plane_states: Vec<Vec<u8>>,
145}
146
147impl Ffv1Decoder {
148 pub fn new() -> Self {
150 Self {
151 config: None,
152 output_queue: Vec::new(),
153 flushing: false,
154 frame_count: 0,
155 plane_states: Vec::new(),
156 }
157 }
158
159 pub fn with_extradata(extradata: &[u8]) -> CodecResult<Self> {
161 let mut dec = Self::new();
162 dec.parse_config(extradata)?;
163 Ok(dec)
164 }
165
166 fn parse_config(&mut self, data: &[u8]) -> CodecResult<()> {
172 if data.len() < 16 {
177 return Err(CodecError::InvalidBitstream(format!(
178 "FFV1 config too short: {} bytes, need at least 16",
179 data.len()
180 )));
181 }
182
183 let version = Ffv1Version::from_u8(data[0])?;
184 let colorspace = Ffv1Colorspace::from_u8(data[1])?;
185 let h_shift = u32::from(data[2]);
186 let v_shift = u32::from(data[3]);
187 let chroma_type = Ffv1ChromaType::from_shifts(h_shift, v_shift)?;
188 let bits_per_raw_sample = data[4];
189 let ec = data[5] != 0;
190 let num_h_slices = u32::from(data[6]);
191 let num_v_slices = u32::from(data[7]);
192
193 let width_bytes: [u8; 4] = data[8..12]
195 .try_into()
196 .map_err(|_| CodecError::InvalidBitstream("bad width bytes".to_string()))?;
197 let height_bytes: [u8; 4] = data[12..16]
198 .try_into()
199 .map_err(|_| CodecError::InvalidBitstream("bad height bytes".to_string()))?;
200
201 let width = u32::from_le_bytes(width_bytes);
202 let height = u32::from_le_bytes(height_bytes);
203
204 let config = Ffv1Config {
205 version,
206 width,
207 height,
208 colorspace,
209 chroma_type,
210 bits_per_raw_sample,
211 num_h_slices,
212 num_v_slices,
213 ec,
214 range_coder_mode: version.uses_range_coder(),
215 state_transition_delta: Vec::new(),
216 };
217 config.validate()?;
218
219 self.init_states(&config);
220 self.config = Some(config);
221 Ok(())
222 }
223
224 fn init_states(&mut self, config: &Ffv1Config) {
226 let plane_count = config.plane_count();
227 self.plane_states.clear();
228 for _ in 0..plane_count {
229 self.plane_states.push(vec![INITIAL_STATE; CONTEXT_COUNT]);
230 }
231 }
232
233 fn reset_states(&mut self) {
235 for states in &mut self.plane_states {
236 for s in states.iter_mut() {
237 *s = INITIAL_STATE;
238 }
239 }
240 }
241
242 fn decode_frame(&mut self, data: &[u8], pts: i64) -> CodecResult<VideoFrame> {
244 let config = self
246 .config
247 .as_ref()
248 .ok_or_else(|| CodecError::DecoderError("FFV1 decoder not configured".to_string()))?;
249
250 let width = config.width;
251 let height = config.height;
252 let plane_count = config.plane_count();
253 let ec = config.ec;
254 let num_slices = config.num_slices();
255 let num_h_slices = config.num_h_slices;
256 let num_v_slices = config.num_v_slices;
257 let max_val = config.max_sample_value();
258 let bps = config.bits_per_raw_sample;
259 let bytes_per_sample = config.bytes_per_sample();
260 let pixel_format = pixel_format_for_config(config);
261
262 let plane_dims: Vec<(u32, u32)> = (0..plane_count)
263 .map(|i| config.plane_dimensions(i))
264 .collect();
265
266 let _ = config;
268
269 let is_keyframe = self.frame_count == 0;
270
271 if is_keyframe {
272 self.reset_states();
273 }
274
275 let mut frame = VideoFrame::new(pixel_format, width, height);
276 frame.timestamp = Timestamp::new(pts, Rational::new(1, 1000));
277 frame.frame_type = if is_keyframe {
278 FrameType::Key
279 } else {
280 FrameType::Inter
281 };
282
283 let planes_data = if num_slices == 1 {
285 self.decode_single_slice(data, ec, plane_count, &plane_dims)?
286 } else {
287 decode_multi_slice_parallel(
288 data,
289 ec,
290 num_slices,
291 num_h_slices,
292 num_v_slices,
293 plane_count,
294 &plane_dims,
295 )?
296 };
297
298 for (plane_idx, plane_lines) in planes_data.iter().enumerate() {
300 let (pw, ph) = plane_dims[plane_idx];
301 let stride_samples = pw as usize;
303 let mut plane_data = vec![0u8; stride_samples * ph as usize * bytes_per_sample];
304
305 for (y, line) in plane_lines.iter().enumerate() {
306 if y >= ph as usize {
307 break;
308 }
309 for (x, &sample) in line.iter().enumerate() {
310 if x >= pw as usize {
311 break;
312 }
313 let s = sample.clamp(0, max_val);
314 let out_idx = y * stride_samples + x;
315 if bps <= 8 {
316 plane_data[out_idx] = s as u8;
317 } else {
318 plane_data[out_idx * 2] = (s & 0xFF) as u8;
319 plane_data[out_idx * 2 + 1] = ((s >> 8) & 0xFF) as u8;
320 }
321 }
322 }
323
324 let stride_bytes = stride_samples * bytes_per_sample;
326 frame
327 .planes
328 .push(Plane::with_dimensions(plane_data, stride_bytes, pw, ph));
329 }
330
331 self.frame_count += 1;
332 Ok(frame)
333 }
334
335 fn decode_single_slice(
340 &mut self,
341 data: &[u8],
342 ec: bool,
343 plane_count: usize,
344 plane_dims: &[(u32, u32)],
345 ) -> CodecResult<Vec<Vec<Vec<i32>>>> {
346 let slice_data = if ec && data.len() >= 4 {
347 let payload = &data[..data.len() - 4];
348 let stored_crc_bytes: [u8; 4] = data[data.len() - 4..]
349 .try_into()
350 .map_err(|_| CodecError::InvalidBitstream("bad CRC bytes".to_string()))?;
351 let stored_crc = u32::from_le_bytes(stored_crc_bytes);
352 let computed_crc = crc32_mpeg2(payload);
353 if stored_crc != computed_crc {
354 return Err(CodecError::InvalidBitstream(format!(
355 "FFV1 slice CRC mismatch: stored={stored_crc:#010X}, computed={computed_crc:#010X}"
356 )));
357 }
358 payload
359 } else {
360 data
361 };
362
363 let headers: Vec<SliceHeader> = plane_dims
365 .iter()
366 .map(|&(pw, ph)| SliceHeader {
367 slice_x: 0,
368 slice_y: 0,
369 slice_width: pw,
370 slice_height: ph,
371 })
372 .collect();
373
374 decode_all_planes_in_slice(slice_data, &headers, &mut self.plane_states)
375 }
376}
377
378fn decode_multi_slice_parallel(
385 data: &[u8],
386 ec: bool,
387 num_slices: u32,
388 num_h_slices: u32,
389 num_v_slices: u32,
390 plane_count: usize,
391 plane_dims: &[(u32, u32)],
392) -> CodecResult<Vec<Vec<Vec<i32>>>> {
393 let slice_data_len = data.len() / (num_slices as usize);
394
395 let slice_descs: Vec<(u32, u32, usize, usize, usize)> = (0..num_v_slices)
397 .flat_map(|sy| {
398 (0..num_h_slices).map(move |sx| {
399 let slice_idx = (sy * num_h_slices + sx) as usize;
400 let start = slice_idx * slice_data_len;
401 let end = if slice_idx + 1 == num_slices as usize {
402 data.len()
403 } else {
404 start + slice_data_len
405 };
406 (sy, sx, slice_idx, start, end)
407 })
408 })
409 .collect();
410
411 let slice_results: Vec<Result<(usize, usize, Vec<Vec<Vec<i32>>>), CodecError>> = slice_descs
413 .par_iter()
414 .map(|&(sy, sx, _slice_idx, start, end)| {
415 let raw_segment = &data[start..end];
417 let slice_segment = if ec && raw_segment.len() >= 4 {
418 let payload = &raw_segment[..raw_segment.len() - 4];
419 let stored_bytes: [u8; 4] = raw_segment[raw_segment.len() - 4..]
420 .try_into()
421 .map_err(|_| {
422 CodecError::InvalidBitstream("bad CRC bytes in slice".to_string())
423 })?;
424 let stored_crc = u32::from_le_bytes(stored_bytes);
425 let computed_crc = crc32_mpeg2(payload);
426 if stored_crc != computed_crc {
427 return Err(CodecError::InvalidBitstream(format!(
428 "FFV1 multi-slice CRC mismatch: stored={stored_crc:#010X}, computed={computed_crc:#010X}"
429 )));
430 }
431 payload
432 } else {
433 raw_segment
434 };
435
436 let mut local_plane_states: Vec<Vec<u8>> = (0..plane_count)
438 .map(|_| vec![INITIAL_STATE; CONTEXT_COUNT])
439 .collect();
440
441 let headers: Vec<SliceHeader> = plane_dims
443 .iter()
444 .enumerate()
445 .map(|(plane_idx, &(pw, ph))| {
446 let slice_pw = pw / num_h_slices;
447 let slice_ph = ph / num_v_slices;
448 let actual_sw = if sx == num_h_slices - 1 {
449 pw - sx * slice_pw
450 } else {
451 slice_pw
452 };
453 let actual_sh = if sy == num_v_slices - 1 {
454 ph - sy * slice_ph
455 } else {
456 slice_ph
457 };
458 let _ = plane_idx; SliceHeader {
460 slice_x: sx * slice_pw,
461 slice_y: sy * slice_ph,
462 slice_width: actual_sw,
463 slice_height: actual_sh,
464 }
465 })
466 .collect();
467
468 let per_plane_rows = decode_all_planes_in_slice(
469 slice_segment,
470 &headers,
471 &mut local_plane_states,
472 )?;
473
474 Ok((sy as usize, sx as usize, per_plane_rows))
475 })
476 .collect();
477
478 let decoded_slices: Vec<(usize, usize, Vec<Vec<Vec<i32>>>)> = slice_results
480 .into_iter()
481 .collect::<Result<Vec<_>, CodecError>>(
482 )?;
483
484 let mut grid: std::collections::HashMap<(usize, usize), Vec<Vec<Vec<i32>>>> =
486 std::collections::HashMap::new();
487 let mut ordered = decoded_slices;
488 for (sy, sx, rows) in ordered.drain(..) {
489 grid.insert((sy, sx), rows);
490 }
491
492 let mut planes_data: Vec<Vec<Vec<i32>>> = (0..plane_count).map(|_| Vec::new()).collect();
494
495 for sy in 0..num_v_slices as usize {
496 for plane_idx in 0..plane_count {
497 let mut plane_band: Vec<Vec<i32>> = Vec::new();
498
499 let (pw, ph) = plane_dims[plane_idx];
500 let slice_ph = ph as usize / num_v_slices as usize;
501 let actual_sh = if sy == num_v_slices as usize - 1 {
502 ph as usize - sy * slice_ph
503 } else {
504 slice_ph
505 };
506
507 for row_in_band in 0..actual_sh {
508 let mut full_row: Vec<i32> = Vec::with_capacity(pw as usize);
510 for sx in 0..num_h_slices as usize {
511 let slice_pw = pw as usize / num_h_slices as usize;
512 let actual_sw = if sx == num_h_slices as usize - 1 {
513 pw as usize - sx * slice_pw
514 } else {
515 slice_pw
516 };
517 let slice_rows = grid.get(&(sy, sx)).ok_or_else(|| {
518 CodecError::Internal(format!("missing slice ({sy}, {sx})"))
519 })?;
520 let slice_plane_rows = slice_rows.get(plane_idx).ok_or_else(|| {
521 CodecError::Internal(format!(
522 "missing plane {plane_idx} in slice ({sy}, {sx})"
523 ))
524 })?;
525 let row = slice_plane_rows.get(row_in_band).ok_or_else(|| {
526 CodecError::Internal(format!(
527 "missing row {row_in_band} in plane {plane_idx} slice ({sy}, {sx})"
528 ))
529 })?;
530 let take = actual_sw.min(row.len());
532 full_row.extend_from_slice(&row[..take]);
533 }
534 plane_band.push(full_row);
535 }
536 planes_data[plane_idx].extend(plane_band);
537 }
538 }
539
540 Ok(planes_data)
541}
542
543impl VideoDecoder for Ffv1Decoder {
544 fn codec(&self) -> CodecId {
545 CodecId::Ffv1
546 }
547
548 fn send_packet(&mut self, data: &[u8], pts: i64) -> CodecResult<()> {
549 if self.flushing {
550 return Err(CodecError::DecoderError(
551 "decoder is flushing, cannot accept new packets".to_string(),
552 ));
553 }
554
555 if self.config.is_none() {
556 return Err(CodecError::DecoderError(
557 "FFV1 decoder not configured: call with_extradata() first".to_string(),
558 ));
559 }
560
561 let frame = self.decode_frame(data, pts)?;
562 self.output_queue.push(frame);
563 Ok(())
564 }
565
566 fn receive_frame(&mut self) -> CodecResult<Option<VideoFrame>> {
567 if self.output_queue.is_empty() {
568 Ok(None)
569 } else {
570 Ok(Some(self.output_queue.remove(0)))
571 }
572 }
573
574 fn flush(&mut self) -> CodecResult<()> {
575 self.flushing = true;
576 Ok(())
577 }
578
579 fn reset(&mut self) {
580 self.output_queue.clear();
581 self.flushing = false;
582 self.frame_count = 0;
583 self.reset_states();
584 }
585
586 fn output_format(&self) -> Option<PixelFormat> {
587 self.config.as_ref().map(pixel_format_for_config)
588 }
589
590 fn dimensions(&self) -> Option<(u32, u32)> {
591 self.config.as_ref().map(|c| (c.width, c.height))
592 }
593}
594
595#[cfg(test)]
596mod tests {
597 use super::*;
598 use crate::traits::VideoDecoder;
599
600 fn make_config_bytes(width: u32, height: u32, bits: u8, h_shift: u8, v_shift: u8) -> Vec<u8> {
601 let mut data = Vec::new();
602 data.push(3); data.push(0); data.push(h_shift);
605 data.push(v_shift);
606 data.push(bits);
607 data.push(1); data.push(1); data.push(1); data.extend_from_slice(&width.to_le_bytes());
611 data.extend_from_slice(&height.to_le_bytes());
612 data
613 }
614
615 fn make_config_bytes_420_8(width: u32, height: u32) -> Vec<u8> {
616 make_config_bytes(width, height, 8, 1, 1)
617 }
618
619 #[test]
620 #[ignore]
621 fn test_decoder_creation() {
622 let dec = Ffv1Decoder::new();
623 assert!(dec.config.is_none());
624 assert_eq!(dec.codec(), CodecId::Ffv1);
625 }
626
627 #[test]
628 #[ignore]
629 fn test_decoder_with_extradata() {
630 let config_data = make_config_bytes_420_8(320, 240);
631 let dec = Ffv1Decoder::with_extradata(&config_data).expect("valid config");
632 assert!(dec.config.is_some());
633 assert_eq!(dec.dimensions(), Some((320, 240)));
634 assert_eq!(dec.output_format(), Some(PixelFormat::Yuv420p));
635 }
636
637 #[test]
638 #[ignore]
639 fn test_decoder_invalid_config() {
640 assert!(Ffv1Decoder::with_extradata(&[0; 5]).is_err());
642 let mut bad = make_config_bytes_420_8(320, 240);
644 bad[0] = 99;
645 assert!(Ffv1Decoder::with_extradata(&bad).is_err());
646 }
647
648 #[test]
649 #[ignore]
650 fn test_decoder_not_configured() {
651 let mut dec = Ffv1Decoder::new();
652 assert!(dec.send_packet(&[0; 100], 0).is_err());
653 }
654
655 #[test]
656 #[ignore]
657 fn test_decoder_reset() {
658 let config_data = make_config_bytes_420_8(16, 16);
659 let mut dec = Ffv1Decoder::with_extradata(&config_data).expect("valid");
660 dec.frame_count = 10;
661 dec.flushing = true;
662 dec.reset();
663 assert_eq!(dec.frame_count, 0);
664 assert!(!dec.flushing);
665 }
666
667 #[test]
668 #[ignore]
669 fn test_decoder_flush() {
670 let config_data = make_config_bytes_420_8(16, 16);
671 let mut dec = Ffv1Decoder::with_extradata(&config_data).expect("valid");
672 dec.flush().expect("flush ok");
673 assert!(dec.flushing);
674 assert!(dec.send_packet(&[0; 100], 0).is_err());
676 }
677
678 #[test]
679 fn test_pixel_format_dispatch_10bit() {
680 let config = Ffv1Config {
681 width: 16,
682 height: 16,
683 bits_per_raw_sample: 10,
684 chroma_type: Ffv1ChromaType::Chroma420,
685 colorspace: Ffv1Colorspace::YCbCr,
686 ..Default::default()
687 };
688 let dec = Ffv1Decoder {
689 config: Some(config),
690 output_queue: Vec::new(),
691 flushing: false,
692 frame_count: 0,
693 plane_states: Vec::new(),
694 };
695 assert_eq!(dec.output_format(), Some(PixelFormat::Yuv420p10le));
696 }
697
698 #[test]
699 fn test_pixel_format_dispatch_12bit_444() {
700 let config = Ffv1Config {
701 width: 16,
702 height: 16,
703 bits_per_raw_sample: 12,
704 chroma_type: Ffv1ChromaType::Chroma444,
705 colorspace: Ffv1Colorspace::YCbCr,
706 ..Default::default()
707 };
708 let dec = Ffv1Decoder {
709 config: Some(config),
710 output_queue: Vec::new(),
711 flushing: false,
712 frame_count: 0,
713 plane_states: Vec::new(),
714 };
715 assert_eq!(dec.output_format(), Some(PixelFormat::Yuv444p12le));
716 }
717}