1#![forbid(unsafe_code)]
20#![allow(dead_code)]
21#![allow(clippy::doc_markdown)]
22#![allow(clippy::cast_possible_truncation)]
23#![allow(clippy::cast_sign_loss)]
24#![allow(clippy::cast_precision_loss)]
25#![allow(clippy::too_many_arguments)]
26
27use super::tile::TileInfo;
28use crate::error::{CodecError, CodecResult};
29use crate::frame::VideoFrame;
30use rayon::prelude::*;
31use std::sync::Arc;
32
33#[derive(Clone, Debug)]
39pub struct TileEncoderConfig {
40 pub tile_cols: u32,
42 pub tile_rows: u32,
44 pub threads: usize,
46 pub sb_size: u32,
48 pub uniform_spacing: bool,
50 pub min_tile_width_sb: u32,
52 pub max_tile_width_sb: u32,
54 pub min_tile_height_sb: u32,
56 pub max_tile_height_sb: u32,
58}
59
60impl Default for TileEncoderConfig {
61 fn default() -> Self {
62 Self {
63 tile_cols: 1,
64 tile_rows: 1,
65 threads: 0,
66 sb_size: 64,
67 uniform_spacing: true,
68 min_tile_width_sb: 1,
69 max_tile_width_sb: 64,
70 min_tile_height_sb: 1,
71 max_tile_height_sb: 64,
72 }
73 }
74}
75
76impl TileEncoderConfig {
77 #[must_use]
82 pub fn auto(width: u32, height: u32, threads: usize) -> Self {
83 let mut config = Self::default();
84 config.threads = threads;
85 config.configure_for_dimensions(width, height);
86 config
87 }
88
89 pub fn with_tile_counts(tile_cols: u32, tile_rows: u32, threads: usize) -> CodecResult<Self> {
95 if tile_cols == 0 || tile_rows == 0 {
96 return Err(CodecError::InvalidParameter(
97 "Tile counts must be positive".to_string(),
98 ));
99 }
100
101 if tile_cols > 64 || tile_rows > 64 {
102 return Err(CodecError::InvalidParameter(
103 "Maximum 64 tile columns/rows".to_string(),
104 ));
105 }
106
107 if !tile_cols.is_power_of_two() || !tile_rows.is_power_of_two() {
109 return Err(CodecError::InvalidParameter(
110 "Tile counts must be power of 2".to_string(),
111 ));
112 }
113
114 Ok(Self {
115 tile_cols,
116 tile_rows,
117 threads,
118 ..Default::default()
119 })
120 }
121
122 pub fn configure_for_dimensions(&mut self, width: u32, height: u32) {
124 let sb_cols = width.div_ceil(self.sb_size);
125 let sb_rows = height.div_ceil(self.sb_size);
126
127 let thread_count = if self.threads == 0 {
129 rayon::current_num_threads()
130 } else {
131 self.threads
132 };
133
134 let target_tiles = thread_count.next_power_of_two() as u32;
136
137 let aspect_ratio = width as f32 / height.max(1) as f32;
139
140 if aspect_ratio > 2.0 {
141 self.tile_cols = (target_tiles as f32).sqrt().ceil() as u32;
143 self.tile_cols = self.tile_cols.next_power_of_two();
144 self.tile_rows = (target_tiles / self.tile_cols).max(1);
145 self.tile_rows = self.tile_rows.next_power_of_two();
146 } else if aspect_ratio < 0.5 {
147 self.tile_rows = (target_tiles as f32).sqrt().ceil() as u32;
149 self.tile_rows = self.tile_rows.next_power_of_two();
150 self.tile_cols = (target_tiles / self.tile_rows).max(1);
151 self.tile_cols = self.tile_cols.next_power_of_two();
152 } else {
153 let sqrt_tiles = (target_tiles as f32).sqrt() as u32;
155 self.tile_cols = sqrt_tiles.next_power_of_two();
156 self.tile_rows = sqrt_tiles.next_power_of_two();
157 }
158
159 self.tile_cols = self.tile_cols.clamp(1, 64.min(sb_cols));
161 self.tile_rows = self.tile_rows.clamp(1, 64.min(sb_rows));
162
163 while self.tile_cols * self.tile_rows > 4096 {
165 if self.tile_cols > self.tile_rows {
166 self.tile_cols /= 2;
167 } else {
168 self.tile_rows /= 2;
169 }
170 }
171 }
172
173 #[must_use]
175 pub const fn tile_count(&self) -> u32 {
176 self.tile_cols * self.tile_rows
177 }
178
179 #[must_use]
181 pub fn thread_count(&self) -> usize {
182 if self.threads == 0 {
183 rayon::current_num_threads()
184 } else {
185 self.threads
186 }
187 }
188
189 pub fn validate(&self) -> CodecResult<()> {
195 if self.tile_cols == 0 || self.tile_rows == 0 {
196 return Err(CodecError::InvalidParameter(
197 "Tile counts must be positive".to_string(),
198 ));
199 }
200
201 if self.tile_cols > 64 || self.tile_rows > 64 {
202 return Err(CodecError::InvalidParameter(
203 "Maximum 64 tile columns/rows".to_string(),
204 ));
205 }
206
207 if self.tile_count() > 4096 {
208 return Err(CodecError::InvalidParameter(
209 "Maximum 4096 total tiles".to_string(),
210 ));
211 }
212
213 if self.sb_size != 64 && self.sb_size != 128 {
214 return Err(CodecError::InvalidParameter(
215 "Superblock size must be 64 or 128".to_string(),
216 ));
217 }
218
219 Ok(())
220 }
221}
222
223#[derive(Clone, Debug)]
229pub struct TileRegion {
230 pub col: u32,
232 pub row: u32,
234 pub x: u32,
236 pub y: u32,
238 pub width: u32,
240 pub height: u32,
242 pub index: u32,
244}
245
246impl TileRegion {
247 #[must_use]
249 pub const fn new(
250 col: u32,
251 row: u32,
252 x: u32,
253 y: u32,
254 width: u32,
255 height: u32,
256 tile_cols: u32,
257 ) -> Self {
258 Self {
259 col,
260 row,
261 x,
262 y,
263 width,
264 height,
265 index: row * tile_cols + col,
266 }
267 }
268
269 #[must_use]
271 pub const fn is_valid(&self) -> bool {
272 self.width > 0 && self.height > 0
273 }
274
275 #[must_use]
277 pub const fn area(&self) -> u32 {
278 self.width * self.height
279 }
280
281 #[must_use]
283 pub const fn is_left_edge(&self) -> bool {
284 self.col == 0
285 }
286
287 #[must_use]
289 pub const fn is_top_edge(&self) -> bool {
290 self.row == 0
291 }
292}
293
294#[derive(Clone, Debug)]
300pub struct TileFrameSplitter {
301 config: TileEncoderConfig,
303 frame_width: u32,
305 frame_height: u32,
307 regions: Vec<TileRegion>,
309}
310
311impl TileFrameSplitter {
312 pub fn new(
318 config: TileEncoderConfig,
319 frame_width: u32,
320 frame_height: u32,
321 ) -> CodecResult<Self> {
322 config.validate()?;
323
324 let mut splitter = Self {
325 config,
326 frame_width,
327 frame_height,
328 regions: Vec::new(),
329 };
330
331 splitter.compute_regions();
332 Ok(splitter)
333 }
334
335 fn compute_regions(&mut self) {
337 self.regions.clear();
338
339 if self.config.uniform_spacing {
340 self.compute_uniform_regions();
341 } else {
342 self.compute_custom_regions();
343 }
344 }
345
346 fn compute_uniform_regions(&mut self) {
348 let tile_width = self.frame_width.div_ceil(self.config.tile_cols);
349 let tile_height = self.frame_height.div_ceil(self.config.tile_rows);
350
351 for row in 0..self.config.tile_rows {
352 for col in 0..self.config.tile_cols {
353 let x = col * tile_width;
354 let y = row * tile_height;
355
356 let width = if col == self.config.tile_cols - 1 {
357 self.frame_width - x
358 } else {
359 tile_width
360 };
361
362 let height = if row == self.config.tile_rows - 1 {
363 self.frame_height - y
364 } else {
365 tile_height
366 };
367
368 self.regions.push(TileRegion::new(
369 col,
370 row,
371 x,
372 y,
373 width,
374 height,
375 self.config.tile_cols,
376 ));
377 }
378 }
379 }
380
381 fn compute_custom_regions(&mut self) {
383 self.compute_uniform_regions();
386 }
387
388 #[must_use]
390 pub fn regions(&self) -> &[TileRegion] {
391 &self.regions
392 }
393
394 #[must_use]
396 pub fn region(&self, index: usize) -> Option<&TileRegion> {
397 self.regions.get(index)
398 }
399
400 #[must_use]
402 pub fn tile_count(&self) -> usize {
403 self.regions.len()
404 }
405}
406
407#[derive(Clone, Debug)]
413pub struct TileEncoder {
414 region: TileRegion,
416 quality: u8,
418 is_keyframe: bool,
420}
421
422impl TileEncoder {
423 #[must_use]
425 pub const fn new(region: TileRegion, quality: u8, is_keyframe: bool) -> Self {
426 Self {
427 region,
428 quality,
429 is_keyframe,
430 }
431 }
432
433 pub fn encode(&self, frame: &VideoFrame) -> CodecResult<TileEncodedData> {
439 if self.region.x + self.region.width > frame.width
441 || self.region.y + self.region.height > frame.height
442 {
443 return Err(CodecError::InvalidParameter(
444 "Tile region exceeds frame bounds".to_string(),
445 ));
446 }
447
448 let tile_data = self.extract_tile_data(frame)?;
450
451 let encoded = self.encode_tile_data(&tile_data)?;
453
454 Ok(TileEncodedData {
455 region: self.region.clone(),
456 data: encoded,
457 size: 0, })
459 }
460
461 fn extract_tile_data(&self, frame: &VideoFrame) -> CodecResult<Vec<u8>> {
463 let mut tile_pixels = Vec::new();
464
465 if let Some(y_plane) = frame.planes.first() {
467 for y in self.region.y..(self.region.y + self.region.height) {
468 let row_start = (y as usize * y_plane.stride) + self.region.x as usize;
469 let row_end = row_start + self.region.width as usize;
470 if row_end <= y_plane.data.len() {
471 tile_pixels.extend_from_slice(&y_plane.data[row_start..row_end]);
472 }
473 }
474 }
475
476 let chroma_x = self.region.x / 2;
478 let chroma_y = self.region.y / 2;
479 let chroma_width = self.region.width / 2;
480 let chroma_height = self.region.height / 2;
481
482 for plane_idx in 1..frame.planes.len() {
483 if let Some(plane) = frame.planes.get(plane_idx) {
484 for y in chroma_y..(chroma_y + chroma_height) {
485 let row_start = (y as usize * plane.stride) + chroma_x as usize;
486 let row_end = row_start + chroma_width as usize;
487 if row_end <= plane.data.len() {
488 tile_pixels.extend_from_slice(&plane.data[row_start..row_end]);
489 }
490 }
491 }
492 }
493
494 Ok(tile_pixels)
495 }
496
497 fn encode_tile_data(&self, _tile_data: &[u8]) -> CodecResult<Vec<u8>> {
499 let mut encoded = Vec::new();
507
508 encoded.push(if self.is_keyframe { 0x80 } else { 0x00 });
510 encoded.push(self.quality);
511
512 encoded.extend_from_slice(&(self.region.width).to_le_bytes());
514 encoded.extend_from_slice(&(self.region.height).to_le_bytes());
515
516 let compressed_size = (self.region.width * self.region.height / 32) as usize;
518 encoded.resize(encoded.len() + compressed_size, 0);
519
520 Ok(encoded)
521 }
522
523 #[must_use]
525 pub const fn region(&self) -> &TileRegion {
526 &self.region
527 }
528}
529
530#[derive(Clone, Debug)]
536pub struct TileEncodedData {
537 pub region: TileRegion,
539 pub data: Vec<u8>,
541 pub size: usize,
543}
544
545impl TileEncodedData {
546 #[must_use]
548 pub const fn index(&self) -> u32 {
549 self.region.index
550 }
551
552 #[must_use]
554 pub fn encoded_size(&self) -> usize {
555 self.data.len()
556 }
557}
558
559#[derive(Debug)]
565pub struct ParallelTileEncoder {
566 config: Arc<TileEncoderConfig>,
568 splitter: TileFrameSplitter,
570 frame_width: u32,
572 frame_height: u32,
574}
575
576impl ParallelTileEncoder {
577 pub fn new(
583 config: TileEncoderConfig,
584 frame_width: u32,
585 frame_height: u32,
586 ) -> CodecResult<Self> {
587 let splitter = TileFrameSplitter::new(config.clone(), frame_width, frame_height)?;
588
589 Ok(Self {
590 config: Arc::new(config),
591 splitter,
592 frame_width,
593 frame_height,
594 })
595 }
596
597 pub fn encode_frame(
603 &self,
604 frame: &VideoFrame,
605 quality: u8,
606 is_keyframe: bool,
607 ) -> CodecResult<Vec<TileEncodedData>> {
608 if frame.width != self.frame_width || frame.height != self.frame_height {
610 return Err(CodecError::InvalidParameter(format!(
611 "Frame dimensions {}x{} don't match encoder {}x{}",
612 frame.width, frame.height, self.frame_width, self.frame_height
613 )));
614 }
615
616 if self.config.threads > 0 {
618 rayon::ThreadPoolBuilder::new()
619 .num_threads(self.config.threads)
620 .build()
621 .map_err(|e| {
622 CodecError::Internal(format!("Failed to create thread pool: {}", e))
623 })?;
624 }
625
626 let encoded_tiles: Vec<CodecResult<TileEncodedData>> = self
628 .splitter
629 .regions()
630 .par_iter()
631 .map(|region| {
632 let encoder = TileEncoder::new(region.clone(), quality, is_keyframe);
633 encoder.encode(frame)
634 })
635 .collect();
636
637 let mut tiles = Vec::with_capacity(encoded_tiles.len());
639 for result in encoded_tiles {
640 tiles.push(result?);
641 }
642
643 tiles.sort_by_key(TileEncodedData::index);
645
646 Ok(tiles)
647 }
648
649 pub fn merge_tiles(&self, tiles: &[TileEncodedData]) -> CodecResult<Vec<u8>> {
655 if tiles.is_empty() {
656 return Ok(Vec::new());
657 }
658
659 let mut merged = Vec::new();
660
661 self.write_tile_group_header(&mut merged, tiles.len() as u32);
663
664 for (i, tile) in tiles.iter().enumerate() {
666 let is_last = i == tiles.len() - 1;
667
668 if !is_last {
669 let size = tile.data.len() as u32;
671 self.write_tile_size(&mut merged, size);
672 }
673
674 merged.extend_from_slice(&tile.data);
676 }
677
678 Ok(merged)
679 }
680
681 fn write_tile_group_header(&self, output: &mut Vec<u8>, _num_tiles: u32) {
683 if self.config.tile_count() > 1 {
686 output.push(0x01); }
688 }
689
690 fn write_tile_size(&self, output: &mut Vec<u8>, size: u32) {
692 output.extend_from_slice(&size.to_le_bytes());
694 }
695
696 #[must_use]
698 pub fn config(&self) -> &TileEncoderConfig {
699 &self.config
700 }
701
702 #[must_use]
704 pub fn tile_count(&self) -> usize {
705 self.splitter.tile_count()
706 }
707
708 #[must_use]
710 pub fn regions(&self) -> &[TileRegion] {
711 self.splitter.regions()
712 }
713}
714
715pub struct TileInfoBuilder;
721
722impl TileInfoBuilder {
723 #[must_use]
725 pub fn from_config(
726 config: &TileEncoderConfig,
727 frame_width: u32,
728 frame_height: u32,
729 ) -> TileInfo {
730 let sb_cols = frame_width.div_ceil(config.sb_size);
731 let sb_rows = frame_height.div_ceil(config.sb_size);
732
733 let tile_width_sb = sb_cols.div_ceil(config.tile_cols);
734 let tile_height_sb = sb_rows.div_ceil(config.tile_rows);
735
736 let mut tile_col_starts = Vec::new();
738 for i in 0..=config.tile_cols {
739 let start = (i * tile_width_sb).min(sb_cols);
740 tile_col_starts.push(start);
741 }
742
743 let mut tile_row_starts = Vec::new();
745 for i in 0..=config.tile_rows {
746 let start = (i * tile_height_sb).min(sb_rows);
747 tile_row_starts.push(start);
748 }
749
750 let tile_cols_log2 = (config.tile_cols as f32).log2() as u8;
751 let tile_rows_log2 = (config.tile_rows as f32).log2() as u8;
752
753 TileInfo {
754 tile_cols: config.tile_cols,
755 tile_rows: config.tile_rows,
756 tile_col_starts,
757 tile_row_starts,
758 context_update_tile_id: 0,
759 tile_size_bytes: 4,
760 uniform_tile_spacing: config.uniform_spacing,
761 tile_cols_log2,
762 tile_rows_log2,
763 min_tile_cols_log2: 0,
764 max_tile_cols_log2: 6,
765 min_tile_rows_log2: 0,
766 max_tile_rows_log2: 6,
767 sb_cols,
768 sb_rows,
769 sb_size: config.sb_size,
770 }
771 }
772}
773
774#[cfg(test)]
779mod tests {
780 use super::*;
781 use oximedia_core::PixelFormat;
782
783 #[test]
784 fn test_tile_encoder_config_default() {
785 let config = TileEncoderConfig::default();
786 assert_eq!(config.tile_cols, 1);
787 assert_eq!(config.tile_rows, 1);
788 assert_eq!(config.tile_count(), 1);
789 }
790
791 #[test]
792 fn test_tile_encoder_config_auto() {
793 let config = TileEncoderConfig::auto(1920, 1080, 4);
794 assert!(config.tile_count() > 0);
795 assert!(config.tile_count() <= 4096);
796 }
797
798 #[test]
799 fn test_tile_encoder_config_manual() {
800 let config = TileEncoderConfig::with_tile_counts(2, 2, 4).expect("should succeed");
801 assert_eq!(config.tile_cols, 2);
802 assert_eq!(config.tile_rows, 2);
803 assert_eq!(config.tile_count(), 4);
804 }
805
806 #[test]
807 fn test_tile_encoder_config_validation() {
808 let config = TileEncoderConfig::default();
809 assert!(config.validate().is_ok());
810
811 let mut invalid = TileEncoderConfig::default();
812 invalid.tile_cols = 0;
813 assert!(invalid.validate().is_err());
814 }
815
816 #[test]
817 fn test_tile_region() {
818 let region = TileRegion::new(0, 0, 0, 0, 640, 480, 2);
819 assert!(region.is_valid());
820 assert_eq!(region.area(), 640 * 480);
821 assert!(region.is_left_edge());
822 assert!(region.is_top_edge());
823 }
824
825 #[test]
826 fn test_tile_frame_splitter() {
827 let config = TileEncoderConfig::with_tile_counts(2, 2, 4).expect("should succeed");
828 let splitter = TileFrameSplitter::new(config, 1920, 1080).expect("should succeed");
829
830 assert_eq!(splitter.tile_count(), 4);
831 assert_eq!(splitter.regions().len(), 4);
832
833 let region = splitter.region(0).expect("should succeed");
835 assert_eq!(region.col, 0);
836 assert_eq!(region.row, 0);
837 assert_eq!(region.x, 0);
838 assert_eq!(region.y, 0);
839 }
840
841 #[test]
842 fn test_tile_encoder() {
843 let region = TileRegion::new(0, 0, 0, 0, 320, 240, 1);
844 let encoder = TileEncoder::new(region, 128, true);
845
846 let mut frame = VideoFrame::new(PixelFormat::Yuv420p, 1920, 1080);
847 frame.allocate();
848
849 let result = encoder.encode(&frame);
850 assert!(result.is_ok());
851
852 let encoded = result.expect("should succeed");
853 assert!(encoded.encoded_size() > 0);
854 }
855
856 #[test]
857 fn test_parallel_tile_encoder() {
858 let config = TileEncoderConfig::with_tile_counts(2, 2, 4).expect("should succeed");
859 let encoder = ParallelTileEncoder::new(config, 1920, 1080).expect("should succeed");
860
861 assert_eq!(encoder.tile_count(), 4);
862 assert_eq!(encoder.regions().len(), 4);
863 }
864
865 #[test]
866 fn test_parallel_encode_frame() {
867 let config = TileEncoderConfig::with_tile_counts(2, 2, 4).expect("should succeed");
868 let encoder = ParallelTileEncoder::new(config, 1920, 1080).expect("should succeed");
869
870 let mut frame = VideoFrame::new(PixelFormat::Yuv420p, 1920, 1080);
871 frame.allocate();
872
873 let result = encoder.encode_frame(&frame, 128, true);
874 assert!(result.is_ok());
875
876 let tiles = result.expect("should succeed");
877 assert_eq!(tiles.len(), 4);
878 }
879
880 #[test]
881 fn test_merge_tiles() {
882 let config = TileEncoderConfig::with_tile_counts(2, 2, 4).expect("should succeed");
883 let encoder = ParallelTileEncoder::new(config, 1920, 1080).expect("should succeed");
884
885 let mut frame = VideoFrame::new(PixelFormat::Yuv420p, 1920, 1080);
886 frame.allocate();
887
888 let tiles = encoder
889 .encode_frame(&frame, 128, true)
890 .expect("should succeed");
891 let merged = encoder.merge_tiles(&tiles);
892
893 assert!(merged.is_ok());
894 assert!(!merged.expect("should succeed").is_empty());
895 }
896
897 #[test]
898 fn test_tile_info_builder() {
899 let config = TileEncoderConfig::with_tile_counts(2, 2, 4).expect("should succeed");
900 let tile_info = TileInfoBuilder::from_config(&config, 1920, 1080);
901
902 assert_eq!(tile_info.tile_cols, 2);
903 assert_eq!(tile_info.tile_rows, 2);
904 assert_eq!(tile_info.tile_count(), 4);
905 }
906
907 #[test]
908 fn test_aspect_ratio_configuration() {
909 let mut config = TileEncoderConfig::default();
911 config.configure_for_dimensions(3840, 1080);
912 assert!(config.tile_cols >= config.tile_rows);
913
914 let mut config = TileEncoderConfig::default();
916 config.configure_for_dimensions(1080, 3840);
917 assert!(config.tile_rows >= config.tile_cols);
918 }
919}