Skip to main content

oximedia_codec/
tile_encoder.rs

1//! Tile-based parallel frame encoding for OxiMedia codecs.
2//!
3//! This module provides pixel-level infrastructure for splitting raw video
4//! frames into rectangular tiles, processing them concurrently, and
5//! reassembling the result into a complete frame.
6//!
7//! Unlike [`crate::tile`] which works with [`crate::frame::VideoFrame`] and
8//! codec-level bitstream output, this module operates on raw `&[u8]` / `Vec<u8>`
9//! pixel buffers and is therefore codec-agnostic.
10//!
11//! # Architecture
12//!
13//! ```text
14//! TileConfig  ─── tile grid parameters (cols, rows, frame size)
15//!     │
16//!     ▼
17//! TileLayout  ─── pre-computed TileRegion grid (handles remainder pixels)
18//!     │
19//!     ▼
20//! ParallelTileEncoder ─── split_frame → parallel encode_fn → merge_tiles
21//! ```
22//!
23//! # Example
24//!
25//! ```
26//! use oximedia_codec::tile_encoder::{TileConfig, ParallelTileEncoder};
27//!
28//! let config = TileConfig::new()
29//!     .tile_cols(2)
30//!     .tile_rows(2)
31//!     .frame_width(64)
32//!     .frame_height(64);
33//!
34//! let encoder = ParallelTileEncoder::new(config);
35//!
36//! // Create a simple 64×64 RGB frame (3 channels).
37//! let frame: Vec<u8> = (0u8..=255).cycle().take(64 * 64 * 3).collect();
38//!
39//! let tiles = encoder.split_frame(&frame, 3);
40//! assert_eq!(tiles.len(), 4);
41//!
42//! let merged = ParallelTileEncoder::merge_tiles(&tiles, 64, 64, 3);
43//! assert_eq!(merged, frame);
44//! ```
45
46use rayon::prelude::*;
47use std::ops::Range;
48
49// =============================================================================
50// TileConfig
51// =============================================================================
52
53/// Configuration for the tile grid and frame dimensions.
54///
55/// Use the builder-pattern methods to configure:
56///
57/// ```
58/// use oximedia_codec::tile_encoder::TileConfig;
59///
60/// let cfg = TileConfig::new()
61///     .tile_cols(4)
62///     .tile_rows(4)
63///     .num_threads(8)
64///     .frame_width(1920)
65///     .frame_height(1080);
66///
67/// assert_eq!(cfg.tile_cols, 4);
68/// assert_eq!(cfg.num_threads, 8);
69/// ```
70#[derive(Clone, Debug, PartialEq, Eq)]
71pub struct TileConfig {
72    /// Number of tile columns (1–64).
73    pub tile_cols: u32,
74    /// Number of tile rows (1–64).
75    pub tile_rows: u32,
76    /// Worker threads for parallel encoding (0 = use Rayon pool size).
77    pub num_threads: usize,
78    /// Frame width in pixels.
79    pub frame_width: u32,
80    /// Frame height in pixels.
81    pub frame_height: u32,
82}
83
84impl TileConfig {
85    /// Create a `TileConfig` with default values.
86    ///
87    /// Defaults: 1 column, 1 row, 0 threads (auto), 0×0 frame.
88    #[must_use]
89    pub fn new() -> Self {
90        Self::default()
91    }
92
93    /// Set the number of tile columns (1–64).
94    #[must_use]
95    pub fn tile_cols(mut self, cols: u32) -> Self {
96        self.tile_cols = cols.clamp(1, 64);
97        self
98    }
99
100    /// Set the number of tile rows (1–64).
101    #[must_use]
102    pub fn tile_rows(mut self, rows: u32) -> Self {
103        self.tile_rows = rows.clamp(1, 64);
104        self
105    }
106
107    /// Set the worker thread count (0 = Rayon auto).
108    #[must_use]
109    pub fn num_threads(mut self, threads: usize) -> Self {
110        self.num_threads = threads;
111        self
112    }
113
114    /// Set the frame width in pixels.
115    #[must_use]
116    pub fn frame_width(mut self, width: u32) -> Self {
117        self.frame_width = width;
118        self
119    }
120
121    /// Set the frame height in pixels.
122    #[must_use]
123    pub fn frame_height(mut self, height: u32) -> Self {
124        self.frame_height = height;
125        self
126    }
127
128    /// Effective thread count (resolves 0 to the Rayon pool size).
129    #[must_use]
130    pub fn thread_count(&self) -> usize {
131        if self.num_threads == 0 {
132            rayon::current_num_threads()
133        } else {
134            self.num_threads
135        }
136    }
137}
138
139impl Default for TileConfig {
140    fn default() -> Self {
141        Self {
142            tile_cols: 1,
143            tile_rows: 1,
144            num_threads: 0,
145            frame_width: 0,
146            frame_height: 0,
147        }
148    }
149}
150
151// =============================================================================
152// TileRegion
153// =============================================================================
154
155/// Pixel coordinates and dimensions of a single tile within a frame.
156///
157/// ```
158/// use oximedia_codec::tile_encoder::TileRegion;
159///
160/// let region = TileRegion::new(1, 0, 512, 0, 512, 288);
161/// assert_eq!(region.area(), 512 * 288);
162/// assert!(region.contains(600, 100));
163/// assert!(!region.contains(200, 100)); // left of tile
164/// ```
165#[derive(Clone, Debug, PartialEq, Eq)]
166pub struct TileRegion {
167    /// Tile column index (0-based).
168    pub col: u32,
169    /// Tile row index (0-based).
170    pub row: u32,
171    /// X pixel offset from the left of the frame.
172    pub x: u32,
173    /// Y pixel offset from the top of the frame.
174    pub y: u32,
175    /// Tile width in pixels.
176    pub width: u32,
177    /// Tile height in pixels.
178    pub height: u32,
179}
180
181impl TileRegion {
182    /// Create a new `TileRegion`.
183    #[must_use]
184    pub const fn new(col: u32, row: u32, x: u32, y: u32, width: u32, height: u32) -> Self {
185        Self {
186            col,
187            row,
188            x,
189            y,
190            width,
191            height,
192        }
193    }
194
195    /// Area of this tile in pixels.
196    #[must_use]
197    pub const fn area(&self) -> u64 {
198        self.width as u64 * self.height as u64
199    }
200
201    /// Returns `true` if the pixel `(px, py)` falls within this tile.
202    #[must_use]
203    pub const fn contains(&self, px: u32, py: u32) -> bool {
204        px >= self.x && px < self.x + self.width && py >= self.y && py < self.y + self.height
205    }
206
207    /// Pixel column range `x..(x + width)`.
208    #[must_use]
209    pub fn pixel_range_x(&self) -> Range<u32> {
210        self.x..(self.x + self.width)
211    }
212
213    /// Pixel row range `y..(y + height)`.
214    #[must_use]
215    pub fn pixel_range_y(&self) -> Range<u32> {
216        self.y..(self.y + self.height)
217    }
218}
219
220// =============================================================================
221// TileLayout
222// =============================================================================
223
224/// A grid of [`TileRegion`]s computed from a [`TileConfig`].
225///
226/// The last tile in each row/column absorbs any remainder pixels so the union
227/// of all tiles exactly covers the full frame with no overlap.
228///
229/// ```
230/// use oximedia_codec::tile_encoder::{TileConfig, TileLayout};
231///
232/// let cfg = TileConfig::new()
233///     .tile_cols(2)
234///     .tile_rows(2)
235///     .frame_width(100)
236///     .frame_height(100);
237///
238/// let layout = TileLayout::new(cfg);
239/// assert_eq!(layout.tile_count(), 4);
240///
241/// // All tiles together cover 100×100 pixels.
242/// let total: u64 = layout.tiles().iter().map(|t| t.area()).sum();
243/// assert_eq!(total, 100 * 100);
244/// ```
245#[derive(Clone, Debug)]
246pub struct TileLayout {
247    /// The configuration used to build this layout.
248    pub config: TileConfig,
249    /// All tile regions in raster order (row-major).
250    pub tiles: Vec<TileRegion>,
251}
252
253impl TileLayout {
254    /// Compute a `TileLayout` from `config`.
255    ///
256    /// Tile boundaries are computed as `frame_width / tile_cols` (integer
257    /// division); the last column and last row absorb the remainder pixels.
258    #[must_use]
259    pub fn new(config: TileConfig) -> Self {
260        let cols = config.tile_cols.max(1);
261        let rows = config.tile_rows.max(1);
262        let fw = config.frame_width;
263        let fh = config.frame_height;
264
265        // Nominal tile sizes (last tile gets the remainder).
266        let nominal_tw = fw / cols;
267        let nominal_th = fh / rows;
268
269        let mut tiles = Vec::with_capacity((cols * rows) as usize);
270
271        for row in 0..rows {
272            for col in 0..cols {
273                let x = col * nominal_tw;
274                let y = row * nominal_th;
275
276                let width = if col == cols - 1 {
277                    fw.saturating_sub(x)
278                } else {
279                    nominal_tw
280                };
281                let height = if row == rows - 1 {
282                    fh.saturating_sub(y)
283                } else {
284                    nominal_th
285                };
286
287                tiles.push(TileRegion::new(col, row, x, y, width, height));
288            }
289        }
290
291        Self { config, tiles }
292    }
293
294    /// Total number of tiles.
295    #[must_use]
296    pub fn tile_count(&self) -> usize {
297        self.tiles.len()
298    }
299
300    /// Return the tile at grid position `(col, row)`, or `None` if out of bounds.
301    #[must_use]
302    pub fn get_tile(&self, col: u32, row: u32) -> Option<&TileRegion> {
303        let cols = self.config.tile_cols;
304        let rows = self.config.tile_rows;
305        if col >= cols || row >= rows {
306            return None;
307        }
308        self.tiles.get((row * cols + col) as usize)
309    }
310
311    /// All tile regions in raster order.
312    #[must_use]
313    pub fn tiles(&self) -> &[TileRegion] {
314        &self.tiles
315    }
316
317    /// Find which tile contains the pixel `(px, py)`.
318    ///
319    /// Returns `None` if the pixel is outside the frame.
320    #[must_use]
321    pub fn tile_for_pixel(&self, px: u32, py: u32) -> Option<&TileRegion> {
322        self.tiles.iter().find(|t| t.contains(px, py))
323    }
324}
325
326// =============================================================================
327// TileBuffer
328// =============================================================================
329
330/// Raw pixel data extracted from (or destined for) a single tile.
331///
332/// ```
333/// use oximedia_codec::tile_encoder::{TileRegion, TileBuffer};
334///
335/// let region = TileRegion::new(0, 0, 0, 0, 4, 4);
336/// let buf = TileBuffer::new(region, 3); // 3 channels (RGB)
337/// assert_eq!(buf.data.len(), 4 * 4 * 3);
338/// assert_eq!(buf.stride, 4 * 3);
339/// ```
340#[derive(Clone, Debug)]
341pub struct TileBuffer {
342    /// Spatial position of this tile in the frame.
343    pub region: TileRegion,
344    /// Raw pixel bytes for this tile (row-major, tightly packed).
345    pub data: Vec<u8>,
346    /// Row stride in bytes (`width * channels`).
347    pub stride: usize,
348    /// Bytes per pixel.
349    pub channels: u8,
350}
351
352impl TileBuffer {
353    /// Allocate an all-zero `TileBuffer` for `region` with `channels` bytes per pixel.
354    #[must_use]
355    pub fn new(region: TileRegion, channels: u8) -> Self {
356        let ch = channels as usize;
357        let stride = region.width as usize * ch;
358        let data = vec![0u8; region.height as usize * stride];
359        Self {
360            region,
361            data,
362            stride,
363            channels,
364        }
365    }
366
367    /// Copy the tile's pixels from `frame` (a packed, row-major buffer).
368    ///
369    /// `frame_stride` is the number of bytes per row in the full frame
370    /// (i.e. `frame_width * channels`).
371    pub fn extract_from_frame(&mut self, frame: &[u8], frame_stride: usize) {
372        let ch = self.channels as usize;
373        let x_byte = self.region.x as usize * ch;
374        let w_bytes = self.region.width as usize * ch;
375
376        for row in 0..self.region.height as usize {
377            let frame_row_start = (self.region.y as usize + row) * frame_stride + x_byte;
378            let tile_row_start = row * self.stride;
379
380            let src_end = (frame_row_start + w_bytes).min(frame.len());
381            let copy_len = src_end.saturating_sub(frame_row_start);
382
383            self.data[tile_row_start..tile_row_start + copy_len]
384                .copy_from_slice(&frame[frame_row_start..src_end]);
385        }
386    }
387
388    /// Write this tile's pixels back into `frame`.
389    ///
390    /// `frame_stride` must match the full frame's row stride.
391    pub fn write_to_frame(&self, frame: &mut [u8], frame_stride: usize) {
392        let ch = self.channels as usize;
393        let x_byte = self.region.x as usize * ch;
394        let w_bytes = self.region.width as usize * ch;
395
396        for row in 0..self.region.height as usize {
397            let frame_row_start = (self.region.y as usize + row) * frame_stride + x_byte;
398            let tile_row_start = row * self.stride;
399
400            let dst_end = (frame_row_start + w_bytes).min(frame.len());
401            let copy_len = dst_end.saturating_sub(frame_row_start);
402
403            frame[frame_row_start..frame_row_start + copy_len]
404                .copy_from_slice(&self.data[tile_row_start..tile_row_start + copy_len]);
405        }
406    }
407}
408
409// =============================================================================
410// ParallelTileEncoder
411// =============================================================================
412
413/// Splits a raw pixel frame into tiles, processes them in parallel, and
414/// reassembles the result.
415///
416/// # Example
417///
418/// ```
419/// use oximedia_codec::tile_encoder::{TileConfig, ParallelTileEncoder};
420///
421/// let config = TileConfig::new()
422///     .tile_cols(2)
423///     .tile_rows(2)
424///     .frame_width(64)
425///     .frame_height(64);
426///
427/// let encoder = ParallelTileEncoder::new(config);
428///
429/// let frame: Vec<u8> = (0u8..=255).cycle().take(64 * 64 * 3).collect();
430/// let tiles = encoder.split_frame(&frame, 3);
431/// assert_eq!(tiles.len(), 4);
432///
433/// // Identity encode: return each tile unchanged.
434/// let processed = encoder
435///     .encode_tiles_parallel(tiles, |tile| Ok(tile))
436///     ?;
437///
438/// let merged = ParallelTileEncoder::merge_tiles(&processed, 64, 64, 3);
439/// assert_eq!(merged, frame);
440/// # Ok::<(), Box<dyn std::error::Error>>(())
441/// ```
442pub struct ParallelTileEncoder {
443    /// Pre-computed tile layout.
444    pub layout: TileLayout,
445}
446
447impl ParallelTileEncoder {
448    /// Create a `ParallelTileEncoder` from `config`.
449    #[must_use]
450    pub fn new(config: TileConfig) -> Self {
451        Self {
452            layout: TileLayout::new(config),
453        }
454    }
455
456    /// Split `frame` into [`TileBuffer`]s, one per tile in the layout.
457    ///
458    /// `channels` is the number of bytes per pixel in `frame`.
459    #[must_use]
460    pub fn split_frame(&self, frame: &[u8], channels: u8) -> Vec<TileBuffer> {
461        let fw = self.layout.config.frame_width;
462        let frame_stride = fw as usize * channels as usize;
463
464        self.layout
465            .tiles
466            .iter()
467            .map(|region| {
468                let mut buf = TileBuffer::new(region.clone(), channels);
469                buf.extract_from_frame(frame, frame_stride);
470                buf
471            })
472            .collect()
473    }
474
475    /// Merge a slice of [`TileBuffer`]s back into a complete frame.
476    ///
477    /// The returned `Vec<u8>` has `frame_width * frame_height * channels` bytes.
478    #[must_use]
479    pub fn merge_tiles(
480        tiles: &[TileBuffer],
481        frame_width: u32,
482        frame_height: u32,
483        channels: u8,
484    ) -> Vec<u8> {
485        let ch = channels as usize;
486        let frame_stride = frame_width as usize * ch;
487        let frame_size = frame_height as usize * frame_stride;
488        let mut frame = vec![0u8; frame_size];
489
490        for tile in tiles {
491            tile.write_to_frame(&mut frame, frame_stride);
492        }
493
494        frame
495    }
496
497    /// Process `tiles` in parallel using `encode_fn`.
498    ///
499    /// Each tile is passed by value to `encode_fn`.  The closure must return
500    /// either a (possibly modified) [`TileBuffer`] or an error string.
501    ///
502    /// Uses Rayon for parallel execution.  The output order matches the input
503    /// order (raster order when produced by `split_frame`).
504    ///
505    /// # Errors
506    ///
507    /// Returns the first error string produced by any invocation of
508    /// `encode_fn`.
509    pub fn encode_tiles_parallel<F>(
510        &self,
511        tiles: Vec<TileBuffer>,
512        encode_fn: F,
513    ) -> Result<Vec<TileBuffer>, String>
514    where
515        F: Fn(TileBuffer) -> Result<TileBuffer, String> + Send + Sync,
516    {
517        let results: Vec<Result<TileBuffer, String>> =
518            tiles.into_par_iter().map(|tile| encode_fn(tile)).collect();
519
520        let mut out = Vec::with_capacity(results.len());
521        for r in results {
522            out.push(r?);
523        }
524        Ok(out)
525    }
526}
527
528// =============================================================================
529// Adaptive Tile Partitioning
530// =============================================================================
531
532/// Content complexity metric for a tile region.
533#[derive(Clone, Debug, PartialEq)]
534pub struct TileComplexity {
535    /// Tile column index.
536    pub col: u32,
537    /// Tile row index.
538    pub row: u32,
539    /// Variance of pixel values (higher = more complex).
540    pub variance: f64,
541    /// Mean absolute difference between adjacent pixels (edge density).
542    pub edge_density: f64,
543    /// Normalised complexity score in [0.0, 1.0].
544    pub score: f64,
545}
546
547/// Analyse content complexity for each tile in a frame.
548///
549/// `frame` is a packed row-major pixel buffer with `channels` bytes per pixel.
550/// Returns a [`TileComplexity`] for every tile in `layout`.
551pub fn analyse_tile_complexity(
552    layout: &TileLayout,
553    frame: &[u8],
554    channels: u8,
555) -> Vec<TileComplexity> {
556    let fw = layout.config.frame_width;
557    let frame_stride = fw as usize * channels as usize;
558
559    let complexities: Vec<TileComplexity> = layout
560        .tiles
561        .iter()
562        .map(|region| {
563            let ch = channels as usize;
564            let w = region.width as usize;
565            let h = region.height as usize;
566            let n = (w * h) as f64;
567
568            if n < 1.0 {
569                return TileComplexity {
570                    col: region.col,
571                    row: region.row,
572                    variance: 0.0,
573                    edge_density: 0.0,
574                    score: 0.0,
575                };
576            }
577
578            // Compute mean and variance of luma (average of channels).
579            let mut sum: f64 = 0.0;
580            let mut sum_sq: f64 = 0.0;
581            let mut edge_sum: f64 = 0.0;
582            let mut edge_count: u64 = 0;
583
584            for row_idx in 0..h {
585                let frame_y = region.y as usize + row_idx;
586                for col_idx in 0..w {
587                    let frame_x = region.x as usize + col_idx;
588                    let base = frame_y * frame_stride + frame_x * ch;
589
590                    // Average across channels for luma approximation.
591                    let mut pixel_sum: u32 = 0;
592                    for c in 0..ch.min(frame.len().saturating_sub(base)) {
593                        pixel_sum += frame[base + c] as u32;
594                    }
595                    let luma = pixel_sum as f64 / ch.max(1) as f64;
596                    sum += luma;
597                    sum_sq += luma * luma;
598
599                    // Horizontal edge detection.
600                    if col_idx + 1 < w {
601                        let next_base = base + ch;
602                        let mut next_sum: u32 = 0;
603                        for c in 0..ch.min(frame.len().saturating_sub(next_base)) {
604                            next_sum += frame[next_base + c] as u32;
605                        }
606                        let next_luma = next_sum as f64 / ch.max(1) as f64;
607                        edge_sum += (luma - next_luma).abs();
608                        edge_count += 1;
609                    }
610
611                    // Vertical edge detection.
612                    if row_idx + 1 < h {
613                        let below_base = (frame_y + 1) * frame_stride + frame_x * ch;
614                        let mut below_sum: u32 = 0;
615                        for c in 0..ch.min(frame.len().saturating_sub(below_base)) {
616                            below_sum += frame[below_base + c] as u32;
617                        }
618                        let below_luma = below_sum as f64 / ch.max(1) as f64;
619                        edge_sum += (luma - below_luma).abs();
620                        edge_count += 1;
621                    }
622                }
623            }
624
625            let mean = sum / n;
626            let variance = (sum_sq / n) - (mean * mean);
627            let edge_density = if edge_count > 0 {
628                edge_sum / edge_count as f64
629            } else {
630                0.0
631            };
632
633            TileComplexity {
634                col: region.col,
635                row: region.row,
636                variance: variance.max(0.0),
637                edge_density,
638                score: 0.0, // filled in below
639            }
640        })
641        .collect();
642
643    // Normalise scores to [0.0, 1.0].
644    let max_var = complexities
645        .iter()
646        .map(|c| c.variance)
647        .fold(0.0_f64, f64::max);
648    let max_edge = complexities
649        .iter()
650        .map(|c| c.edge_density)
651        .fold(0.0_f64, f64::max);
652
653    complexities
654        .into_iter()
655        .map(|mut c| {
656            let norm_var = if max_var > 0.0 {
657                c.variance / max_var
658            } else {
659                0.0
660            };
661            let norm_edge = if max_edge > 0.0 {
662                c.edge_density / max_edge
663            } else {
664                0.0
665            };
666            c.score = (0.6 * norm_var + 0.4 * norm_edge).clamp(0.0, 1.0);
667            c
668        })
669        .collect()
670}
671
672/// Decides whether a tile should be split into sub-tiles based on complexity.
673///
674/// Returns a suggested partition: `(sub_cols, sub_rows)` for each tile.
675/// Simple tiles get `(1,1)`, complex tiles get up to `(max_split, max_split)`.
676pub fn adaptive_tile_partition(
677    complexities: &[TileComplexity],
678    threshold: f64,
679    max_split: u32,
680) -> Vec<(u32, u32)> {
681    let max_split = max_split.max(1).min(8);
682    complexities
683        .iter()
684        .map(|c| {
685            if c.score > threshold {
686                // Scale split factor by how far above threshold.
687                let factor = ((c.score - threshold) / (1.0 - threshold.min(0.999))
688                    * max_split as f64)
689                    .ceil() as u32;
690                let splits = factor.clamp(2, max_split);
691                (splits, splits)
692            } else {
693                (1, 1)
694            }
695        })
696        .collect()
697}
698
699// =============================================================================
700// Tile-Level Rate Control
701// =============================================================================
702
703/// Bit budget allocation for a single tile.
704#[derive(Clone, Debug, PartialEq)]
705pub struct TileBitBudget {
706    /// Tile column index.
707    pub col: u32,
708    /// Tile row index.
709    pub row: u32,
710    /// Allocated bits for this tile.
711    pub bits: u64,
712    /// Quality parameter (lower = higher quality, range depends on codec).
713    pub qp: f64,
714}
715
716/// Allocate a total bit budget across tiles based on complexity.
717///
718/// More complex tiles receive proportionally more bits.  The `total_bits`
719/// budget is distributed according to each tile's complexity score, with
720/// a minimum floor of `min_bits_per_tile`.
721pub fn allocate_tile_bits(
722    complexities: &[TileComplexity],
723    total_bits: u64,
724    min_bits_per_tile: u64,
725    base_qp: f64,
726) -> Vec<TileBitBudget> {
727    if complexities.is_empty() {
728        return Vec::new();
729    }
730
731    // Ensure minimum allocation is possible.
732    let min_total = min_bits_per_tile * complexities.len() as u64;
733    let distributable = total_bits.saturating_sub(min_total);
734
735    // Weight by complexity score (add small epsilon to avoid zero-weight).
736    let weights: Vec<f64> = complexities.iter().map(|c| c.score + 0.01).collect();
737    let total_weight: f64 = weights.iter().sum();
738
739    complexities
740        .iter()
741        .zip(weights.iter())
742        .map(|(c, &w)| {
743            let share = if total_weight > 0.0 {
744                (w / total_weight * distributable as f64) as u64
745            } else {
746                distributable / complexities.len() as u64
747            };
748            let bits = min_bits_per_tile + share;
749
750            // QP adjustment: lower complexity → higher QP (save bits).
751            // Higher complexity → lower QP (spend bits for quality).
752            let qp_delta = (1.0 - c.score) * 6.0 - 3.0; // range [-3, +3]
753            let qp = (base_qp + qp_delta).clamp(0.0, 51.0);
754
755            TileBitBudget {
756                col: c.col,
757                row: c.row,
758                bits,
759                qp,
760            }
761        })
762        .collect()
763}
764
765// =============================================================================
766// Tile Dependency Tracking
767// =============================================================================
768
769/// The type of dependency one tile has on another.
770#[derive(Clone, Debug, PartialEq, Eq)]
771pub enum TileDependencyKind {
772    /// Motion vector crosses into adjacent tile.
773    MotionVector,
774    /// In-loop filter requires border pixels from neighbour.
775    LoopFilter,
776    /// Entropy context is shared with the tile to the left.
777    EntropyContext,
778}
779
780/// A dependency edge from one tile to another.
781#[derive(Clone, Debug, PartialEq, Eq)]
782pub struct TileDependency {
783    /// Source tile (col, row).
784    pub from: (u32, u32),
785    /// Target tile (col, row) that `from` depends on.
786    pub to: (u32, u32),
787    /// Kind of dependency.
788    pub kind: TileDependencyKind,
789}
790
791/// Dependency graph for a tile layout.
792#[derive(Clone, Debug)]
793pub struct TileDependencyGraph {
794    /// All dependency edges.
795    pub edges: Vec<TileDependency>,
796    /// Number of tile columns.
797    pub cols: u32,
798    /// Number of tile rows.
799    pub rows: u32,
800}
801
802impl TileDependencyGraph {
803    /// Build a dependency graph for the given layout.
804    ///
805    /// By default, each tile depends on its left neighbour (entropy context)
806    /// and its top neighbour (loop filter boundary).  The caller can add
807    /// motion-vector dependencies afterwards.
808    pub fn build(layout: &TileLayout) -> Self {
809        let cols = layout.config.tile_cols;
810        let rows = layout.config.tile_rows;
811        let mut edges = Vec::new();
812
813        for row in 0..rows {
814            for col in 0..cols {
815                // Left neighbour: entropy context dependency.
816                if col > 0 {
817                    edges.push(TileDependency {
818                        from: (col, row),
819                        to: (col - 1, row),
820                        kind: TileDependencyKind::EntropyContext,
821                    });
822                }
823                // Top neighbour: loop-filter dependency.
824                if row > 0 {
825                    edges.push(TileDependency {
826                        from: (col, row),
827                        to: (col, row - 1),
828                        kind: TileDependencyKind::LoopFilter,
829                    });
830                }
831            }
832        }
833
834        Self { edges, cols, rows }
835    }
836
837    /// Add a motion-vector dependency between two tiles.
838    pub fn add_mv_dependency(&mut self, from: (u32, u32), to: (u32, u32)) {
839        if from.0 < self.cols && from.1 < self.rows && to.0 < self.cols && to.1 < self.rows {
840            self.edges.push(TileDependency {
841                from,
842                to,
843                kind: TileDependencyKind::MotionVector,
844            });
845        }
846    }
847
848    /// Return all tiles that `(col, row)` depends on.
849    pub fn dependencies_of(&self, col: u32, row: u32) -> Vec<&TileDependency> {
850        self.edges.iter().filter(|e| e.from == (col, row)).collect()
851    }
852
853    /// Return tiles that can be encoded independently (no incoming dependencies
854    /// from tiles that haven't been encoded yet).
855    ///
856    /// `encoded` is a set of already-encoded tile coordinates.
857    pub fn ready_tiles(&self, encoded: &[(u32, u32)]) -> Vec<(u32, u32)> {
858        let mut ready = Vec::new();
859        for row in 0..self.rows {
860            for col in 0..self.cols {
861                let pos = (col, row);
862                if encoded.contains(&pos) {
863                    continue;
864                }
865                let deps = self.dependencies_of(col, row);
866                let all_met = deps.iter().all(|d| encoded.contains(&d.to));
867                if all_met {
868                    ready.push(pos);
869                }
870            }
871        }
872        ready
873    }
874}
875
876// =============================================================================
877// Tile Work Queue (parallel encode with dependency awareness)
878// =============================================================================
879
880/// A work item for the tile encode queue.
881#[derive(Clone, Debug)]
882pub struct TileWorkItem {
883    /// Tile coordinate.
884    pub pos: (u32, u32),
885    /// Tile buffer to encode.
886    pub buffer: TileBuffer,
887    /// Bit budget (if rate control is active).
888    pub bit_budget: Option<TileBitBudget>,
889}
890
891/// Encodes tiles in dependency-aware waves using Rayon.
892///
893/// Tiles with satisfied dependencies are encoded in parallel waves.
894/// Returns encoded tile buffers in the order they were submitted.
895pub fn encode_tiles_wavefront<F>(
896    graph: &TileDependencyGraph,
897    mut work_items: Vec<TileWorkItem>,
898    encode_fn: F,
899) -> Result<Vec<TileBuffer>, String>
900where
901    F: Fn(TileWorkItem) -> Result<TileBuffer, String> + Send + Sync,
902{
903    let total = work_items.len();
904    let mut encoded_positions: Vec<(u32, u32)> = Vec::with_capacity(total);
905    let mut results: Vec<Option<TileBuffer>> = (0..total).map(|_| None).collect();
906
907    // Build a position → index map.
908    let pos_to_idx: std::collections::HashMap<(u32, u32), usize> = work_items
909        .iter()
910        .enumerate()
911        .map(|(i, w)| (w.pos, i))
912        .collect();
913
914    while encoded_positions.len() < total {
915        let ready = graph.ready_tiles(&encoded_positions);
916        if ready.is_empty() && encoded_positions.len() < total {
917            return Err("dependency deadlock: no tiles ready but not all encoded".to_string());
918        }
919
920        // Collect work items for this wave.
921        let wave_items: Vec<(usize, TileWorkItem)> = ready
922            .iter()
923            .filter_map(|pos| {
924                let idx = pos_to_idx.get(pos).copied();
925                idx.map(|i| {
926                    // Replace with a placeholder (empty buffer).
927                    let item = std::mem::replace(
928                        &mut work_items[i],
929                        TileWorkItem {
930                            pos: *pos,
931                            buffer: TileBuffer::new(TileRegion::new(0, 0, 0, 0, 0, 0), 1),
932                            bit_budget: None,
933                        },
934                    );
935                    (i, item)
936                })
937            })
938            .collect();
939
940        let wave_results: Vec<(usize, Result<TileBuffer, String>)> = wave_items
941            .into_par_iter()
942            .map(|(i, item)| (i, encode_fn(item)))
943            .collect();
944
945        for (i, result) in wave_results {
946            let buf = result?;
947            let pos = work_items[i].pos;
948            results[i] = Some(buf);
949            encoded_positions.push(pos);
950        }
951    }
952
953    // Collect results.
954    results
955        .into_iter()
956        .enumerate()
957        .map(|(i, r)| r.ok_or_else(|| format!("tile {} was not encoded", i)))
958        .collect()
959}
960
961// =============================================================================
962// Tile Quality Analysis
963// =============================================================================
964
965/// Per-tile quality metrics.
966#[derive(Clone, Debug, PartialEq)]
967pub struct TileQualityMetrics {
968    /// Tile column index.
969    pub col: u32,
970    /// Tile row index.
971    pub row: u32,
972    /// Estimated PSNR in dB (peak signal-to-noise ratio).
973    pub psnr_db: f64,
974    /// Estimated SSIM (structural similarity) in [0.0, 1.0].
975    pub ssim: f64,
976    /// Mean squared error.
977    pub mse: f64,
978}
979
980/// Compute quality metrics between original and reconstructed tile buffers.
981///
982/// Both buffers must have the same dimensions and channel count.
983pub fn compute_tile_quality(
984    original: &TileBuffer,
985    reconstructed: &TileBuffer,
986) -> Result<TileQualityMetrics, String> {
987    if original.data.len() != reconstructed.data.len() {
988        return Err("tile buffer sizes do not match".to_string());
989    }
990    if original.data.is_empty() {
991        return Ok(TileQualityMetrics {
992            col: original.region.col,
993            row: original.region.row,
994            psnr_db: f64::INFINITY,
995            ssim: 1.0,
996            mse: 0.0,
997        });
998    }
999
1000    let n = original.data.len() as f64;
1001
1002    // MSE
1003    let mse: f64 = original
1004        .data
1005        .iter()
1006        .zip(reconstructed.data.iter())
1007        .map(|(&a, &b)| {
1008            let diff = a as f64 - b as f64;
1009            diff * diff
1010        })
1011        .sum::<f64>()
1012        / n;
1013
1014    // PSNR
1015    let max_val = 255.0_f64;
1016    let psnr_db = if mse > 0.0 {
1017        10.0 * (max_val * max_val / mse).log10()
1018    } else {
1019        f64::INFINITY
1020    };
1021
1022    // Simplified SSIM (block-level approximation)
1023    let ssim = compute_ssim_approx(&original.data, &reconstructed.data);
1024
1025    Ok(TileQualityMetrics {
1026        col: original.region.col,
1027        row: original.region.row,
1028        psnr_db,
1029        ssim,
1030        mse,
1031    })
1032}
1033
1034/// Simplified SSIM computation between two equal-length byte slices.
1035///
1036/// Uses the standard SSIM formula with C1 = (0.01*255)^2, C2 = (0.03*255)^2.
1037fn compute_ssim_approx(a: &[u8], b: &[u8]) -> f64 {
1038    let n = a.len() as f64;
1039    if n < 1.0 {
1040        return 1.0;
1041    }
1042
1043    let c1: f64 = (0.01 * 255.0) * (0.01 * 255.0);
1044    let c2: f64 = (0.03 * 255.0) * (0.03 * 255.0);
1045
1046    let mut sum_a: f64 = 0.0;
1047    let mut sum_b: f64 = 0.0;
1048    let mut sum_a2: f64 = 0.0;
1049    let mut sum_b2: f64 = 0.0;
1050    let mut sum_ab: f64 = 0.0;
1051
1052    for (&va, &vb) in a.iter().zip(b.iter()) {
1053        let fa = va as f64;
1054        let fb = vb as f64;
1055        sum_a += fa;
1056        sum_b += fb;
1057        sum_a2 += fa * fa;
1058        sum_b2 += fb * fb;
1059        sum_ab += fa * fb;
1060    }
1061
1062    let mu_a = sum_a / n;
1063    let mu_b = sum_b / n;
1064    let sigma_a2 = (sum_a2 / n) - mu_a * mu_a;
1065    let sigma_b2 = (sum_b2 / n) - mu_b * mu_b;
1066    let sigma_ab = (sum_ab / n) - mu_a * mu_b;
1067
1068    let numerator = (2.0 * mu_a * mu_b + c1) * (2.0 * sigma_ab + c2);
1069    let denominator = (mu_a * mu_a + mu_b * mu_b + c1) * (sigma_a2 + sigma_b2 + c2);
1070
1071    if denominator > 0.0 {
1072        (numerator / denominator).clamp(-1.0, 1.0)
1073    } else {
1074        1.0
1075    }
1076}
1077
1078/// Compute quality metrics for all tiles by comparing original and reconstructed frames.
1079pub fn analyse_frame_quality(
1080    layout: &TileLayout,
1081    original_frame: &[u8],
1082    reconstructed_frame: &[u8],
1083    channels: u8,
1084) -> Result<Vec<TileQualityMetrics>, String> {
1085    let fw = layout.config.frame_width;
1086    let frame_stride = fw as usize * channels as usize;
1087
1088    layout
1089        .tiles
1090        .iter()
1091        .map(|region| {
1092            let mut orig_buf = TileBuffer::new(region.clone(), channels);
1093            let mut recon_buf = TileBuffer::new(region.clone(), channels);
1094            orig_buf.extract_from_frame(original_frame, frame_stride);
1095            recon_buf.extract_from_frame(reconstructed_frame, frame_stride);
1096            compute_tile_quality(&orig_buf, &recon_buf)
1097        })
1098        .collect()
1099}
1100
1101// =============================================================================
1102// Tests
1103// =============================================================================
1104
1105#[cfg(test)]
1106mod tests {
1107    use super::*;
1108
1109    // -----------------------------------------------------------------------
1110    // TileConfig
1111    // -----------------------------------------------------------------------
1112
1113    #[test]
1114    fn test_tile_config_default() {
1115        let cfg = TileConfig::default();
1116        assert_eq!(cfg.tile_cols, 1);
1117        assert_eq!(cfg.tile_rows, 1);
1118        assert_eq!(cfg.num_threads, 0);
1119        assert_eq!(cfg.frame_width, 0);
1120        assert_eq!(cfg.frame_height, 0);
1121    }
1122
1123    #[test]
1124    fn test_tile_config_builder() {
1125        let cfg = TileConfig::new()
1126            .tile_cols(4)
1127            .tile_rows(3)
1128            .num_threads(8)
1129            .frame_width(1920)
1130            .frame_height(1080);
1131
1132        assert_eq!(cfg.tile_cols, 4);
1133        assert_eq!(cfg.tile_rows, 3);
1134        assert_eq!(cfg.num_threads, 8);
1135        assert_eq!(cfg.frame_width, 1920);
1136        assert_eq!(cfg.frame_height, 1080);
1137    }
1138
1139    #[test]
1140    fn test_tile_config_clamp_cols() {
1141        // Values > 64 are clamped.
1142        let cfg = TileConfig::new().tile_cols(100);
1143        assert_eq!(cfg.tile_cols, 64);
1144    }
1145
1146    #[test]
1147    fn test_tile_config_thread_count_auto() {
1148        let cfg = TileConfig::new().num_threads(0);
1149        assert!(cfg.thread_count() >= 1);
1150    }
1151
1152    #[test]
1153    fn test_tile_config_thread_count_explicit() {
1154        let cfg = TileConfig::new().num_threads(4);
1155        assert_eq!(cfg.thread_count(), 4);
1156    }
1157
1158    // -----------------------------------------------------------------------
1159    // TileRegion
1160    // -----------------------------------------------------------------------
1161
1162    #[test]
1163    fn test_tile_region_area() {
1164        let r = TileRegion::new(0, 0, 0, 0, 100, 50);
1165        assert_eq!(r.area(), 5000);
1166    }
1167
1168    #[test]
1169    fn test_tile_region_contains() {
1170        let r = TileRegion::new(1, 0, 50, 0, 50, 50);
1171        assert!(r.contains(50, 0));
1172        assert!(r.contains(99, 49));
1173        assert!(!r.contains(49, 0)); // left of region
1174        assert!(!r.contains(100, 0)); // right boundary (exclusive)
1175        assert!(!r.contains(50, 50)); // bottom boundary (exclusive)
1176    }
1177
1178    #[test]
1179    fn test_tile_region_pixel_ranges() {
1180        let r = TileRegion::new(0, 1, 0, 100, 200, 80);
1181        assert_eq!(r.pixel_range_x(), 0..200);
1182        assert_eq!(r.pixel_range_y(), 100..180);
1183    }
1184
1185    // -----------------------------------------------------------------------
1186    // TileLayout – divisible dimensions
1187    // -----------------------------------------------------------------------
1188
1189    #[test]
1190    fn test_tile_layout_2x2_divisible() {
1191        let cfg = TileConfig::new()
1192            .tile_cols(2)
1193            .tile_rows(2)
1194            .frame_width(100)
1195            .frame_height(100);
1196
1197        let layout = TileLayout::new(cfg);
1198        assert_eq!(layout.tile_count(), 4);
1199
1200        // All tiles should be 50×50 for an evenly-divisible frame.
1201        for tile in layout.tiles() {
1202            assert_eq!(tile.width, 50);
1203            assert_eq!(tile.height, 50);
1204        }
1205
1206        // Total area must equal frame area.
1207        let total: u64 = layout.tiles().iter().map(|t| t.area()).sum();
1208        assert_eq!(total, 100 * 100);
1209    }
1210
1211    #[test]
1212    fn test_tile_layout_get_tile() {
1213        let cfg = TileConfig::new()
1214            .tile_cols(2)
1215            .tile_rows(2)
1216            .frame_width(100)
1217            .frame_height(100);
1218
1219        let layout = TileLayout::new(cfg);
1220
1221        let tl = layout.get_tile(0, 0).expect("should succeed");
1222        assert_eq!((tl.x, tl.y), (0, 0));
1223
1224        let tr = layout.get_tile(1, 0).expect("should succeed");
1225        assert_eq!(tr.x, 50);
1226
1227        let bl = layout.get_tile(0, 1).expect("should succeed");
1228        assert_eq!(bl.y, 50);
1229
1230        assert!(layout.get_tile(2, 0).is_none());
1231    }
1232
1233    // -----------------------------------------------------------------------
1234    // TileLayout – non-divisible dimensions
1235    // -----------------------------------------------------------------------
1236
1237    #[test]
1238    fn test_tile_layout_2x2_non_divisible() {
1239        // 101×101 with 2×2 tiles: nominal 50×50, last col/row gets remainder.
1240        let cfg = TileConfig::new()
1241            .tile_cols(2)
1242            .tile_rows(2)
1243            .frame_width(101)
1244            .frame_height(101);
1245
1246        let layout = TileLayout::new(cfg);
1247        assert_eq!(layout.tile_count(), 4);
1248
1249        // Top-left: 50×50
1250        let tl = layout.get_tile(0, 0).expect("should succeed");
1251        assert_eq!(tl.width, 50);
1252        assert_eq!(tl.height, 50);
1253
1254        // Top-right: 51×50 (gets the 1-pixel remainder in x)
1255        let tr = layout.get_tile(1, 0).expect("should succeed");
1256        assert_eq!(tr.width, 51);
1257        assert_eq!(tr.height, 50);
1258
1259        // Bottom-left: 50×51
1260        let bl = layout.get_tile(0, 1).expect("should succeed");
1261        assert_eq!(bl.width, 50);
1262        assert_eq!(bl.height, 51);
1263
1264        // Bottom-right: 51×51
1265        let br = layout.get_tile(1, 1).expect("should succeed");
1266        assert_eq!(br.width, 51);
1267        assert_eq!(br.height, 51);
1268
1269        // Total area == 101×101
1270        let total: u64 = layout.tiles().iter().map(|t| t.area()).sum();
1271        assert_eq!(total, 101 * 101);
1272    }
1273
1274    #[test]
1275    fn test_tile_layout_non_divisible_coverage() {
1276        // Verify every pixel is covered exactly once.
1277        let fw = 97u32;
1278        let fh = 83u32;
1279        let cfg = TileConfig::new()
1280            .tile_cols(3)
1281            .tile_rows(3)
1282            .frame_width(fw)
1283            .frame_height(fh);
1284
1285        let layout = TileLayout::new(cfg);
1286        let mut counts = vec![0u32; (fw * fh) as usize];
1287
1288        for tile in layout.tiles() {
1289            for py in tile.pixel_range_y() {
1290                for px in tile.pixel_range_x() {
1291                    counts[(py * fw + px) as usize] += 1;
1292                }
1293            }
1294        }
1295
1296        assert!(
1297            counts.iter().all(|&c| c == 1),
1298            "some pixels are covered 0 or 2+ times"
1299        );
1300    }
1301
1302    // -----------------------------------------------------------------------
1303    // TileLayout – tile_for_pixel
1304    // -----------------------------------------------------------------------
1305
1306    #[test]
1307    fn test_tile_for_pixel() {
1308        let cfg = TileConfig::new()
1309            .tile_cols(2)
1310            .tile_rows(2)
1311            .frame_width(100)
1312            .frame_height(100);
1313
1314        let layout = TileLayout::new(cfg);
1315
1316        let t = layout.tile_for_pixel(25, 25).expect("should succeed");
1317        assert_eq!((t.col, t.row), (0, 0));
1318
1319        let t = layout.tile_for_pixel(75, 25).expect("should succeed");
1320        assert_eq!((t.col, t.row), (1, 0));
1321
1322        let t = layout.tile_for_pixel(25, 75).expect("should succeed");
1323        assert_eq!((t.col, t.row), (0, 1));
1324
1325        let t = layout.tile_for_pixel(75, 75).expect("should succeed");
1326        assert_eq!((t.col, t.row), (1, 1));
1327
1328        // Out-of-frame pixel.
1329        assert!(layout.tile_for_pixel(200, 200).is_none());
1330    }
1331
1332    // -----------------------------------------------------------------------
1333    // TileBuffer
1334    // -----------------------------------------------------------------------
1335
1336    #[test]
1337    fn test_tile_buffer_new() {
1338        let region = TileRegion::new(0, 0, 0, 0, 8, 6);
1339        let buf = TileBuffer::new(region, 3);
1340        assert_eq!(buf.stride, 8 * 3);
1341        assert_eq!(buf.data.len(), 8 * 6 * 3);
1342        assert!(buf.data.iter().all(|&b| b == 0));
1343    }
1344
1345    #[test]
1346    fn test_tile_buffer_extract() {
1347        // 4×4 single-channel frame: pixels 0..16
1348        let frame: Vec<u8> = (0u8..16).collect();
1349        let region = TileRegion::new(0, 0, 1, 1, 2, 2); // 2×2 tile at offset (1,1)
1350        let mut buf = TileBuffer::new(region, 1);
1351        buf.extract_from_frame(&frame, 4); // frame stride = 4
1352
1353        // Row 1, col 1 → index 5; row 1, col 2 → index 6
1354        // Row 2, col 1 → index 9; row 2, col 2 → index 10
1355        assert_eq!(buf.data, vec![5, 6, 9, 10]);
1356    }
1357
1358    #[test]
1359    fn test_tile_buffer_write_back() {
1360        let region = TileRegion::new(0, 0, 1, 1, 2, 2);
1361        let mut buf = TileBuffer::new(region, 1);
1362        buf.data = vec![5, 6, 9, 10];
1363
1364        let mut frame = vec![0u8; 16];
1365        buf.write_to_frame(&mut frame, 4);
1366
1367        assert_eq!(frame[5], 5);
1368        assert_eq!(frame[6], 6);
1369        assert_eq!(frame[9], 9);
1370        assert_eq!(frame[10], 10);
1371        // Other pixels untouched.
1372        assert_eq!(frame[0], 0);
1373    }
1374
1375    // -----------------------------------------------------------------------
1376    // ParallelTileEncoder – split and merge roundtrip
1377    // -----------------------------------------------------------------------
1378
1379    #[test]
1380    fn test_split_merge_roundtrip_divisible() {
1381        let fw = 64u32;
1382        let fh = 64u32;
1383        let channels = 3u8;
1384
1385        let config = TileConfig::new()
1386            .tile_cols(4)
1387            .tile_rows(4)
1388            .frame_width(fw)
1389            .frame_height(fh);
1390
1391        let encoder = ParallelTileEncoder::new(config);
1392
1393        // Create a unique frame.
1394        let frame: Vec<u8> = (0u8..=255)
1395            .cycle()
1396            .take((fw * fh * channels as u32) as usize)
1397            .collect();
1398        let tiles = encoder.split_frame(&frame, channels);
1399        assert_eq!(tiles.len(), 16);
1400
1401        let merged = ParallelTileEncoder::merge_tiles(&tiles, fw, fh, channels);
1402        assert_eq!(merged, frame, "roundtrip failed for divisible dimensions");
1403    }
1404
1405    #[test]
1406    fn test_split_merge_roundtrip_non_divisible() {
1407        let fw = 101u32;
1408        let fh = 99u32;
1409        let channels = 1u8;
1410
1411        let config = TileConfig::new()
1412            .tile_cols(3)
1413            .tile_rows(3)
1414            .frame_width(fw)
1415            .frame_height(fh);
1416
1417        let encoder = ParallelTileEncoder::new(config);
1418
1419        let frame: Vec<u8> = (0u8..=255).cycle().take((fw * fh) as usize).collect();
1420        let tiles = encoder.split_frame(&frame, channels);
1421
1422        let merged = ParallelTileEncoder::merge_tiles(&tiles, fw, fh, channels);
1423        assert_eq!(
1424            merged, frame,
1425            "roundtrip failed for non-divisible dimensions"
1426        );
1427    }
1428
1429    // -----------------------------------------------------------------------
1430    // ParallelTileEncoder – encode_tiles_parallel
1431    // -----------------------------------------------------------------------
1432
1433    #[test]
1434    fn test_encode_tiles_parallel_identity() {
1435        let fw = 64u32;
1436        let fh = 64u32;
1437        let channels = 3u8;
1438
1439        let config = TileConfig::new()
1440            .tile_cols(2)
1441            .tile_rows(2)
1442            .frame_width(fw)
1443            .frame_height(fh);
1444
1445        let encoder = ParallelTileEncoder::new(config);
1446
1447        let frame: Vec<u8> = (0u8..=255)
1448            .cycle()
1449            .take((fw * fh * channels as u32) as usize)
1450            .collect();
1451        let tiles = encoder.split_frame(&frame, channels);
1452
1453        // Identity encode: return each tile unchanged.
1454        let processed = encoder
1455            .encode_tiles_parallel(tiles, |tile| Ok(tile))
1456            .expect("should succeed");
1457
1458        let merged = ParallelTileEncoder::merge_tiles(&processed, fw, fh, channels);
1459        assert_eq!(merged, frame, "parallel identity encode broke the frame");
1460    }
1461
1462    #[test]
1463    fn test_encode_tiles_parallel_error_propagates() {
1464        let config = TileConfig::new()
1465            .tile_cols(2)
1466            .tile_rows(2)
1467            .frame_width(64)
1468            .frame_height(64);
1469
1470        let encoder = ParallelTileEncoder::new(config);
1471        let frame = vec![0u8; 64 * 64 * 3];
1472        let tiles = encoder.split_frame(&frame, 3);
1473
1474        let result = encoder.encode_tiles_parallel(tiles, |_| Err("deliberate error".to_string()));
1475        assert!(result.is_err());
1476    }
1477
1478    #[test]
1479    fn test_encode_tiles_parallel_transform() {
1480        // Invert all pixel values and check the result.
1481        let fw = 32u32;
1482        let fh = 32u32;
1483        let channels = 1u8;
1484
1485        let config = TileConfig::new()
1486            .tile_cols(2)
1487            .tile_rows(2)
1488            .frame_width(fw)
1489            .frame_height(fh);
1490
1491        let encoder = ParallelTileEncoder::new(config);
1492
1493        let frame: Vec<u8> = (0u8..=255).cycle().take((fw * fh) as usize).collect();
1494        let tiles = encoder.split_frame(&frame, channels);
1495
1496        let inverted = encoder
1497            .encode_tiles_parallel(tiles, |mut tile| {
1498                for b in &mut tile.data {
1499                    *b = 255 - *b;
1500                }
1501                Ok(tile)
1502            })
1503            .expect("should succeed");
1504
1505        let merged = ParallelTileEncoder::merge_tiles(&inverted, fw, fh, channels);
1506        let expected: Vec<u8> = frame.iter().map(|&b| 255 - b).collect();
1507        assert_eq!(merged, expected, "inversion result mismatch");
1508    }
1509
1510    // -----------------------------------------------------------------------
1511    // Tile Complexity Analysis
1512    // -----------------------------------------------------------------------
1513
1514    #[test]
1515    fn test_analyse_tile_complexity_uniform() {
1516        let cfg = TileConfig::new()
1517            .tile_cols(2)
1518            .tile_rows(2)
1519            .frame_width(8)
1520            .frame_height(8);
1521        let layout = TileLayout::new(cfg);
1522        // Uniform frame: all pixels = 128.
1523        let frame = vec![128u8; 8 * 8];
1524        let complexities = analyse_tile_complexity(&layout, &frame, 1);
1525        assert_eq!(complexities.len(), 4);
1526        for c in &complexities {
1527            assert!(
1528                c.variance < 1.0,
1529                "uniform frame should have near-zero variance"
1530            );
1531            assert!(
1532                c.edge_density < 1.0,
1533                "uniform frame should have near-zero edge density"
1534            );
1535        }
1536    }
1537
1538    #[test]
1539    fn test_analyse_tile_complexity_gradient() {
1540        let cfg = TileConfig::new()
1541            .tile_cols(1)
1542            .tile_rows(1)
1543            .frame_width(16)
1544            .frame_height(16);
1545        let layout = TileLayout::new(cfg);
1546        // Gradient frame: increasing pixel values.
1547        let frame: Vec<u8> = (0..16 * 16).map(|i| (i % 256) as u8).collect();
1548        let complexities = analyse_tile_complexity(&layout, &frame, 1);
1549        assert_eq!(complexities.len(), 1);
1550        assert!(
1551            complexities[0].variance > 0.0,
1552            "gradient should have non-zero variance"
1553        );
1554        assert!(
1555            complexities[0].edge_density > 0.0,
1556            "gradient should have non-zero edge density"
1557        );
1558    }
1559
1560    #[test]
1561    fn test_analyse_complexity_score_normalised() {
1562        let cfg = TileConfig::new()
1563            .tile_cols(2)
1564            .tile_rows(2)
1565            .frame_width(16)
1566            .frame_height(16);
1567        let layout = TileLayout::new(cfg);
1568        // Mixed frame: top-left noisy, rest uniform.
1569        let mut frame = vec![128u8; 16 * 16];
1570        for y in 0..8 {
1571            for x in 0..8 {
1572                frame[y * 16 + x] = ((x * 31 + y * 17) % 256) as u8;
1573            }
1574        }
1575        let complexities = analyse_tile_complexity(&layout, &frame, 1);
1576        for c in &complexities {
1577            assert!(
1578                c.score >= 0.0 && c.score <= 1.0,
1579                "score out of range: {}",
1580                c.score
1581            );
1582        }
1583    }
1584
1585    // -----------------------------------------------------------------------
1586    // Adaptive Tile Partitioning
1587    // -----------------------------------------------------------------------
1588
1589    #[test]
1590    fn test_adaptive_partition_below_threshold() {
1591        let complexities = vec![
1592            TileComplexity {
1593                col: 0,
1594                row: 0,
1595                variance: 10.0,
1596                edge_density: 5.0,
1597                score: 0.2,
1598            },
1599            TileComplexity {
1600                col: 1,
1601                row: 0,
1602                variance: 15.0,
1603                edge_density: 7.0,
1604                score: 0.3,
1605            },
1606        ];
1607        let partitions = adaptive_tile_partition(&complexities, 0.5, 4);
1608        assert_eq!(partitions, vec![(1, 1), (1, 1)]);
1609    }
1610
1611    #[test]
1612    fn test_adaptive_partition_above_threshold() {
1613        let complexities = vec![TileComplexity {
1614            col: 0,
1615            row: 0,
1616            variance: 100.0,
1617            edge_density: 50.0,
1618            score: 0.9,
1619        }];
1620        let partitions = adaptive_tile_partition(&complexities, 0.5, 4);
1621        assert!(partitions[0].0 >= 2, "high-complexity tile should be split");
1622        assert!(partitions[0].1 >= 2, "high-complexity tile should be split");
1623    }
1624
1625    #[test]
1626    fn test_adaptive_partition_max_split_clamped() {
1627        let complexities = vec![TileComplexity {
1628            col: 0,
1629            row: 0,
1630            variance: 1000.0,
1631            edge_density: 500.0,
1632            score: 1.0,
1633        }];
1634        let partitions = adaptive_tile_partition(&complexities, 0.0, 4);
1635        assert!(partitions[0].0 <= 4);
1636        assert!(partitions[0].1 <= 4);
1637    }
1638
1639    // -----------------------------------------------------------------------
1640    // Tile-Level Rate Control
1641    // -----------------------------------------------------------------------
1642
1643    #[test]
1644    fn test_allocate_tile_bits_proportional() {
1645        let complexities = vec![
1646            TileComplexity {
1647                col: 0,
1648                row: 0,
1649                variance: 100.0,
1650                edge_density: 50.0,
1651                score: 0.8,
1652            },
1653            TileComplexity {
1654                col: 1,
1655                row: 0,
1656                variance: 10.0,
1657                edge_density: 5.0,
1658                score: 0.2,
1659            },
1660        ];
1661        let budgets = allocate_tile_bits(&complexities, 10000, 100, 28.0);
1662        assert_eq!(budgets.len(), 2);
1663        // Higher complexity tile should get more bits.
1664        assert!(budgets[0].bits > budgets[1].bits);
1665        // Total should be close to the budget (rounding may cause +-1 per tile).
1666        let total: u64 = budgets.iter().map(|b| b.bits).sum();
1667        let diff = (total as i64 - 10000_i64).unsigned_abs();
1668        assert!(
1669            diff <= budgets.len() as u64,
1670            "total {} too far from 10000",
1671            total
1672        );
1673    }
1674
1675    #[test]
1676    fn test_allocate_tile_bits_minimum_floor() {
1677        let complexities = vec![TileComplexity {
1678            col: 0,
1679            row: 0,
1680            variance: 0.0,
1681            edge_density: 0.0,
1682            score: 0.0,
1683        }];
1684        let budgets = allocate_tile_bits(&complexities, 1000, 500, 28.0);
1685        assert!(budgets[0].bits >= 500);
1686    }
1687
1688    #[test]
1689    fn test_allocate_tile_bits_qp_range() {
1690        let complexities = vec![
1691            TileComplexity {
1692                col: 0,
1693                row: 0,
1694                variance: 0.0,
1695                edge_density: 0.0,
1696                score: 0.0,
1697            },
1698            TileComplexity {
1699                col: 1,
1700                row: 0,
1701                variance: 100.0,
1702                edge_density: 50.0,
1703                score: 1.0,
1704            },
1705        ];
1706        let budgets = allocate_tile_bits(&complexities, 10000, 100, 28.0);
1707        for b in &budgets {
1708            assert!(b.qp >= 0.0 && b.qp <= 51.0, "QP out of range: {}", b.qp);
1709        }
1710    }
1711
1712    #[test]
1713    fn test_allocate_tile_bits_empty() {
1714        let budgets = allocate_tile_bits(&[], 10000, 100, 28.0);
1715        assert!(budgets.is_empty());
1716    }
1717
1718    // -----------------------------------------------------------------------
1719    // Tile Dependency Tracking
1720    // -----------------------------------------------------------------------
1721
1722    #[test]
1723    fn test_dependency_graph_build_2x2() {
1724        let cfg = TileConfig::new()
1725            .tile_cols(2)
1726            .tile_rows(2)
1727            .frame_width(64)
1728            .frame_height(64);
1729        let layout = TileLayout::new(cfg);
1730        let graph = TileDependencyGraph::build(&layout);
1731
1732        // (0,0): no deps; (1,0): left; (0,1): top; (1,1): left + top
1733        assert_eq!(graph.dependencies_of(0, 0).len(), 0);
1734        assert_eq!(graph.dependencies_of(1, 0).len(), 1);
1735        assert_eq!(graph.dependencies_of(0, 1).len(), 1);
1736        assert_eq!(graph.dependencies_of(1, 1).len(), 2);
1737    }
1738
1739    #[test]
1740    fn test_dependency_graph_ready_tiles() {
1741        let cfg = TileConfig::new()
1742            .tile_cols(2)
1743            .tile_rows(2)
1744            .frame_width(64)
1745            .frame_height(64);
1746        let layout = TileLayout::new(cfg);
1747        let graph = TileDependencyGraph::build(&layout);
1748
1749        // Initially only (0,0) is ready (no dependencies).
1750        let ready = graph.ready_tiles(&[]);
1751        assert!(ready.contains(&(0, 0)));
1752        assert!(!ready.contains(&(1, 1)));
1753
1754        // After encoding (0,0), (1,0) and (0,1) become ready.
1755        let ready = graph.ready_tiles(&[(0, 0)]);
1756        assert!(ready.contains(&(1, 0)));
1757        assert!(ready.contains(&(0, 1)));
1758    }
1759
1760    #[test]
1761    fn test_dependency_graph_add_mv() {
1762        let cfg = TileConfig::new()
1763            .tile_cols(2)
1764            .tile_rows(2)
1765            .frame_width(64)
1766            .frame_height(64);
1767        let layout = TileLayout::new(cfg);
1768        let mut graph = TileDependencyGraph::build(&layout);
1769
1770        let before = graph.dependencies_of(0, 0).len();
1771        graph.add_mv_dependency((0, 0), (1, 1));
1772        let after = graph.dependencies_of(0, 0).len();
1773        assert_eq!(after, before + 1);
1774    }
1775
1776    // -----------------------------------------------------------------------
1777    // Tile Quality Analysis
1778    // -----------------------------------------------------------------------
1779
1780    #[test]
1781    fn test_tile_quality_identical() {
1782        let region = TileRegion::new(0, 0, 0, 0, 4, 4);
1783        let mut buf = TileBuffer::new(region, 1);
1784        buf.data = vec![100; 16];
1785        let metrics = compute_tile_quality(&buf, &buf).expect("should succeed");
1786        assert_eq!(metrics.mse, 0.0);
1787        assert!(metrics.psnr_db.is_infinite());
1788        assert!((metrics.ssim - 1.0).abs() < 1e-6);
1789    }
1790
1791    #[test]
1792    fn test_tile_quality_different() {
1793        let region = TileRegion::new(0, 0, 0, 0, 4, 4);
1794        let mut orig = TileBuffer::new(region.clone(), 1);
1795        let mut recon = TileBuffer::new(region, 1);
1796        orig.data = vec![100; 16];
1797        recon.data = vec![110; 16];
1798
1799        let metrics = compute_tile_quality(&orig, &recon).expect("should succeed");
1800        assert!(metrics.mse > 0.0);
1801        assert!(metrics.psnr_db > 0.0 && metrics.psnr_db < 100.0);
1802        assert!(metrics.ssim < 1.0);
1803    }
1804
1805    #[test]
1806    fn test_tile_quality_size_mismatch() {
1807        let r1 = TileRegion::new(0, 0, 0, 0, 4, 4);
1808        let r2 = TileRegion::new(0, 0, 0, 0, 4, 2);
1809        let buf1 = TileBuffer::new(r1, 1);
1810        let buf2 = TileBuffer::new(r2, 1);
1811        let result = compute_tile_quality(&buf1, &buf2);
1812        assert!(result.is_err());
1813    }
1814
1815    #[test]
1816    fn test_analyse_frame_quality_full() {
1817        let cfg = TileConfig::new()
1818            .tile_cols(2)
1819            .tile_rows(2)
1820            .frame_width(8)
1821            .frame_height(8);
1822        let layout = TileLayout::new(cfg);
1823        let original: Vec<u8> = (0..64).collect();
1824        let reconstructed = original.clone();
1825        let metrics =
1826            analyse_frame_quality(&layout, &original, &reconstructed, 1).expect("should succeed");
1827        assert_eq!(metrics.len(), 4);
1828        for m in &metrics {
1829            assert_eq!(m.mse, 0.0);
1830        }
1831    }
1832
1833    // -----------------------------------------------------------------------
1834    // Wavefront Encoding
1835    // -----------------------------------------------------------------------
1836
1837    #[test]
1838    fn test_wavefront_encode_2x2() {
1839        let cfg = TileConfig::new()
1840            .tile_cols(2)
1841            .tile_rows(2)
1842            .frame_width(8)
1843            .frame_height(8);
1844        let layout = TileLayout::new(cfg);
1845        let graph = TileDependencyGraph::build(&layout);
1846
1847        let encoder = ParallelTileEncoder::new(
1848            TileConfig::new()
1849                .tile_cols(2)
1850                .tile_rows(2)
1851                .frame_width(8)
1852                .frame_height(8),
1853        );
1854        let frame: Vec<u8> = (0..64).collect();
1855        let tiles = encoder.split_frame(&frame, 1);
1856
1857        let work_items: Vec<TileWorkItem> = tiles
1858            .into_iter()
1859            .map(|buf| TileWorkItem {
1860                pos: (buf.region.col, buf.region.row),
1861                buffer: buf,
1862                bit_budget: None,
1863            })
1864            .collect();
1865
1866        let results = encode_tiles_wavefront(&graph, work_items, |item| Ok(item.buffer))
1867            .expect("should succeed");
1868        assert_eq!(results.len(), 4);
1869    }
1870}