Skip to main content

mozjpeg_rs/
encode.rs

1//! JPEG encoder pipeline.
2//!
3//! This module provides two encoder types:
4//!
5//! - [`Encoder`]: Full-featured encoder with trellis quantization, progressive mode,
6//!   and Huffman optimization. Batch encoding only.
7//! - [`StreamingEncoder`]: Streaming-capable encoder without optimizations.
8//!   Supports both batch and scanline-by-scanline encoding.
9//!
10//! Both implement the [`Encode`] trait for batch encoding.
11//!
12//! # Examples
13//!
14//! ```ignore
15//! use mozjpeg_rs::{Encoder, Preset};
16//!
17//! // Full-featured batch encoding
18//! let jpeg = Encoder::new(Preset::default())
19//!     .quality(85)
20//!     .encode_rgb(&pixels, width, height)?;
21//!
22//! // Streaming encoding (memory-efficient for large images)
23//! let mut stream = Encoder::streaming()
24//!     .quality(85)
25//!     .start(width, height, file)?;
26//! for row in scanlines.chunks(16) {
27//!     stream.write_scanlines(row)?;
28//! }
29//! stream.finish()?;
30//! ```
31
32use std::io::Write;
33use std::sync::atomic::{AtomicBool, Ordering};
34use std::time::{Duration, Instant};
35
36use crate::bitstream::BitWriter;
37use crate::color::convert_rgb_to_ycbcr_c_compat;
38use crate::consts::{DCTSIZE, DCTSIZE2, QuantTableIdx};
39use crate::deringing::preprocess_deringing;
40use crate::entropy::{EntropyEncoder, ProgressiveEncoder, ProgressiveSymbolCounter, SymbolCounter};
41use crate::error::{Error, Result};
42use crate::huffman::DerivedTable;
43use crate::huffman::FrequencyCounter;
44use crate::marker::MarkerWriter;
45use crate::progressive::{generate_baseline_scan, generate_mozjpeg_max_compression_scans};
46use crate::quant::{create_quant_tables, quantize_block_raw};
47use crate::sample;
48use crate::scan_optimize::{ScanSearchConfig, ScanSelector, generate_search_scans};
49use crate::scan_trial::ScanTrialEncoder;
50use crate::simd::SimdOps;
51#[cfg(target_arch = "x86_64")]
52use crate::simd::x86_64::entropy::SimdEntropyEncoder;
53use crate::trellis::trellis_quantize_block;
54use crate::types::{Limits, PixelDensity, Preset, Subsampling, TrellisConfig};
55
56mod helpers;
57mod streaming;
58
59pub(crate) use helpers::{
60    create_components, create_std_ac_chroma_table, create_std_ac_luma_table,
61    create_std_dc_chroma_table, create_std_dc_luma_table, create_ycbcr_components,
62    natural_to_zigzag, run_dc_trellis_by_row, try_alloc_vec, try_alloc_vec_array, write_dht_marker,
63    write_sos_marker,
64};
65pub use streaming::{EncodingStream, StreamingEncoder};
66
67// ============================================================================
68// Cancellation Support
69// ============================================================================
70
71/// Internal context for cancellation checking during encoding.
72///
73/// This is passed through the encoding pipeline to allow periodic
74/// cancellation checks without function signature changes everywhere.
75///
76/// Implements [`enough::Stop`] so it can be used interchangeably with
77/// any cooperative cancellation source.
78#[derive(Clone, Copy)]
79pub(crate) struct CancellationContext<'a> {
80    /// Optional cancellation flag - if set to true, encoding should abort.
81    pub cancel: Option<&'a AtomicBool>,
82    /// Optional deadline - if current time exceeds this, encoding should abort.
83    pub deadline: Option<Instant>,
84}
85
86impl<'a> CancellationContext<'a> {
87    /// Create a context with no cancellation (always succeeds).
88    #[allow(dead_code)]
89    pub const fn none() -> Self {
90        Self {
91            cancel: None,
92            deadline: None,
93        }
94    }
95
96    /// Create a context from optional cancel flag and timeout.
97    #[allow(dead_code)]
98    pub fn new(cancel: Option<&'a AtomicBool>, timeout: Option<Duration>) -> Self {
99        Self {
100            cancel,
101            deadline: timeout.map(|d| Instant::now() + d),
102        }
103    }
104
105    /// Check if cancellation has been requested.
106    ///
107    /// Returns `Ok(())` if encoding should continue, or `Err` if cancelled/timed out.
108    #[inline]
109    pub fn check(&self) -> Result<()> {
110        if let Some(c) = self.cancel
111            && c.load(Ordering::Relaxed)
112        {
113            return Err(Error::Cancelled);
114        }
115        if let Some(d) = self.deadline
116            && Instant::now() > d
117        {
118            return Err(Error::TimedOut);
119        }
120        Ok(())
121    }
122
123    /// Check cancellation every N iterations (to reduce overhead).
124    ///
125    /// Only performs the check when `iteration % interval == 0`.
126    #[inline]
127    #[allow(dead_code)]
128    pub fn check_periodic(&self, iteration: usize, interval: usize) -> Result<()> {
129        if iteration.is_multiple_of(interval) {
130            self.check()
131        } else {
132            Ok(())
133        }
134    }
135}
136
137impl enough::Stop for CancellationContext<'_> {
138    fn check(&self) -> std::result::Result<(), enough::StopReason> {
139        if let Some(c) = self.cancel
140            && c.load(Ordering::Relaxed)
141        {
142            return Err(enough::StopReason::Cancelled);
143        }
144        if let Some(d) = self.deadline
145            && Instant::now() > d
146        {
147            return Err(enough::StopReason::TimedOut);
148        }
149        Ok(())
150    }
151
152    fn may_stop(&self) -> bool {
153        self.cancel.is_some() || self.deadline.is_some()
154    }
155}
156
157// ============================================================================
158// Encode Trait (internal, for potential future streaming API)
159// ============================================================================
160
161/// Trait for JPEG encoding (batch mode).
162///
163/// Implemented by both [`Encoder`] and [`StreamingEncoder`].
164#[allow(dead_code)]
165pub trait Encode {
166    /// Encode RGB image data to JPEG.
167    ///
168    /// # Arguments
169    /// * `rgb_data` - RGB pixel data (3 bytes per pixel, row-major order)
170    /// * `width` - Image width in pixels
171    /// * `height` - Image height in pixels
172    fn encode_rgb(&self, rgb_data: &[u8], width: u32, height: u32) -> Result<Vec<u8>>;
173
174    /// Encode grayscale image data to JPEG.
175    ///
176    /// # Arguments
177    /// * `gray_data` - Grayscale pixel data (1 byte per pixel, row-major order)
178    /// * `width` - Image width in pixels
179    /// * `height` - Image height in pixels
180    fn encode_gray(&self, gray_data: &[u8], width: u32, height: u32) -> Result<Vec<u8>>;
181}
182
183/// JPEG encoder with configurable quality and features.
184#[derive(Debug, Clone)]
185pub struct Encoder {
186    /// Quality level (1-100)
187    quality: u8,
188    /// Enable progressive mode
189    progressive: bool,
190    /// Chroma subsampling mode
191    subsampling: Subsampling,
192    /// Quantization table variant
193    quant_table_idx: QuantTableIdx,
194    /// Custom luminance quantization table (overrides quant_table_idx if set)
195    custom_luma_qtable: Option<[u16; DCTSIZE2]>,
196    /// Custom chrominance quantization table (overrides quant_table_idx if set)
197    custom_chroma_qtable: Option<[u16; DCTSIZE2]>,
198    /// Trellis quantization configuration
199    trellis: TrellisConfig,
200    /// Force baseline-compatible output
201    force_baseline: bool,
202    /// Optimize Huffman tables (requires 2-pass)
203    optimize_huffman: bool,
204    /// Enable overshoot deringing (reduces ringing on white backgrounds)
205    overshoot_deringing: bool,
206    /// Use C mozjpeg-compatible color conversion for exact parity.
207    /// Produces bytewise-identical YCbCr values to C mozjpeg.
208    /// Enabled by default; use `.fast_color()` for faster yuv crate (~40% faster, ±1 rounding).
209    c_compat_color: bool,
210    /// Optimize progressive scan configuration (tries multiple configs, picks smallest)
211    optimize_scans: bool,
212    /// Restart interval in MCUs (0 = disabled)
213    restart_interval: u16,
214    /// Pixel density for JFIF APP0 marker
215    pixel_density: PixelDensity,
216    /// EXIF data to embed (raw TIFF structure, without "Exif\0\0" header)
217    exif_data: Option<Vec<u8>>,
218    /// ICC color profile to embed (will be chunked into APP2 markers)
219    icc_profile: Option<Vec<u8>>,
220    /// Custom APP markers to embed (marker number 0-15, data)
221    custom_markers: Vec<(u8, Vec<u8>)>,
222    /// SIMD operations dispatch (detected once at construction)
223    simd: SimdOps,
224    /// Smoothing factor (0-100, 0 = disabled)
225    /// Applies a weighted average filter to reduce fine-scale noise.
226    /// Useful for converting dithered images (like GIFs) to JPEG.
227    smoothing: u8,
228    /// Resource limits (dimensions, memory, ICC size)
229    limits: Limits,
230}
231
232impl Default for Encoder {
233    fn default() -> Self {
234        Self::new(Preset::default())
235    }
236}
237
238impl Encoder {
239    /// Create an encoder with the specified preset.
240    ///
241    /// # Arguments
242    ///
243    /// * `preset` - Encoding preset (see [`Preset`] for details):
244    ///   - [`BaselineFastest`](Preset::BaselineFastest): No optimizations, fastest encoding
245    ///   - [`BaselineBalanced`](Preset::BaselineBalanced): Baseline with all optimizations
246    ///   - [`ProgressiveBalanced`](Preset::ProgressiveBalanced): Progressive with optimizations (default)
247    ///   - [`ProgressiveSmallest`](Preset::ProgressiveSmallest): Maximum compression
248    ///
249    /// # Preset Comparison
250    ///
251    /// | Preset | Time | Size | Best For |
252    /// |--------|------|------|----------|
253    /// | `BaselineFastest` | ~2ms | baseline | Real-time, thumbnails |
254    /// | `BaselineBalanced` | ~7ms | -13% | Sequential playback |
255    /// | `ProgressiveBalanced` | ~9ms | -13% | Web images (default) |
256    /// | `ProgressiveSmallest` | ~21ms | -14% | Storage, archival |
257    ///
258    /// *Benchmarks: 512×512 Q75 image*
259    ///
260    /// # Example
261    ///
262    /// ```no_run
263    /// use mozjpeg_rs::{Encoder, Preset};
264    ///
265    /// let pixels: Vec<u8> = vec![128; 256 * 256 * 3];
266    ///
267    /// // Default: progressive with good balance
268    /// let jpeg = Encoder::new(Preset::default())
269    ///     .quality(85)
270    ///     .encode_rgb(&pixels, 256, 256)
271    ///     .unwrap();
272    ///
273    /// // Fastest for real-time applications
274    /// let jpeg = Encoder::new(Preset::BaselineFastest)
275    ///     .quality(80)
276    ///     .encode_rgb(&pixels, 256, 256)
277    ///     .unwrap();
278    ///
279    /// // Maximum compression (matches C mozjpeg)
280    /// let jpeg = Encoder::new(Preset::ProgressiveSmallest)
281    ///     .quality(85)
282    ///     .encode_rgb(&pixels, 256, 256)
283    ///     .unwrap();
284    /// ```
285    pub fn new(preset: Preset) -> Self {
286        match preset {
287            Preset::BaselineFastest => Self::fastest(),
288            Preset::BaselineBalanced => Self::baseline_optimized(),
289            Preset::ProgressiveBalanced => Self::progressive_balanced(),
290            Preset::ProgressiveSmallest => Self::max_compression(),
291        }
292    }
293
294    /// Create an encoder with the most optimized baseline (non-progressive) settings.
295    ///
296    /// This is the recommended starting point for most use cases. It produces
297    /// sequential (non-progressive) JPEGs with all mozjpeg optimizations enabled:
298    /// trellis quantization, Huffman optimization, and overshoot deringing.
299    ///
300    /// # Default Settings
301    ///
302    /// | Setting | Value | Notes |
303    /// |---------|-------|-------|
304    /// | quality | 75 | Good balance of size/quality |
305    /// | progressive | **false** | Sequential baseline JPEG |
306    /// | optimize_scans | **false** | N/A for baseline mode |
307    /// | subsampling | 4:2:0 | Standard chroma subsampling |
308    /// | trellis | **enabled** | AC + DC trellis quantization |
309    /// | optimize_huffman | **true** | 2-pass for optimal Huffman tables |
310    /// | overshoot_deringing | **true** | Reduces ringing on hard edges |
311    /// | quant_tables | ImageMagick | Same as C mozjpeg default |
312    /// | force_baseline | false | Allows 16-bit DQT at very low Q |
313    ///
314    /// # Comparison with C mozjpeg
315    ///
316    /// **Important:** This differs from C mozjpeg's `jpeg_set_defaults()`!
317    ///
318    /// C mozjpeg uses `JCP_MAX_COMPRESSION` profile by default, which enables
319    /// progressive mode and optimize_scans. This produces ~20% smaller files
320    /// but with slower encoding and progressive rendering.
321    ///
322    /// | Setting | `baseline_optimized()` | C mozjpeg default |
323    /// |---------|------------------------|-------------------|
324    /// | progressive | **false** | true |
325    /// | optimize_scans | **false** | true |
326    /// | trellis | true | true |
327    /// | deringing | true | true |
328    ///
329    /// To match C mozjpeg's default behavior, use [`max_compression()`](Self::max_compression).
330    ///
331    /// # Example
332    ///
333    /// ```no_run
334    /// use mozjpeg_rs::Encoder;
335    ///
336    /// let pixels: Vec<u8> = vec![128; 256 * 256 * 3];
337    /// let jpeg = Encoder::baseline_optimized()
338    ///     .quality(85)
339    ///     .encode_rgb(&pixels, 256, 256)
340    ///     .unwrap();
341    /// ```
342    pub fn baseline_optimized() -> Self {
343        Self {
344            quality: 75,
345            progressive: false,
346            subsampling: Subsampling::S420,
347            quant_table_idx: QuantTableIdx::ImageMagick,
348            custom_luma_qtable: None,
349            custom_chroma_qtable: None,
350            trellis: TrellisConfig::default(),
351            force_baseline: false,
352            optimize_huffman: true,
353            overshoot_deringing: true,
354            c_compat_color: true,
355            optimize_scans: false,
356            restart_interval: 0,
357            pixel_density: PixelDensity::default(),
358            exif_data: None,
359            icc_profile: None,
360            custom_markers: Vec::new(),
361            simd: SimdOps::detect(),
362            smoothing: 0,
363            limits: Limits::none(),
364        }
365    }
366
367    /// Create encoder with maximum compression (matches C mozjpeg defaults).
368    ///
369    /// This matches the `JCP_MAX_COMPRESSION` profile used by C mozjpeg's
370    /// `jpeg_set_defaults()` and the `mozjpeg` crate.
371    ///
372    /// # Settings (differences from `new()` in **bold**)
373    ///
374    /// | Setting | Value | Notes |
375    /// |---------|-------|-------|
376    /// | quality | 75 | Same as `new()` |
377    /// | progressive | **true** | Multi-scan progressive JPEG |
378    /// | optimize_scans | **true** | Tries multiple scan configs |
379    /// | subsampling | 4:2:0 | Same as `new()` |
380    /// | trellis | enabled | Same as `new()` |
381    /// | optimize_huffman | true | Same as `new()` |
382    /// | overshoot_deringing | true | Same as `new()` |
383    ///
384    /// # File Size Comparison
385    ///
386    /// Typical results at Q75 (256×256 image):
387    /// - `Encoder::baseline_optimized()`: ~650 bytes (baseline)
388    /// - `Encoder::max_compression()`: ~520 bytes (**~20% smaller**)
389    ///
390    /// # Example
391    ///
392    /// ```no_run
393    /// use mozjpeg_rs::Encoder;
394    ///
395    /// // Match C mozjpeg's default compression
396    /// let pixels: Vec<u8> = vec![128; 256 * 256 * 3];
397    /// let jpeg = Encoder::max_compression()
398    ///     .quality(85)
399    ///     .encode_rgb(&pixels, 256, 256)
400    ///     .unwrap();
401    /// ```
402    pub fn max_compression() -> Self {
403        Self {
404            quality: 75,
405            progressive: true,
406            subsampling: Subsampling::S420,
407            quant_table_idx: QuantTableIdx::ImageMagick,
408            custom_luma_qtable: None,
409            custom_chroma_qtable: None,
410            trellis: TrellisConfig::default(),
411            force_baseline: false,
412            optimize_huffman: true,
413            overshoot_deringing: true,
414            c_compat_color: true,
415            optimize_scans: true,
416            restart_interval: 0,
417            pixel_density: PixelDensity::default(),
418            exif_data: None,
419            icc_profile: None,
420            custom_markers: Vec::new(),
421            simd: SimdOps::detect(),
422            smoothing: 0,
423            limits: Limits::none(),
424        }
425    }
426
427    /// Create encoder with progressive mode and all optimizations except optimize_scans.
428    ///
429    /// This is the **recommended default** for most use cases. It provides:
430    /// - Progressive rendering (blurry-to-sharp loading)
431    /// - All mozjpeg optimizations (trellis, Huffman, deringing)
432    /// - Good balance between file size and encoding speed
433    ///
434    /// # Settings
435    ///
436    /// | Setting | Value | Notes |
437    /// |---------|-------|-------|
438    /// | progressive | **true** | Multi-scan progressive JPEG |
439    /// | optimize_scans | **false** | Uses fixed 9-scan config |
440    /// | trellis | enabled | AC + DC trellis quantization |
441    /// | optimize_huffman | true | 2-pass for optimal tables |
442    /// | overshoot_deringing | true | Reduces ringing on hard edges |
443    ///
444    /// # vs `max_compression()`
445    ///
446    /// This preset omits `optimize_scans` which:
447    /// - Saves ~100% encoding time (9ms vs 21ms at 512×512)
448    /// - Loses only ~1% file size reduction
449    ///
450    /// Use `max_compression()` only when file size is critical.
451    ///
452    /// # Example
453    ///
454    /// ```no_run
455    /// use mozjpeg_rs::Encoder;
456    ///
457    /// let pixels: Vec<u8> = vec![128; 256 * 256 * 3];
458    /// let jpeg = Encoder::progressive_balanced()
459    ///     .quality(85)
460    ///     .encode_rgb(&pixels, 256, 256)
461    ///     .unwrap();
462    /// ```
463    pub fn progressive_balanced() -> Self {
464        Self {
465            quality: 75,
466            progressive: true,
467            subsampling: Subsampling::S420,
468            quant_table_idx: QuantTableIdx::ImageMagick,
469            custom_luma_qtable: None,
470            custom_chroma_qtable: None,
471            trellis: TrellisConfig::default(),
472            force_baseline: false,
473            optimize_huffman: true,
474            overshoot_deringing: true,
475            c_compat_color: true,
476            optimize_scans: false, // Key difference from max_compression()
477            restart_interval: 0,
478            pixel_density: PixelDensity::default(),
479            exif_data: None,
480            icc_profile: None,
481            custom_markers: Vec::new(),
482            simd: SimdOps::detect(),
483            smoothing: 0,
484            limits: Limits::none(),
485        }
486    }
487
488    /// Create encoder with fastest settings (libjpeg-turbo compatible).
489    ///
490    /// Disables all mozjpeg-specific optimizations for maximum encoding speed.
491    /// Output is compatible with standard libjpeg/libjpeg-turbo.
492    ///
493    /// # Settings (differences from `new()` in **bold**)
494    ///
495    /// | Setting | Value | Notes |
496    /// |---------|-------|-------|
497    /// | quality | 75 | Same as `new()` |
498    /// | progressive | false | Same as `new()` |
499    /// | trellis | **disabled** | No trellis quantization |
500    /// | optimize_huffman | **false** | Uses default Huffman tables |
501    /// | overshoot_deringing | **false** | No deringing filter |
502    /// | force_baseline | **true** | 8-bit DQT only |
503    ///
504    /// # Performance
505    ///
506    /// Encoding is ~4-10x faster than `new()`, but files are ~10-20% larger.
507    ///
508    /// # Example
509    ///
510    /// ```no_run
511    /// use mozjpeg_rs::Encoder;
512    ///
513    /// // Fast encoding for real-time applications
514    /// let pixels: Vec<u8> = vec![128; 256 * 256 * 3];
515    /// let jpeg = Encoder::fastest()
516    ///     .quality(80)
517    ///     .encode_rgb(&pixels, 256, 256)
518    ///     .unwrap();
519    /// ```
520    pub fn fastest() -> Self {
521        Self {
522            quality: 75,
523            progressive: false,
524            subsampling: Subsampling::S420,
525            quant_table_idx: QuantTableIdx::ImageMagick,
526            custom_luma_qtable: None,
527            custom_chroma_qtable: None,
528            trellis: TrellisConfig::disabled(),
529            force_baseline: true,
530            optimize_huffman: false,
531            overshoot_deringing: false,
532            c_compat_color: true,
533            optimize_scans: false,
534            restart_interval: 0,
535            pixel_density: PixelDensity::default(),
536            exif_data: None,
537            icc_profile: None,
538            custom_markers: Vec::new(),
539            simd: SimdOps::detect(),
540            smoothing: 0,
541            limits: Limits::none(),
542        }
543    }
544
545    /// Set quality level (1-100).
546    ///
547    /// Higher values produce larger, higher-quality images.
548    pub fn quality(mut self, quality: u8) -> Self {
549        self.quality = quality.clamp(1, 100);
550        self
551    }
552
553    /// Enable or disable progressive mode.
554    pub fn progressive(mut self, enable: bool) -> Self {
555        self.progressive = enable;
556        self
557    }
558
559    /// Set chroma subsampling mode.
560    pub fn subsampling(mut self, mode: Subsampling) -> Self {
561        self.subsampling = mode;
562        self
563    }
564
565    /// Set quantization table variant.
566    pub fn quant_tables(mut self, idx: QuantTableIdx) -> Self {
567        self.quant_table_idx = idx;
568        self
569    }
570
571    /// Configure trellis quantization.
572    pub fn trellis(mut self, config: TrellisConfig) -> Self {
573        self.trellis = config;
574        self
575    }
576
577    /// Force baseline-compatible output.
578    pub fn force_baseline(mut self, enable: bool) -> Self {
579        self.force_baseline = enable;
580        self
581    }
582
583    /// Enable Huffman table optimization.
584    pub fn optimize_huffman(mut self, enable: bool) -> Self {
585        self.optimize_huffman = enable;
586        self
587    }
588
589    /// Enable overshoot deringing.
590    ///
591    /// Reduces visible ringing artifacts near hard edges, especially on white
592    /// backgrounds. Works by allowing encoded values to "overshoot" above 255
593    /// (which will clamp back to 255 when decoded) to create smoother waveforms.
594    ///
595    /// This is a mozjpeg-specific feature that can improve visual quality at
596    /// minimal file size cost. Enabled by default.
597    pub fn overshoot_deringing(mut self, enable: bool) -> Self {
598        self.overshoot_deringing = enable;
599        self
600    }
601
602    /// Use faster color conversion with the `yuv` crate.
603    ///
604    /// - `fast_color(true)` — ~40% faster RGB→YCbCr using the `yuv` crate (±1 rounding vs C mozjpeg)
605    /// - `fast_color(false)` — exact C mozjpeg parity, bytewise identical output (default)
606    ///
607    /// The ±1 rounding differences are invisible in decoded images but may cause
608    /// slightly different file sizes (typically <1% for baseline mode).
609    ///
610    /// # Example
611    ///
612    /// ```
613    /// use mozjpeg_rs::{Encoder, Preset};
614    ///
615    /// // Faster color conversion
616    /// let encoder = Encoder::new(Preset::default()).quality(85).fast_color(true);
617    ///
618    /// // Exact C mozjpeg parity (default, explicit)
619    /// let encoder = Encoder::new(Preset::default()).quality(85).fast_color(false);
620    /// ```
621    #[cfg(feature = "fast-yuv")]
622    pub fn fast_color(mut self, enable: bool) -> Self {
623        self.c_compat_color = !enable;
624        self
625    }
626
627    /// Legacy API for color conversion mode.
628    ///
629    /// **Deprecated:** Use [`fast_color()`](Self::fast_color) instead.
630    /// - `c_compat_color(true)` = `fast_color(false)` (exact C parity)
631    /// - `c_compat_color(false)` = `fast_color(true)` (faster yuv crate)
632    #[deprecated(since = "0.7.0", note = "Use fast_color() instead")]
633    pub fn c_compat_color(mut self, enable: bool) -> Self {
634        self.c_compat_color = enable;
635        self
636    }
637
638    /// Override SIMD operations dispatch for testing alternative DCT implementations.
639    ///
640    /// The default dispatch selects the best available i32-based DCT. Use this to
641    /// test experimental paths like [`SimdOps::avx2_i16()`], which uses 16-bit packed
642    /// SIMD and is vulnerable to overflow with [`overshoot_deringing`](Self::overshoot_deringing)
643    /// enabled (see [mozilla/mozjpeg#453](https://github.com/mozilla/mozjpeg/pull/453)).
644    ///
645    /// Test patterns for this bug are in the codec-corpus at
646    /// `imageflow/test_inputs/dct_overflow_patterns/`.
647    pub fn simd_ops(mut self, ops: SimdOps) -> Self {
648        self.simd = ops;
649        self
650    }
651
652    /// Enable or disable scan optimization for progressive mode.
653    ///
654    /// When enabled, the encoder tries multiple scan configurations and
655    /// picks the one that produces the smallest output. This can improve
656    /// compression by 1-3% but increases encoding time.
657    ///
658    /// Only has effect when progressive mode is enabled.
659    pub fn optimize_scans(mut self, enable: bool) -> Self {
660        self.optimize_scans = enable;
661        self
662    }
663
664    /// Set input smoothing factor (0-100).
665    ///
666    /// Applies a weighted average filter to reduce fine-scale noise in the
667    /// input image before encoding. This is particularly useful for converting
668    /// dithered images (like GIFs) to JPEG.
669    ///
670    /// - 0 = disabled (default)
671    /// - 10-50 = recommended for dithered images
672    /// - Higher values = more smoothing (may blur the image)
673    ///
674    /// # Example
675    /// ```
676    /// use mozjpeg_rs::Encoder;
677    ///
678    /// // Convert a dithered GIF to JPEG with smoothing
679    /// let encoder = Encoder::baseline_optimized()
680    ///     .quality(85)
681    ///     .smoothing(30);
682    /// ```
683    pub fn smoothing(mut self, factor: u8) -> Self {
684        self.smoothing = factor.min(100);
685        self
686    }
687
688    /// Set restart interval in MCUs.
689    ///
690    /// Restart markers are inserted every N MCUs, which can help with
691    /// error recovery and parallel decoding. Set to 0 to disable (default).
692    ///
693    /// Common values: 0 (disabled), or image width in MCUs for row-by-row restarts.
694    pub fn restart_interval(mut self, interval: u16) -> Self {
695        self.restart_interval = interval;
696        self
697    }
698
699    /// Set EXIF data to embed in the JPEG.
700    ///
701    /// # Arguments
702    /// * `data` - Raw EXIF data (TIFF structure). The "Exif\0\0" header
703    ///   will be added automatically.
704    ///
705    /// Pass empty or call without this method to omit EXIF data.
706    pub fn exif_data(mut self, data: Vec<u8>) -> Self {
707        self.exif_data = if data.is_empty() { None } else { Some(data) };
708        self
709    }
710
711    /// Set pixel density for the JFIF APP0 marker.
712    ///
713    /// This specifies the physical pixel density (DPI/DPC) or aspect ratio.
714    /// Note that most software ignores JFIF density in favor of EXIF metadata.
715    ///
716    /// # Example
717    /// ```
718    /// use mozjpeg_rs::{Encoder, PixelDensity};
719    ///
720    /// let encoder = Encoder::baseline_optimized()
721    ///     .pixel_density(PixelDensity::dpi(300, 300)); // 300 DPI
722    /// ```
723    pub fn pixel_density(mut self, density: PixelDensity) -> Self {
724        self.pixel_density = density;
725        self
726    }
727
728    /// Set ICC color profile to embed.
729    ///
730    /// The profile will be embedded in APP2 markers with the standard
731    /// "ICC_PROFILE" identifier. Large profiles are automatically chunked.
732    ///
733    /// # Arguments
734    /// * `profile` - Raw ICC profile data
735    pub fn icc_profile(mut self, profile: Vec<u8>) -> Self {
736        self.icc_profile = if profile.is_empty() {
737            None
738        } else {
739            Some(profile)
740        };
741        self
742    }
743
744    /// Add a custom APP marker.
745    ///
746    /// # Arguments
747    /// * `app_num` - APP marker number (0-15, e.g., 1 for EXIF, 2 for ICC)
748    /// * `data` - Raw marker data (including any identifier prefix)
749    ///
750    /// Multiple markers with the same number are allowed.
751    /// Markers are written in the order they are added.
752    pub fn add_marker(mut self, app_num: u8, data: Vec<u8>) -> Self {
753        if app_num <= 15 && !data.is_empty() {
754            self.custom_markers.push((app_num, data));
755        }
756        self
757    }
758
759    /// Set custom luminance quantization table.
760    ///
761    /// This overrides the table selected by `quant_tables()`.
762    /// Values should be in natural (row-major) order, not zigzag.
763    ///
764    /// # Arguments
765    /// * `table` - 64 quantization values (quality scaling still applies)
766    pub fn custom_luma_qtable(mut self, table: [u16; DCTSIZE2]) -> Self {
767        self.custom_luma_qtable = Some(table);
768        self
769    }
770
771    /// Set custom chrominance quantization table.
772    ///
773    /// This overrides the table selected by `quant_tables()`.
774    /// Values should be in natural (row-major) order, not zigzag.
775    ///
776    /// # Arguments
777    /// * `table` - 64 quantization values (quality scaling still applies)
778    pub fn custom_chroma_qtable(mut self, table: [u16; DCTSIZE2]) -> Self {
779        self.custom_chroma_qtable = Some(table);
780        self
781    }
782
783    // =========================================================================
784    // Resource Limits
785    // =========================================================================
786
787    /// Set resource limits for the encoder.
788    ///
789    /// Limits can restrict:
790    /// - Maximum image width and height
791    /// - Maximum pixel count (width × height)
792    /// - Maximum estimated memory allocation
793    /// - Maximum ICC profile size
794    ///
795    /// # Example
796    /// ```
797    /// use mozjpeg_rs::{Encoder, Preset, Limits};
798    ///
799    /// let limits = Limits::default()
800    ///     .max_width(4096)
801    ///     .max_height(4096)
802    ///     .max_pixel_count(16_000_000)
803    ///     .max_alloc_bytes(100 * 1024 * 1024);
804    ///
805    /// let encoder = Encoder::new(Preset::default())
806    ///     .limits(limits);
807    /// ```
808    pub fn limits(mut self, limits: Limits) -> Self {
809        self.limits = limits;
810        self
811    }
812
813    /// Check all resource limits before encoding.
814    ///
815    /// # Arguments
816    /// * `width` - Image width
817    /// * `height` - Image height
818    /// * `is_gray` - True for grayscale images (affects memory estimate)
819    fn check_limits(&self, width: u32, height: u32, is_gray: bool) -> Result<()> {
820        let limits = &self.limits;
821
822        // Check dimension limits
823        if (limits.max_width > 0 && width > limits.max_width)
824            || (limits.max_height > 0 && height > limits.max_height)
825        {
826            return Err(Error::DimensionLimitExceeded {
827                width,
828                height,
829                max_width: limits.max_width,
830                max_height: limits.max_height,
831            });
832        }
833
834        // Check pixel count limit
835        if limits.max_pixel_count > 0 {
836            let pixel_count = width as u64 * height as u64;
837            if pixel_count > limits.max_pixel_count {
838                return Err(Error::PixelCountExceeded {
839                    pixel_count,
840                    limit: limits.max_pixel_count,
841                });
842            }
843        }
844
845        // Check allocation limit
846        if limits.max_alloc_bytes > 0 {
847            let estimate = if is_gray {
848                self.estimate_resources_gray(width, height)
849            } else {
850                self.estimate_resources(width, height)
851            };
852            if estimate.peak_memory_bytes > limits.max_alloc_bytes {
853                return Err(Error::AllocationLimitExceeded {
854                    estimated: estimate.peak_memory_bytes,
855                    limit: limits.max_alloc_bytes,
856                });
857            }
858        }
859
860        // Check ICC profile size limit
861        if limits.max_icc_profile_bytes > 0
862            && let Some(ref icc) = self.icc_profile
863            && icc.len() > limits.max_icc_profile_bytes
864        {
865            return Err(Error::IccProfileTooLarge {
866                size: icc.len(),
867                limit: limits.max_icc_profile_bytes,
868            });
869        }
870
871        Ok(())
872    }
873
874    // =========================================================================
875    // Aliases for rimage/CLI-style naming
876    // =========================================================================
877
878    /// Set baseline mode (opposite of progressive).
879    ///
880    /// When `true`, produces a sequential JPEG (non-progressive).
881    /// This is equivalent to `progressive(false)`.
882    ///
883    /// # Example
884    /// ```
885    /// use mozjpeg_rs::Encoder;
886    ///
887    /// // These are equivalent:
888    /// let enc1 = Encoder::baseline_optimized().baseline(true);
889    /// let enc2 = Encoder::baseline_optimized().progressive(false);
890    /// ```
891    #[inline]
892    pub fn baseline(self, enable: bool) -> Self {
893        self.progressive(!enable)
894    }
895
896    /// Enable or disable Huffman coding optimization.
897    ///
898    /// Alias for [`optimize_huffman()`](Self::optimize_huffman).
899    /// This name matches mozjpeg's CLI flag naming.
900    #[inline]
901    pub fn optimize_coding(self, enable: bool) -> Self {
902        self.optimize_huffman(enable)
903    }
904
905    /// Set chroma subsampling mode.
906    ///
907    /// Alias for [`subsampling()`](Self::subsampling).
908    #[inline]
909    pub fn chroma_subsampling(self, mode: Subsampling) -> Self {
910        self.subsampling(mode)
911    }
912
913    /// Set quantization table variant.
914    ///
915    /// Alias for [`quant_tables()`](Self::quant_tables).
916    #[inline]
917    pub fn qtable(self, idx: QuantTableIdx) -> Self {
918        self.quant_tables(idx)
919    }
920
921    // =========================================================================
922    // Resource Estimation
923    // =========================================================================
924
925    /// Estimate resource usage for encoding an RGB image of the given dimensions.
926    ///
927    /// Returns peak memory usage (in bytes) and a relative CPU cost multiplier.
928    /// Useful for scheduling, enforcing resource limits, or providing feedback.
929    ///
930    /// # Arguments
931    /// * `width` - Image width in pixels
932    /// * `height` - Image height in pixels
933    ///
934    /// # Example
935    ///
936    /// ```
937    /// use mozjpeg_rs::{Encoder, Preset};
938    ///
939    /// let encoder = Encoder::new(Preset::ProgressiveBalanced).quality(85);
940    /// let estimate = encoder.estimate_resources(1920, 1080);
941    ///
942    /// println!("Peak memory: {} MB", estimate.peak_memory_bytes / 1_000_000);
943    /// println!("Relative CPU cost: {:.1}x", estimate.cpu_cost_multiplier);
944    /// ```
945    pub fn estimate_resources(&self, width: u32, height: u32) -> crate::types::ResourceEstimate {
946        let width = width as usize;
947        let height = height as usize;
948        let pixels = width * height;
949
950        // Calculate chroma dimensions based on subsampling
951        let (h_samp, v_samp) = self.subsampling.luma_factors();
952        let chroma_width = (width + h_samp as usize - 1) / h_samp as usize;
953        let chroma_height = (height + v_samp as usize - 1) / v_samp as usize;
954        let chroma_pixels = chroma_width * chroma_height;
955
956        // MCU-aligned dimensions
957        let mcu_h = 8 * h_samp as usize;
958        let mcu_v = 8 * v_samp as usize;
959        let mcu_width = (width + mcu_h - 1) / mcu_h * mcu_h;
960        let mcu_height = (height + mcu_v - 1) / mcu_v * mcu_v;
961
962        // Block counts
963        let y_blocks = (mcu_width / 8) * (mcu_height / 8);
964        let chroma_block_w = (chroma_width + 7) / 8;
965        let chroma_block_h = (chroma_height + 7) / 8;
966        let chroma_blocks = chroma_block_w * chroma_block_h;
967        let total_blocks = y_blocks + 2 * chroma_blocks;
968
969        // --- Memory estimation ---
970        let mut memory: usize = 0;
971
972        // Color conversion buffers (Y, Cb, Cr planes)
973        memory += 3 * pixels;
974
975        // Chroma subsampled buffers
976        memory += 2 * chroma_pixels;
977
978        // MCU-padded buffers
979        memory += mcu_width * mcu_height; // Y
980        let mcu_chroma_w = (chroma_width + 7) / 8 * 8;
981        let mcu_chroma_h = (chroma_height + 7) / 8 * 8;
982        memory += 2 * mcu_chroma_w * mcu_chroma_h; // Cb, Cr
983
984        // Block storage (needed for progressive or optimize_huffman)
985        let needs_block_storage = self.progressive || self.optimize_huffman;
986        if needs_block_storage {
987            // i16[64] per block = 128 bytes
988            memory += total_blocks * 128;
989        }
990
991        // Raw DCT storage (needed for DC trellis)
992        if self.trellis.dc_enabled {
993            // i32[64] per block = 256 bytes
994            memory += total_blocks * 256;
995        }
996
997        // Output buffer estimate (varies by quality, ~0.3-1.0x input for typical images)
998        // Use a conservative estimate based on quality
999        let output_ratio = if self.quality >= 95 {
1000            0.8
1001        } else if self.quality >= 85 {
1002            0.5
1003        } else if self.quality >= 75 {
1004            0.3
1005        } else {
1006            0.2
1007        };
1008        memory += (pixels as f64 * 3.0 * output_ratio) as usize;
1009
1010        // --- CPU cost estimation ---
1011        // Reference: BaselineFastest Q75 = 1.0
1012        let mut cpu_cost = 1.0;
1013
1014        // Trellis AC quantization is the biggest CPU factor
1015        if self.trellis.enabled {
1016            cpu_cost += 3.5;
1017        }
1018
1019        // DC trellis adds extra work
1020        if self.trellis.dc_enabled {
1021            cpu_cost += 0.5;
1022        }
1023
1024        // Huffman optimization (frequency counting pass)
1025        if self.optimize_huffman {
1026            cpu_cost += 0.3;
1027        }
1028
1029        // Progressive mode (multiple scan encoding)
1030        if self.progressive {
1031            cpu_cost += 1.5;
1032        }
1033
1034        // optimize_scans (trial encoding many scan configurations)
1035        if self.optimize_scans {
1036            cpu_cost += 3.0;
1037        }
1038
1039        // High quality increases trellis work (more candidates to evaluate)
1040        // This matters most when trellis is enabled
1041        if self.trellis.enabled && self.quality >= 85 {
1042            let quality_factor = 1.0 + (self.quality as f64 - 85.0) / 30.0;
1043            cpu_cost *= quality_factor;
1044        }
1045
1046        crate::types::ResourceEstimate {
1047            peak_memory_bytes: memory,
1048            cpu_cost_multiplier: cpu_cost,
1049            block_count: total_blocks,
1050        }
1051    }
1052
1053    /// Estimate resource usage for encoding a grayscale image.
1054    ///
1055    /// Similar to [`estimate_resources`](Self::estimate_resources) but for single-channel images.
1056    pub fn estimate_resources_gray(
1057        &self,
1058        width: u32,
1059        height: u32,
1060    ) -> crate::types::ResourceEstimate {
1061        let width = width as usize;
1062        let height = height as usize;
1063        let pixels = width * height;
1064
1065        // MCU-aligned dimensions (always 8x8 for grayscale)
1066        let mcu_width = (width + 7) / 8 * 8;
1067        let mcu_height = (height + 7) / 8 * 8;
1068
1069        // Block count
1070        let blocks = (mcu_width / 8) * (mcu_height / 8);
1071
1072        // --- Memory estimation ---
1073        let mut memory: usize = 0;
1074
1075        // MCU-padded buffer
1076        memory += mcu_width * mcu_height;
1077
1078        // Block storage (needed for progressive or optimize_huffman)
1079        let needs_block_storage = self.progressive || self.optimize_huffman;
1080        if needs_block_storage {
1081            memory += blocks * 128;
1082        }
1083
1084        // Raw DCT storage (needed for DC trellis)
1085        if self.trellis.dc_enabled {
1086            memory += blocks * 256;
1087        }
1088
1089        // Output buffer estimate
1090        let output_ratio = if self.quality >= 95 {
1091            0.8
1092        } else if self.quality >= 85 {
1093            0.5
1094        } else if self.quality >= 75 {
1095            0.3
1096        } else {
1097            0.2
1098        };
1099        memory += (pixels as f64 * output_ratio) as usize;
1100
1101        // --- CPU cost (same formula, but less work due to single channel) ---
1102        let mut cpu_cost = 1.0;
1103
1104        if self.trellis.enabled {
1105            cpu_cost += 3.5;
1106        }
1107        if self.trellis.dc_enabled {
1108            cpu_cost += 0.5;
1109        }
1110        if self.optimize_huffman {
1111            cpu_cost += 0.3;
1112        }
1113        if self.progressive {
1114            cpu_cost += 1.0; // Less for grayscale (fewer scans)
1115        }
1116        if self.optimize_scans {
1117            cpu_cost += 2.0; // Less for grayscale
1118        }
1119        if self.trellis.enabled && self.quality >= 85 {
1120            let quality_factor = 1.0 + (self.quality as f64 - 85.0) / 30.0;
1121            cpu_cost *= quality_factor;
1122        }
1123
1124        // Grayscale is ~1/3 the work of RGB (single channel)
1125        cpu_cost /= 3.0;
1126
1127        crate::types::ResourceEstimate {
1128            peak_memory_bytes: memory,
1129            cpu_cost_multiplier: cpu_cost,
1130            block_count: blocks,
1131        }
1132    }
1133
1134    // =========================================================================
1135    // Encoding
1136    // =========================================================================
1137
1138    /// Encode RGB image data to JPEG.
1139    ///
1140    /// # Arguments
1141    /// * `rgb_data` - RGB pixel data (3 bytes per pixel, row-major)
1142    /// * `width` - Image width in pixels
1143    /// * `height` - Image height in pixels
1144    ///
1145    /// # Returns
1146    /// JPEG-encoded data as a `Vec<u8>`.
1147    pub fn encode_rgb(&self, rgb_data: &[u8], width: u32, height: u32) -> Result<Vec<u8>> {
1148        // Validate dimensions: must be non-zero
1149        if width == 0 || height == 0 {
1150            return Err(Error::InvalidDimensions { width, height });
1151        }
1152
1153        // Check all resource limits
1154        self.check_limits(width, height, false)?;
1155
1156        // Use checked arithmetic to prevent overflow
1157        let expected_len = (width as usize)
1158            .checked_mul(height as usize)
1159            .and_then(|n| n.checked_mul(3))
1160            .ok_or(Error::InvalidDimensions { width, height })?;
1161
1162        if rgb_data.len() != expected_len {
1163            return Err(Error::BufferSizeMismatch {
1164                expected: expected_len,
1165                actual: rgb_data.len(),
1166            });
1167        }
1168
1169        // Apply smoothing if enabled
1170        let rgb_data = if self.smoothing > 0 {
1171            std::borrow::Cow::Owned(crate::smooth::smooth_rgb(
1172                rgb_data,
1173                width,
1174                height,
1175                self.smoothing,
1176            ))
1177        } else {
1178            std::borrow::Cow::Borrowed(rgb_data)
1179        };
1180
1181        let mut output = Vec::new();
1182        self.encode_rgb_to_writer(&rgb_data, width, height, &mut output)?;
1183        Ok(output)
1184    }
1185
1186    /// Encode grayscale image data to JPEG.
1187    ///
1188    /// # Arguments
1189    /// * `gray_data` - Grayscale pixel data (1 byte per pixel, row-major)
1190    /// * `width` - Image width in pixels
1191    /// * `height` - Image height in pixels
1192    ///
1193    /// # Returns
1194    /// JPEG-encoded data as a `Vec<u8>`.
1195    pub fn encode_gray(&self, gray_data: &[u8], width: u32, height: u32) -> Result<Vec<u8>> {
1196        // Validate dimensions: must be non-zero
1197        if width == 0 || height == 0 {
1198            return Err(Error::InvalidDimensions { width, height });
1199        }
1200
1201        // Check all resource limits
1202        self.check_limits(width, height, true)?;
1203
1204        // Use checked arithmetic to prevent overflow
1205        let expected_len = (width as usize)
1206            .checked_mul(height as usize)
1207            .ok_or(Error::InvalidDimensions { width, height })?;
1208
1209        if gray_data.len() != expected_len {
1210            return Err(Error::BufferSizeMismatch {
1211                expected: expected_len,
1212                actual: gray_data.len(),
1213            });
1214        }
1215
1216        // Apply smoothing if enabled
1217        let gray_data = if self.smoothing > 0 {
1218            std::borrow::Cow::Owned(crate::smooth::smooth_grayscale(
1219                gray_data,
1220                width,
1221                height,
1222                self.smoothing,
1223            ))
1224        } else {
1225            std::borrow::Cow::Borrowed(gray_data)
1226        };
1227
1228        let mut output = Vec::new();
1229        self.encode_gray_to_writer(&gray_data, width, height, &mut output)?;
1230        Ok(output)
1231    }
1232
1233    /// Encode RGB image data with row stride to JPEG.
1234    ///
1235    /// Use this when your pixel buffer has padding between rows (e.g., memory-aligned
1236    /// buffers, cropped regions without copying, GPU textures).
1237    ///
1238    /// # Arguments
1239    /// * `rgb_data` - RGB pixel data with optional row padding
1240    /// * `width` - Image width in pixels
1241    /// * `height` - Image height in pixels
1242    /// * `stride` - Number of bytes between the start of consecutive rows.
1243    ///   Must be >= `width * 3`. A stride of `width * 3` means tightly packed rows.
1244    ///
1245    /// # Returns
1246    /// JPEG-encoded data as a `Vec<u8>`.
1247    ///
1248    /// # Example
1249    /// ```no_run
1250    /// use mozjpeg_rs::{Encoder, Preset};
1251    ///
1252    /// // 100x100 image with rows padded to 320 bytes (for 64-byte alignment)
1253    /// let stride = 320;
1254    /// let buffer: Vec<u8> = vec![128; stride * 100];
1255    ///
1256    /// let jpeg = Encoder::new(Preset::default())
1257    ///     .quality(85)
1258    ///     .encode_rgb_strided(&buffer, 100, 100, stride)?;
1259    /// # Ok::<(), mozjpeg_rs::Error>(())
1260    /// ```
1261    pub fn encode_rgb_strided(
1262        &self,
1263        rgb_data: &[u8],
1264        width: u32,
1265        height: u32,
1266        stride: usize,
1267    ) -> Result<Vec<u8>> {
1268        let width_usize = width as usize;
1269        let height_usize = height as usize;
1270        let row_bytes = width_usize
1271            .checked_mul(3)
1272            .ok_or(Error::InvalidDimensions { width, height })?;
1273
1274        // Validate stride
1275        if stride < row_bytes {
1276            return Err(Error::InvalidStride {
1277                stride,
1278                minimum: row_bytes,
1279            });
1280        }
1281
1282        // If stride equals row_bytes, use the fast path
1283        if stride == row_bytes {
1284            return self.encode_rgb(rgb_data, width, height);
1285        }
1286
1287        // Validate buffer size
1288        let expected_len = stride
1289            .checked_mul(height_usize.saturating_sub(1))
1290            .and_then(|n| n.checked_add(row_bytes))
1291            .ok_or(Error::InvalidDimensions { width, height })?;
1292
1293        if rgb_data.len() < expected_len {
1294            return Err(Error::BufferSizeMismatch {
1295                expected: expected_len,
1296                actual: rgb_data.len(),
1297            });
1298        }
1299
1300        // Copy to contiguous buffer
1301        let mut contiguous = try_alloc_vec(0u8, row_bytes * height_usize)?;
1302        for y in 0..height_usize {
1303            let src_start = y * stride;
1304            let dst_start = y * row_bytes;
1305            contiguous[dst_start..dst_start + row_bytes]
1306                .copy_from_slice(&rgb_data[src_start..src_start + row_bytes]);
1307        }
1308
1309        self.encode_rgb(&contiguous, width, height)
1310    }
1311
1312    /// Encode grayscale image data with row stride to JPEG.
1313    ///
1314    /// Use this when your pixel buffer has padding between rows (e.g., memory-aligned
1315    /// buffers, cropped regions without copying).
1316    ///
1317    /// # Arguments
1318    /// * `gray_data` - Grayscale pixel data with optional row padding
1319    /// * `width` - Image width in pixels
1320    /// * `height` - Image height in pixels
1321    /// * `stride` - Number of bytes between the start of consecutive rows.
1322    ///   Must be >= `width`. A stride of `width` means tightly packed rows.
1323    ///
1324    /// # Returns
1325    /// JPEG-encoded data as a `Vec<u8>`.
1326    pub fn encode_gray_strided(
1327        &self,
1328        gray_data: &[u8],
1329        width: u32,
1330        height: u32,
1331        stride: usize,
1332    ) -> Result<Vec<u8>> {
1333        let width_usize = width as usize;
1334        let height_usize = height as usize;
1335
1336        // Validate stride
1337        if stride < width_usize {
1338            return Err(Error::InvalidStride {
1339                stride,
1340                minimum: width_usize,
1341            });
1342        }
1343
1344        // If stride equals width, use the fast path
1345        if stride == width_usize {
1346            return self.encode_gray(gray_data, width, height);
1347        }
1348
1349        // Validate buffer size
1350        let expected_len = stride
1351            .checked_mul(height_usize.saturating_sub(1))
1352            .and_then(|n| n.checked_add(width_usize))
1353            .ok_or(Error::InvalidDimensions { width, height })?;
1354
1355        if gray_data.len() < expected_len {
1356            return Err(Error::BufferSizeMismatch {
1357                expected: expected_len,
1358                actual: gray_data.len(),
1359            });
1360        }
1361
1362        // Copy to contiguous buffer
1363        let mut contiguous = try_alloc_vec(0u8, width_usize * height_usize)?;
1364        for y in 0..height_usize {
1365            let src_start = y * stride;
1366            let dst_start = y * width_usize;
1367            contiguous[dst_start..dst_start + width_usize]
1368                .copy_from_slice(&gray_data[src_start..src_start + width_usize]);
1369        }
1370
1371        self.encode_gray(&contiguous, width, height)
1372    }
1373
1374    /// Encode RGB image data to JPEG with cooperative cancellation.
1375    ///
1376    /// Accepts any [`Stop`](enough::Stop) implementation for cooperative
1377    /// cancellation. Use [`Unstoppable`](enough::Unstoppable) when
1378    /// cancellation is not needed (zero runtime cost).
1379    ///
1380    /// # Arguments
1381    /// * `rgb_data` - RGB pixel data (3 bytes per pixel, row-major)
1382    /// * `width` - Image width in pixels
1383    /// * `height` - Image height in pixels
1384    /// * `stop` - Cancellation token checked before, during, and after encoding.
1385    ///
1386    /// # Returns
1387    /// * `Ok(Vec<u8>)` - JPEG-encoded data
1388    /// * `Err(Error::Cancelled)` - If stopped via the token
1389    /// * `Err(Error::TimedOut)` - If a deadline-based stop timed out
1390    ///
1391    /// # Example
1392    /// ```no_run
1393    /// use mozjpeg_rs::{Encoder, Preset, Unstoppable};
1394    ///
1395    /// let encoder = Encoder::new(Preset::ProgressiveBalanced);
1396    /// let pixels: Vec<u8> = vec![128; 1920 * 1080 * 3];
1397    ///
1398    /// let result = encoder.encode_rgb_with_stop(&pixels, 1920, 1080, &Unstoppable);
1399    /// ```
1400    pub fn encode_rgb_with_stop(
1401        &self,
1402        rgb_data: &[u8],
1403        width: u32,
1404        height: u32,
1405        stop: &dyn enough::Stop,
1406    ) -> Result<Vec<u8>> {
1407        // Validate dimensions
1408        if width == 0 || height == 0 {
1409            return Err(Error::InvalidDimensions { width, height });
1410        }
1411
1412        // Check all resource limits
1413        self.check_limits(width, height, false)?;
1414
1415        // Check buffer size
1416        let expected_len = (width as usize)
1417            .checked_mul(height as usize)
1418            .and_then(|n| n.checked_mul(3))
1419            .ok_or(Error::InvalidDimensions { width, height })?;
1420
1421        if rgb_data.len() != expected_len {
1422            return Err(Error::BufferSizeMismatch {
1423                expected: expected_len,
1424                actual: rgb_data.len(),
1425            });
1426        }
1427
1428        // Check for immediate cancellation
1429        stop.check()?;
1430
1431        // Apply smoothing if enabled
1432        let rgb_data = if self.smoothing > 0 {
1433            std::borrow::Cow::Owned(crate::smooth::smooth_rgb(
1434                rgb_data,
1435                width,
1436                height,
1437                self.smoothing,
1438            ))
1439        } else {
1440            std::borrow::Cow::Borrowed(rgb_data)
1441        };
1442
1443        let mut output = Vec::new();
1444        // Check cancellation before and after encoding.
1445        stop.check()?;
1446        self.encode_rgb_to_writer(&rgb_data, width, height, &mut output)?;
1447        stop.check()?;
1448
1449        Ok(output)
1450    }
1451
1452    /// Encode grayscale image data to JPEG with cooperative cancellation.
1453    ///
1454    /// Accepts any [`Stop`](enough::Stop) implementation. Use
1455    /// [`Unstoppable`](enough::Unstoppable) when cancellation is not needed.
1456    ///
1457    /// # Arguments
1458    /// * `gray_data` - Grayscale pixel data (1 byte per pixel, row-major)
1459    /// * `width` - Image width in pixels
1460    /// * `height` - Image height in pixels
1461    /// * `stop` - Cancellation token checked before, during, and after encoding.
1462    ///
1463    /// # Returns
1464    /// * `Ok(Vec<u8>)` - JPEG-encoded data
1465    /// * `Err(Error::Cancelled)` - If stopped via the token
1466    /// * `Err(Error::TimedOut)` - If a deadline-based stop timed out
1467    pub fn encode_gray_with_stop(
1468        &self,
1469        gray_data: &[u8],
1470        width: u32,
1471        height: u32,
1472        stop: &dyn enough::Stop,
1473    ) -> Result<Vec<u8>> {
1474        // Validate dimensions
1475        if width == 0 || height == 0 {
1476            return Err(Error::InvalidDimensions { width, height });
1477        }
1478
1479        // Check all resource limits
1480        self.check_limits(width, height, true)?;
1481
1482        // Check buffer size
1483        let expected_len = (width as usize)
1484            .checked_mul(height as usize)
1485            .ok_or(Error::InvalidDimensions { width, height })?;
1486
1487        if gray_data.len() != expected_len {
1488            return Err(Error::BufferSizeMismatch {
1489                expected: expected_len,
1490                actual: gray_data.len(),
1491            });
1492        }
1493
1494        // Check for immediate cancellation
1495        stop.check()?;
1496
1497        // Apply smoothing if enabled
1498        let gray_data = if self.smoothing > 0 {
1499            std::borrow::Cow::Owned(crate::smooth::smooth_grayscale(
1500                gray_data,
1501                width,
1502                height,
1503                self.smoothing,
1504            ))
1505        } else {
1506            std::borrow::Cow::Borrowed(gray_data)
1507        };
1508
1509        let mut output = Vec::new();
1510        // Check cancellation before and after encoding.
1511        stop.check()?;
1512        self.encode_gray_to_writer(&gray_data, width, height, &mut output)?;
1513        stop.check()?;
1514
1515        Ok(output)
1516    }
1517
1518    /// Encode RGBA image data to JPEG.
1519    ///
1520    /// Reads 4 bytes per pixel (R, G, B, A), ignores alpha, and encodes to JPEG.
1521    /// No intermediate RGB buffer is needed — the color conversion reads RGBA directly.
1522    ///
1523    /// To encode BGRA data, swizzle to RGBA first with `garb::bytes::bgra_to_rgba`
1524    /// (or an in-place variant).
1525    ///
1526    /// # Arguments
1527    /// * `rgba_data` - RGBA pixel data (4 bytes per pixel, row-major)
1528    /// * `width` - Image width in pixels
1529    /// * `height` - Image height in pixels
1530    ///
1531    /// # Returns
1532    /// JPEG-encoded data as a `Vec<u8>`.
1533    pub fn encode_rgba(&self, rgba_data: &[u8], width: u32, height: u32) -> Result<Vec<u8>> {
1534        if width == 0 || height == 0 {
1535            return Err(Error::InvalidDimensions { width, height });
1536        }
1537        self.check_limits(width, height, false)?;
1538
1539        let expected_len = (width as usize)
1540            .checked_mul(height as usize)
1541            .and_then(|n| n.checked_mul(4))
1542            .ok_or(Error::InvalidDimensions { width, height })?;
1543
1544        if rgba_data.len() != expected_len {
1545            return Err(Error::BufferSizeMismatch {
1546                expected: expected_len,
1547                actual: rgba_data.len(),
1548            });
1549        }
1550
1551        let mut output = Vec::new();
1552        self.encode_rgba_to_writer(rgba_data, width, height, &mut output)?;
1553        Ok(output)
1554    }
1555
1556    /// Encode RGBA image data to JPEG with cooperative cancellation.
1557    ///
1558    /// Reads 4 bytes per pixel (R, G, B, A), ignores alpha.
1559    pub fn encode_rgba_with_stop(
1560        &self,
1561        rgba_data: &[u8],
1562        width: u32,
1563        height: u32,
1564        stop: &dyn enough::Stop,
1565    ) -> Result<Vec<u8>> {
1566        if width == 0 || height == 0 {
1567            return Err(Error::InvalidDimensions { width, height });
1568        }
1569        self.check_limits(width, height, false)?;
1570
1571        let expected_len = (width as usize)
1572            .checked_mul(height as usize)
1573            .and_then(|n| n.checked_mul(4))
1574            .ok_or(Error::InvalidDimensions { width, height })?;
1575
1576        if rgba_data.len() != expected_len {
1577            return Err(Error::BufferSizeMismatch {
1578                expected: expected_len,
1579                actual: rgba_data.len(),
1580            });
1581        }
1582
1583        stop.check()?;
1584        let mut output = Vec::new();
1585        self.encode_rgba_to_writer(rgba_data, width, height, &mut output)?;
1586        stop.check()?;
1587        Ok(output)
1588    }
1589
1590    /// Encode RGB image data to JPEG with cancellation and timeout support.
1591    ///
1592    /// **Deprecated:** Use [`encode_rgb_with_stop()`](Self::encode_rgb_with_stop) instead,
1593    /// which accepts any [`Stop`](enough::Stop) implementation.
1594    #[deprecated(
1595        since = "0.10.0",
1596        note = "Use encode_rgb_with_stop() with any enough::Stop implementation"
1597    )]
1598    pub fn encode_rgb_cancellable(
1599        &self,
1600        rgb_data: &[u8],
1601        width: u32,
1602        height: u32,
1603        cancel: Option<&AtomicBool>,
1604        timeout: Option<Duration>,
1605    ) -> Result<Vec<u8>> {
1606        let ctx = CancellationContext::new(cancel, timeout);
1607        self.encode_rgb_with_stop(rgb_data, width, height, &ctx)
1608    }
1609
1610    /// Encode grayscale image data to JPEG with cancellation and timeout support.
1611    ///
1612    /// **Deprecated:** Use [`encode_gray_with_stop()`](Self::encode_gray_with_stop) instead,
1613    /// which accepts any [`Stop`](enough::Stop) implementation.
1614    #[deprecated(
1615        since = "0.10.0",
1616        note = "Use encode_gray_with_stop() with any enough::Stop implementation"
1617    )]
1618    pub fn encode_gray_cancellable(
1619        &self,
1620        gray_data: &[u8],
1621        width: u32,
1622        height: u32,
1623        cancel: Option<&AtomicBool>,
1624        timeout: Option<Duration>,
1625    ) -> Result<Vec<u8>> {
1626        let ctx = CancellationContext::new(cancel, timeout);
1627        self.encode_gray_with_stop(gray_data, width, height, &ctx)
1628    }
1629
1630    /// Encode grayscale image data to a writer.
1631    pub fn encode_gray_to_writer<W: Write>(
1632        &self,
1633        gray_data: &[u8],
1634        width: u32,
1635        height: u32,
1636        output: W,
1637    ) -> Result<()> {
1638        let width = width as usize;
1639        let height = height as usize;
1640
1641        // For grayscale, Y plane is the input directly (no conversion needed)
1642        let y_plane = gray_data;
1643
1644        // Grayscale uses 1x1 sampling
1645        let (mcu_width, mcu_height) = sample::mcu_aligned_dimensions(width, height, 1, 1);
1646
1647        let mcu_y_size = mcu_width
1648            .checked_mul(mcu_height)
1649            .ok_or(Error::AllocationFailed)?;
1650        let mut y_mcu = try_alloc_vec(0u8, mcu_y_size)?;
1651        sample::expand_to_mcu(y_plane, width, height, &mut y_mcu, mcu_width, mcu_height);
1652
1653        // Create quantization table (only luma needed)
1654        let luma_qtable = if let Some(ref custom) = self.custom_luma_qtable {
1655            crate::quant::create_quant_table(custom, self.quality, self.force_baseline)
1656        } else {
1657            let (luma, _) =
1658                create_quant_tables(self.quality, self.quant_table_idx, self.force_baseline);
1659            luma
1660        };
1661
1662        // Create Huffman tables (only luma needed)
1663        let dc_luma_huff = create_std_dc_luma_table();
1664        let ac_luma_huff = create_std_ac_luma_table();
1665        let dc_luma_derived = DerivedTable::from_huff_table(&dc_luma_huff, true)?;
1666        let ac_luma_derived = DerivedTable::from_huff_table(&ac_luma_huff, false)?;
1667
1668        // Single component for grayscale
1669        let components = create_components(Subsampling::Gray);
1670
1671        // Write JPEG file
1672        let mut marker_writer = MarkerWriter::new(output);
1673
1674        // SOI
1675        marker_writer.write_soi()?;
1676
1677        // APP0 (JFIF) with pixel density
1678        marker_writer.write_jfif_app0(
1679            self.pixel_density.unit as u8,
1680            self.pixel_density.x,
1681            self.pixel_density.y,
1682        )?;
1683
1684        // EXIF (if present)
1685        if let Some(ref exif) = self.exif_data {
1686            marker_writer.write_app1_exif(exif)?;
1687        }
1688
1689        // ICC profile (if present)
1690        if let Some(ref icc) = self.icc_profile {
1691            marker_writer.write_icc_profile(icc)?;
1692        }
1693
1694        // Custom APP markers
1695        for (app_num, data) in &self.custom_markers {
1696            marker_writer.write_app(*app_num, data)?;
1697        }
1698
1699        // DQT (only luma table for grayscale)
1700        let luma_qtable_zz = natural_to_zigzag(&luma_qtable.values);
1701        marker_writer.write_dqt(0, &luma_qtable_zz, false)?;
1702
1703        // SOF (baseline or progressive)
1704        marker_writer.write_sof(
1705            self.progressive,
1706            8,
1707            height as u16,
1708            width as u16,
1709            &components,
1710        )?;
1711
1712        // DRI (restart interval)
1713        if self.restart_interval > 0 {
1714            marker_writer.write_dri(self.restart_interval)?;
1715        }
1716
1717        // DHT (only luma tables for grayscale) - written later for progressive
1718        if !self.progressive && !self.optimize_huffman {
1719            marker_writer
1720                .write_dht_multiple(&[(0, false, &dc_luma_huff), (0, true, &ac_luma_huff)])?;
1721        }
1722
1723        let mcu_rows = mcu_height / DCTSIZE;
1724        let mcu_cols = mcu_width / DCTSIZE;
1725        let num_blocks = mcu_rows
1726            .checked_mul(mcu_cols)
1727            .ok_or(Error::AllocationFailed)?;
1728
1729        if self.progressive {
1730            // Progressive mode: collect all blocks, then encode multiple scans
1731            let mut y_blocks = try_alloc_vec_array::<i16, DCTSIZE2>(num_blocks)?;
1732            let mut dct_block = [0i16; DCTSIZE2];
1733
1734            // Optionally collect raw DCT for DC trellis
1735            let dc_trellis_enabled = self.trellis.enabled && self.trellis.dc_enabled;
1736            let mut y_raw_dct = if dc_trellis_enabled {
1737                Some(try_alloc_vec_array::<i32, DCTSIZE2>(num_blocks)?)
1738            } else {
1739                None
1740            };
1741
1742            // Collect all blocks
1743            for mcu_row in 0..mcu_rows {
1744                for mcu_col in 0..mcu_cols {
1745                    let block_idx = mcu_row * mcu_cols + mcu_col;
1746                    self.process_block_to_storage_with_raw(
1747                        &y_mcu,
1748                        mcu_width,
1749                        mcu_row,
1750                        mcu_col,
1751                        &luma_qtable.values,
1752                        &ac_luma_derived,
1753                        &mut y_blocks[block_idx],
1754                        &mut dct_block,
1755                        y_raw_dct.as_mut().map(|v| v[block_idx].as_mut_slice()),
1756                    )?;
1757                }
1758            }
1759
1760            // Run DC trellis optimization if enabled
1761            if dc_trellis_enabled && let Some(ref y_raw) = y_raw_dct {
1762                run_dc_trellis_by_row(
1763                    y_raw,
1764                    &mut y_blocks,
1765                    luma_qtable.values[0],
1766                    &dc_luma_derived,
1767                    self.trellis.lambda_log_scale1,
1768                    self.trellis.lambda_log_scale2,
1769                    mcu_rows,
1770                    mcu_cols,
1771                    mcu_cols,
1772                    1,
1773                    1,
1774                    self.trellis.delta_dc_weight,
1775                );
1776            }
1777
1778            // Run EOB optimization if enabled (cross-block EOBRUN optimization)
1779            if self.trellis.enabled && self.trellis.eob_opt {
1780                use crate::trellis::{estimate_block_eob_info, optimize_eob_runs};
1781
1782                // Estimate EOB info for each block
1783                let eob_info: Vec<_> = y_blocks
1784                    .iter()
1785                    .map(|block| estimate_block_eob_info(block, &ac_luma_derived, 1, 63))
1786                    .collect();
1787
1788                // Optimize EOB runs across all blocks
1789                optimize_eob_runs(&mut y_blocks, &eob_info, &ac_luma_derived, 1, 63);
1790            }
1791
1792            // Generate progressive scan script for grayscale (1 component)
1793            let scans = generate_mozjpeg_max_compression_scans(1);
1794
1795            // Build optimized Huffman tables
1796            let mut dc_freq = FrequencyCounter::new();
1797            let mut dc_counter = ProgressiveSymbolCounter::new();
1798            for scan in &scans {
1799                let is_dc_first_scan = scan.ss == 0 && scan.se == 0 && scan.ah == 0;
1800                if is_dc_first_scan {
1801                    // Count DC symbols using progressive counter
1802                    for block in &y_blocks {
1803                        dc_counter.count_dc_first(block, 0, scan.al, &mut dc_freq);
1804                    }
1805                }
1806            }
1807
1808            let opt_dc_huff = dc_freq.generate_table()?;
1809            let opt_dc_derived = DerivedTable::from_huff_table(&opt_dc_huff, true)?;
1810
1811            // Write DC Huffman table upfront
1812            marker_writer.write_dht_multiple(&[(0, false, &opt_dc_huff)])?;
1813
1814            // Encode each scan
1815            let output = marker_writer.into_inner();
1816            let mut bit_writer = BitWriter::new(output);
1817
1818            for scan in &scans {
1819                let is_dc_scan = scan.ss == 0 && scan.se == 0;
1820
1821                if is_dc_scan {
1822                    // DC scan
1823                    marker_writer = MarkerWriter::new(bit_writer.into_inner());
1824                    marker_writer.write_sos(scan, &components)?;
1825                    bit_writer = BitWriter::new(marker_writer.into_inner());
1826
1827                    let mut prog_encoder = ProgressiveEncoder::new(&mut bit_writer);
1828
1829                    if scan.ah == 0 {
1830                        // DC first scan
1831                        for block in &y_blocks {
1832                            prog_encoder.encode_dc_first(block, 0, &opt_dc_derived, scan.al)?;
1833                        }
1834                    } else {
1835                        // DC refinement scan
1836                        for block in &y_blocks {
1837                            prog_encoder.encode_dc_refine(block, scan.al)?;
1838                        }
1839                    }
1840
1841                    prog_encoder.finish_scan(None)?;
1842                } else {
1843                    // AC scan - generate per-scan Huffman table
1844                    let mut ac_freq = FrequencyCounter::new();
1845                    let mut ac_counter = ProgressiveSymbolCounter::new();
1846
1847                    for block in &y_blocks {
1848                        if scan.ah == 0 {
1849                            ac_counter.count_ac_first(
1850                                block,
1851                                scan.ss,
1852                                scan.se,
1853                                scan.al,
1854                                &mut ac_freq,
1855                            );
1856                        } else {
1857                            ac_counter.count_ac_refine(
1858                                block,
1859                                scan.ss,
1860                                scan.se,
1861                                scan.ah,
1862                                scan.al,
1863                                &mut ac_freq,
1864                            );
1865                        }
1866                    }
1867                    ac_counter.finish_scan(Some(&mut ac_freq));
1868
1869                    let opt_ac_huff = ac_freq.generate_table()?;
1870                    let opt_ac_derived = DerivedTable::from_huff_table(&opt_ac_huff, false)?;
1871
1872                    // Write AC Huffman table and SOS
1873                    marker_writer = MarkerWriter::new(bit_writer.into_inner());
1874                    marker_writer.write_dht_multiple(&[(0, true, &opt_ac_huff)])?;
1875                    marker_writer.write_sos(scan, &components)?;
1876                    bit_writer = BitWriter::new(marker_writer.into_inner());
1877
1878                    let mut prog_encoder = ProgressiveEncoder::new(&mut bit_writer);
1879
1880                    for block in &y_blocks {
1881                        if scan.ah == 0 {
1882                            prog_encoder.encode_ac_first(
1883                                block,
1884                                scan.ss,
1885                                scan.se,
1886                                scan.al,
1887                                &opt_ac_derived,
1888                            )?;
1889                        } else {
1890                            prog_encoder.encode_ac_refine(
1891                                block,
1892                                scan.ss,
1893                                scan.se,
1894                                scan.ah,
1895                                scan.al,
1896                                &opt_ac_derived,
1897                            )?;
1898                        }
1899                    }
1900
1901                    prog_encoder.finish_scan(Some(&opt_ac_derived))?;
1902                }
1903            }
1904
1905            let mut output = bit_writer.into_inner();
1906            output.write_all(&[0xFF, 0xD9])?; // EOI
1907        } else if self.optimize_huffman {
1908            // 2-pass: collect blocks, count frequencies, then encode
1909            let mut y_blocks = try_alloc_vec_array::<i16, DCTSIZE2>(num_blocks)?;
1910            let mut dct_block = [0i16; DCTSIZE2];
1911
1912            // Collect all blocks using the same process as RGB encoding
1913            for mcu_row in 0..mcu_rows {
1914                for mcu_col in 0..mcu_cols {
1915                    let block_idx = mcu_row * mcu_cols + mcu_col;
1916                    self.process_block_to_storage_with_raw(
1917                        &y_mcu,
1918                        mcu_width,
1919                        mcu_row,
1920                        mcu_col,
1921                        &luma_qtable.values,
1922                        &ac_luma_derived,
1923                        &mut y_blocks[block_idx],
1924                        &mut dct_block,
1925                        None, // No raw DCT storage needed for grayscale
1926                    )?;
1927                }
1928            }
1929
1930            // Count frequencies using SymbolCounter
1931            let mut dc_freq = FrequencyCounter::new();
1932            let mut ac_freq = FrequencyCounter::new();
1933            let mut counter = SymbolCounter::new();
1934            for block in &y_blocks {
1935                counter.count_block(block, 0, &mut dc_freq, &mut ac_freq);
1936            }
1937
1938            // Generate optimized tables
1939            let opt_dc_huff = dc_freq.generate_table()?;
1940            let opt_ac_huff = ac_freq.generate_table()?;
1941            let opt_dc_derived = DerivedTable::from_huff_table(&opt_dc_huff, true)?;
1942            let opt_ac_derived = DerivedTable::from_huff_table(&opt_ac_huff, false)?;
1943
1944            // Write optimized Huffman tables
1945            marker_writer
1946                .write_dht_multiple(&[(0, false, &opt_dc_huff), (0, true, &opt_ac_huff)])?;
1947
1948            // Write SOS and encode
1949            let scans = generate_baseline_scan(1);
1950            marker_writer.write_sos(&scans[0], &components)?;
1951
1952            let output = marker_writer.into_inner();
1953            let mut bit_writer = BitWriter::new(output);
1954            let mut encoder = EntropyEncoder::new(&mut bit_writer);
1955
1956            // Restart marker support for grayscale (each block = 1 MCU)
1957            let restart_interval = self.restart_interval as usize;
1958            let mut restart_num = 0u8;
1959
1960            for (mcu_count, block) in y_blocks.iter().enumerate() {
1961                // Emit restart marker if needed
1962                if restart_interval > 0
1963                    && mcu_count > 0
1964                    && mcu_count.is_multiple_of(restart_interval)
1965                {
1966                    encoder.emit_restart(restart_num)?;
1967                    restart_num = restart_num.wrapping_add(1) & 0x07;
1968                }
1969                encoder.encode_block(block, 0, &opt_dc_derived, &opt_ac_derived)?;
1970            }
1971
1972            bit_writer.flush()?;
1973            let mut output = bit_writer.into_inner();
1974            output.write_all(&[0xFF, 0xD9])?; // EOI
1975        } else {
1976            // Single-pass encoding
1977            let scans = generate_baseline_scan(1);
1978            marker_writer.write_sos(&scans[0], &components)?;
1979
1980            let output = marker_writer.into_inner();
1981            let mut bit_writer = BitWriter::new(output);
1982            let mut encoder = EntropyEncoder::new(&mut bit_writer);
1983            let mut dct_block = [0i16; DCTSIZE2];
1984            let mut quant_block = [0i16; DCTSIZE2];
1985
1986            // Restart marker support
1987            let restart_interval = self.restart_interval as usize;
1988            let mut mcu_count = 0usize;
1989            let mut restart_num = 0u8;
1990
1991            for mcu_row in 0..mcu_rows {
1992                for mcu_col in 0..mcu_cols {
1993                    // Emit restart marker if needed
1994                    if restart_interval > 0
1995                        && mcu_count > 0
1996                        && mcu_count.is_multiple_of(restart_interval)
1997                    {
1998                        encoder.emit_restart(restart_num)?;
1999                        restart_num = restart_num.wrapping_add(1) & 0x07;
2000                    }
2001
2002                    // Process block directly to quant_block
2003                    self.process_block_to_storage_with_raw(
2004                        &y_mcu,
2005                        mcu_width,
2006                        mcu_row,
2007                        mcu_col,
2008                        &luma_qtable.values,
2009                        &ac_luma_derived,
2010                        &mut quant_block,
2011                        &mut dct_block,
2012                        None,
2013                    )?;
2014                    encoder.encode_block(&quant_block, 0, &dc_luma_derived, &ac_luma_derived)?;
2015                    mcu_count += 1;
2016                }
2017            }
2018
2019            bit_writer.flush()?;
2020            let mut output = bit_writer.into_inner();
2021            output.write_all(&[0xFF, 0xD9])?; // EOI
2022        }
2023
2024        Ok(())
2025    }
2026
2027    /// Encode pre-converted planar YCbCr image data to JPEG.
2028    ///
2029    /// This method accepts tightly packed YCbCr data (no row padding).
2030    /// For strided data, use [`encode_ycbcr_planar_strided`](Self::encode_ycbcr_planar_strided).
2031    ///
2032    /// # Arguments
2033    /// * `y` - Luma plane (width × height bytes, tightly packed)
2034    /// * `cb` - Cb chroma plane (chroma_width × chroma_height bytes)
2035    /// * `cr` - Cr chroma plane (chroma_width × chroma_height bytes)
2036    /// * `width` - Image width in pixels
2037    /// * `height` - Image height in pixels
2038    ///
2039    /// The chroma plane dimensions depend on the subsampling mode:
2040    /// - 4:4:4: chroma_width = width, chroma_height = height
2041    /// - 4:2:2: chroma_width = ceil(width/2), chroma_height = height
2042    /// - 4:2:0: chroma_width = ceil(width/2), chroma_height = ceil(height/2)
2043    ///
2044    /// # Returns
2045    /// JPEG-encoded data as a `Vec<u8>`.
2046    ///
2047    /// # Errors
2048    /// Returns an error if plane sizes don't match expected dimensions.
2049    pub fn encode_ycbcr_planar(
2050        &self,
2051        y: &[u8],
2052        cb: &[u8],
2053        cr: &[u8],
2054        width: u32,
2055        height: u32,
2056    ) -> Result<Vec<u8>> {
2057        // For packed data, stride equals width
2058        let (luma_h, luma_v) = self.subsampling.luma_factors();
2059        let (chroma_width, _) = sample::subsampled_dimensions(
2060            width as usize,
2061            height as usize,
2062            luma_h as usize,
2063            luma_v as usize,
2064        );
2065        self.encode_ycbcr_planar_strided(
2066            y,
2067            width as usize,
2068            cb,
2069            chroma_width,
2070            cr,
2071            chroma_width,
2072            width,
2073            height,
2074        )
2075    }
2076
2077    /// Encode pre-converted planar YCbCr image data to a writer.
2078    ///
2079    /// See [`encode_ycbcr_planar`](Self::encode_ycbcr_planar) for details.
2080    pub fn encode_ycbcr_planar_to_writer<W: Write>(
2081        &self,
2082        y: &[u8],
2083        cb: &[u8],
2084        cr: &[u8],
2085        width: u32,
2086        height: u32,
2087        output: W,
2088    ) -> Result<()> {
2089        // For packed data, stride equals width
2090        let (luma_h, luma_v) = self.subsampling.luma_factors();
2091        let (chroma_width, _) = sample::subsampled_dimensions(
2092            width as usize,
2093            height as usize,
2094            luma_h as usize,
2095            luma_v as usize,
2096        );
2097        self.encode_ycbcr_planar_strided_to_writer(
2098            y,
2099            width as usize,
2100            cb,
2101            chroma_width,
2102            cr,
2103            chroma_width,
2104            width,
2105            height,
2106            output,
2107        )
2108    }
2109
2110    /// Encode pre-converted planar YCbCr image data with arbitrary strides.
2111    ///
2112    /// This method accepts YCbCr data that has already been:
2113    /// 1. Converted from RGB to YCbCr color space
2114    /// 2. Downsampled according to the encoder's subsampling mode
2115    ///
2116    /// Use this when you have YCbCr data from video decoders or other sources
2117    /// that may have row padding (stride > width).
2118    ///
2119    /// # Arguments
2120    /// * `y` - Luma plane data
2121    /// * `y_stride` - Bytes per row in luma plane (must be >= width)
2122    /// * `cb` - Cb chroma plane data
2123    /// * `cb_stride` - Bytes per row in Cb plane (must be >= chroma_width)
2124    /// * `cr` - Cr chroma plane data
2125    /// * `cr_stride` - Bytes per row in Cr plane (must be >= chroma_width)
2126    /// * `width` - Image width in pixels
2127    /// * `height` - Image height in pixels
2128    ///
2129    /// The chroma plane dimensions depend on the subsampling mode:
2130    /// - 4:4:4: chroma_width = width, chroma_height = height
2131    /// - 4:2:2: chroma_width = ceil(width/2), chroma_height = height
2132    /// - 4:2:0: chroma_width = ceil(width/2), chroma_height = ceil(height/2)
2133    ///
2134    /// # Returns
2135    /// JPEG-encoded data as a `Vec<u8>`.
2136    ///
2137    /// # Errors
2138    /// Returns an error if:
2139    /// - Strides are less than the required width
2140    /// - Plane sizes don't match stride × height
2141    #[allow(clippy::too_many_arguments)]
2142    pub fn encode_ycbcr_planar_strided(
2143        &self,
2144        y: &[u8],
2145        y_stride: usize,
2146        cb: &[u8],
2147        cb_stride: usize,
2148        cr: &[u8],
2149        cr_stride: usize,
2150        width: u32,
2151        height: u32,
2152    ) -> Result<Vec<u8>> {
2153        let mut output = Vec::new();
2154        self.encode_ycbcr_planar_strided_to_writer(
2155            y,
2156            y_stride,
2157            cb,
2158            cb_stride,
2159            cr,
2160            cr_stride,
2161            width,
2162            height,
2163            &mut output,
2164        )?;
2165        Ok(output)
2166    }
2167
2168    /// Encode pre-converted planar YCbCr image data with arbitrary strides to a writer.
2169    ///
2170    /// See [`encode_ycbcr_planar_strided`](Self::encode_ycbcr_planar_strided) for details.
2171    #[allow(clippy::too_many_arguments)]
2172    pub fn encode_ycbcr_planar_strided_to_writer<W: Write>(
2173        &self,
2174        y: &[u8],
2175        y_stride: usize,
2176        cb: &[u8],
2177        cb_stride: usize,
2178        cr: &[u8],
2179        cr_stride: usize,
2180        width: u32,
2181        height: u32,
2182        output: W,
2183    ) -> Result<()> {
2184        let width = width as usize;
2185        let height = height as usize;
2186
2187        // Validate dimensions
2188        if width == 0 || height == 0 {
2189            return Err(Error::InvalidDimensions {
2190                width: width as u32,
2191                height: height as u32,
2192            });
2193        }
2194
2195        // Validate Y stride
2196        if y_stride < width {
2197            return Err(Error::InvalidSamplingFactor {
2198                h: y_stride as u8,
2199                v: width as u8,
2200            });
2201        }
2202
2203        let (luma_h, luma_v) = self.subsampling.luma_factors();
2204        let (chroma_width, chroma_height) =
2205            sample::subsampled_dimensions(width, height, luma_h as usize, luma_v as usize);
2206
2207        // Validate chroma strides
2208        if cb_stride < chroma_width {
2209            return Err(Error::InvalidSamplingFactor {
2210                h: cb_stride as u8,
2211                v: chroma_width as u8,
2212            });
2213        }
2214        if cr_stride < chroma_width {
2215            return Err(Error::InvalidSamplingFactor {
2216                h: cr_stride as u8,
2217                v: chroma_width as u8,
2218            });
2219        }
2220
2221        // Calculate expected plane sizes (stride × height)
2222        let y_size = y_stride
2223            .checked_mul(height)
2224            .ok_or(Error::InvalidDimensions {
2225                width: width as u32,
2226                height: height as u32,
2227            })?;
2228        let cb_size = cb_stride
2229            .checked_mul(chroma_height)
2230            .ok_or(Error::AllocationFailed)?;
2231        let cr_size = cr_stride
2232            .checked_mul(chroma_height)
2233            .ok_or(Error::AllocationFailed)?;
2234
2235        // Validate Y plane size
2236        if y.len() < y_size {
2237            return Err(Error::BufferSizeMismatch {
2238                expected: y_size,
2239                actual: y.len(),
2240            });
2241        }
2242
2243        // Validate Cb plane size
2244        if cb.len() < cb_size {
2245            return Err(Error::BufferSizeMismatch {
2246                expected: cb_size,
2247                actual: cb.len(),
2248            });
2249        }
2250
2251        // Validate Cr plane size
2252        if cr.len() < cr_size {
2253            return Err(Error::BufferSizeMismatch {
2254                expected: cr_size,
2255                actual: cr.len(),
2256            });
2257        }
2258
2259        // Expand planes to MCU-aligned dimensions
2260        let (mcu_width, mcu_height) =
2261            sample::mcu_aligned_dimensions(width, height, luma_h as usize, luma_v as usize);
2262        let (mcu_chroma_w, mcu_chroma_h) =
2263            (mcu_width / luma_h as usize, mcu_height / luma_v as usize);
2264
2265        let mcu_y_size = mcu_width
2266            .checked_mul(mcu_height)
2267            .ok_or(Error::AllocationFailed)?;
2268        let mcu_chroma_size = mcu_chroma_w
2269            .checked_mul(mcu_chroma_h)
2270            .ok_or(Error::AllocationFailed)?;
2271        let mut y_mcu = try_alloc_vec(0u8, mcu_y_size)?;
2272        let mut cb_mcu = try_alloc_vec(0u8, mcu_chroma_size)?;
2273        let mut cr_mcu = try_alloc_vec(0u8, mcu_chroma_size)?;
2274
2275        sample::expand_to_mcu_strided(
2276            y, width, y_stride, height, &mut y_mcu, mcu_width, mcu_height,
2277        );
2278        sample::expand_to_mcu_strided(
2279            cb,
2280            chroma_width,
2281            cb_stride,
2282            chroma_height,
2283            &mut cb_mcu,
2284            mcu_chroma_w,
2285            mcu_chroma_h,
2286        );
2287        sample::expand_to_mcu_strided(
2288            cr,
2289            chroma_width,
2290            cr_stride,
2291            chroma_height,
2292            &mut cr_mcu,
2293            mcu_chroma_w,
2294            mcu_chroma_h,
2295        );
2296
2297        // Encode using shared helper
2298        self.encode_ycbcr_mcu_to_writer(
2299            &y_mcu,
2300            &cb_mcu,
2301            &cr_mcu,
2302            width,
2303            height,
2304            mcu_width,
2305            mcu_height,
2306            chroma_width,
2307            chroma_height,
2308            mcu_chroma_w,
2309            mcu_chroma_h,
2310            output,
2311        )
2312    }
2313
2314    /// Encode RGB image data to a writer.
2315    pub fn encode_rgb_to_writer<W: Write>(
2316        &self,
2317        rgb_data: &[u8],
2318        width: u32,
2319        height: u32,
2320        output: W,
2321    ) -> Result<()> {
2322        let width = width as usize;
2323        let height = height as usize;
2324
2325        let num_pixels = width.checked_mul(height).ok_or(Error::InvalidDimensions {
2326            width: width as u32,
2327            height: height as u32,
2328        })?;
2329
2330        let mut y_plane = try_alloc_vec(0u8, num_pixels)?;
2331        let mut cb_plane = try_alloc_vec(0u8, num_pixels)?;
2332        let mut cr_plane = try_alloc_vec(0u8, num_pixels)?;
2333
2334        if self.c_compat_color {
2335            convert_rgb_to_ycbcr_c_compat(
2336                rgb_data,
2337                &mut y_plane,
2338                &mut cb_plane,
2339                &mut cr_plane,
2340                num_pixels,
2341            );
2342        } else {
2343            (self.simd.color_convert_rgb_to_ycbcr)(
2344                rgb_data,
2345                &mut y_plane,
2346                &mut cb_plane,
2347                &mut cr_plane,
2348                num_pixels,
2349            );
2350        }
2351
2352        self.encode_ycbcr_planes_to_writer(&y_plane, &cb_plane, &cr_plane, width, height, output)
2353    }
2354
2355    /// Encode RGBA image data to a writer.
2356    ///
2357    /// Reads 4 bytes per pixel (R, G, B, A), ignores alpha, and encodes to JPEG.
2358    /// Uses the C mozjpeg-compatible color conversion (no intermediate RGB buffer).
2359    pub fn encode_rgba_to_writer<W: Write>(
2360        &self,
2361        rgba_data: &[u8],
2362        width: u32,
2363        height: u32,
2364        output: W,
2365    ) -> Result<()> {
2366        let width = width as usize;
2367        let height = height as usize;
2368
2369        let num_pixels = width.checked_mul(height).ok_or(Error::InvalidDimensions {
2370            width: width as u32,
2371            height: height as u32,
2372        })?;
2373
2374        let mut y_plane = try_alloc_vec(0u8, num_pixels)?;
2375        let mut cb_plane = try_alloc_vec(0u8, num_pixels)?;
2376        let mut cr_plane = try_alloc_vec(0u8, num_pixels)?;
2377
2378        crate::color::convert_rgba_to_ycbcr_c_compat(
2379            rgba_data,
2380            &mut y_plane,
2381            &mut cb_plane,
2382            &mut cr_plane,
2383            num_pixels,
2384        );
2385
2386        self.encode_ycbcr_planes_to_writer(&y_plane, &cb_plane, &cr_plane, width, height, output)
2387    }
2388
2389    /// Internal helper: downsample, MCU-align, and encode Y/Cb/Cr planes.
2390    fn encode_ycbcr_planes_to_writer<W: Write>(
2391        &self,
2392        y_plane: &[u8],
2393        cb_plane: &[u8],
2394        cr_plane: &[u8],
2395        width: usize,
2396        height: usize,
2397        output: W,
2398    ) -> Result<()> {
2399        let (luma_h, luma_v) = self.subsampling.luma_factors();
2400        let (chroma_width, chroma_height) =
2401            sample::subsampled_dimensions(width, height, luma_h as usize, luma_v as usize);
2402
2403        let chroma_size = chroma_width
2404            .checked_mul(chroma_height)
2405            .ok_or(Error::AllocationFailed)?;
2406        let mut cb_subsampled = try_alloc_vec(0u8, chroma_size)?;
2407        let mut cr_subsampled = try_alloc_vec(0u8, chroma_size)?;
2408
2409        sample::downsample_plane(
2410            cb_plane,
2411            width,
2412            height,
2413            luma_h as usize,
2414            luma_v as usize,
2415            &mut cb_subsampled,
2416        );
2417        sample::downsample_plane(
2418            cr_plane,
2419            width,
2420            height,
2421            luma_h as usize,
2422            luma_v as usize,
2423            &mut cr_subsampled,
2424        );
2425
2426        let (mcu_width, mcu_height) =
2427            sample::mcu_aligned_dimensions(width, height, luma_h as usize, luma_v as usize);
2428        let (mcu_chroma_w, mcu_chroma_h) =
2429            (mcu_width / luma_h as usize, mcu_height / luma_v as usize);
2430
2431        let mcu_y_size = mcu_width
2432            .checked_mul(mcu_height)
2433            .ok_or(Error::AllocationFailed)?;
2434        let mcu_chroma_size = mcu_chroma_w
2435            .checked_mul(mcu_chroma_h)
2436            .ok_or(Error::AllocationFailed)?;
2437        let mut y_mcu = try_alloc_vec(0u8, mcu_y_size)?;
2438        let mut cb_mcu = try_alloc_vec(0u8, mcu_chroma_size)?;
2439        let mut cr_mcu = try_alloc_vec(0u8, mcu_chroma_size)?;
2440
2441        sample::expand_to_mcu(y_plane, width, height, &mut y_mcu, mcu_width, mcu_height);
2442        sample::expand_to_mcu(
2443            &cb_subsampled,
2444            chroma_width,
2445            chroma_height,
2446            &mut cb_mcu,
2447            mcu_chroma_w,
2448            mcu_chroma_h,
2449        );
2450        sample::expand_to_mcu(
2451            &cr_subsampled,
2452            chroma_width,
2453            chroma_height,
2454            &mut cr_mcu,
2455            mcu_chroma_w,
2456            mcu_chroma_h,
2457        );
2458
2459        self.encode_ycbcr_mcu_to_writer(
2460            &y_mcu,
2461            &cb_mcu,
2462            &cr_mcu,
2463            width,
2464            height,
2465            mcu_width,
2466            mcu_height,
2467            chroma_width,
2468            chroma_height,
2469            mcu_chroma_w,
2470            mcu_chroma_h,
2471            output,
2472        )
2473    }
2474
2475    /// Internal helper: Encode MCU-aligned YCbCr planes to JPEG.
2476    ///
2477    /// This is the shared encoding logic used by both `encode_rgb_to_writer`
2478    /// and `encode_ycbcr_planar_to_writer`.
2479    #[allow(clippy::too_many_arguments)]
2480    fn encode_ycbcr_mcu_to_writer<W: Write>(
2481        &self,
2482        y_mcu: &[u8],
2483        cb_mcu: &[u8],
2484        cr_mcu: &[u8],
2485        width: usize,
2486        height: usize,
2487        mcu_width: usize,
2488        mcu_height: usize,
2489        chroma_width: usize,
2490        chroma_height: usize,
2491        mcu_chroma_w: usize,
2492        mcu_chroma_h: usize,
2493        output: W,
2494    ) -> Result<()> {
2495        let (luma_h, luma_v) = self.subsampling.luma_factors();
2496
2497        // Step 4: Create quantization tables
2498        let (luma_qtable, chroma_qtable) = {
2499            let (default_luma, default_chroma) =
2500                create_quant_tables(self.quality, self.quant_table_idx, self.force_baseline);
2501            let luma = if let Some(ref custom) = self.custom_luma_qtable {
2502                crate::quant::create_quant_table(custom, self.quality, self.force_baseline)
2503            } else {
2504                default_luma
2505            };
2506            let chroma = if let Some(ref custom) = self.custom_chroma_qtable {
2507                crate::quant::create_quant_table(custom, self.quality, self.force_baseline)
2508            } else {
2509                default_chroma
2510            };
2511            (luma, chroma)
2512        };
2513
2514        // Step 5: Create Huffman tables (standard tables)
2515        let dc_luma_huff = create_std_dc_luma_table();
2516        let dc_chroma_huff = create_std_dc_chroma_table();
2517        let ac_luma_huff = create_std_ac_luma_table();
2518        let ac_chroma_huff = create_std_ac_chroma_table();
2519
2520        let dc_luma_derived = DerivedTable::from_huff_table(&dc_luma_huff, true)?;
2521        let dc_chroma_derived = DerivedTable::from_huff_table(&dc_chroma_huff, true)?;
2522        let ac_luma_derived = DerivedTable::from_huff_table(&ac_luma_huff, false)?;
2523        let ac_chroma_derived = DerivedTable::from_huff_table(&ac_chroma_huff, false)?;
2524
2525        // Step 6: Set up components
2526        let components = create_ycbcr_components(self.subsampling);
2527
2528        // Step 7: Write JPEG file
2529        let mut marker_writer = MarkerWriter::new(output);
2530
2531        // SOI
2532        marker_writer.write_soi()?;
2533
2534        // APP0 (JFIF) with pixel density
2535        marker_writer.write_jfif_app0(
2536            self.pixel_density.unit as u8,
2537            self.pixel_density.x,
2538            self.pixel_density.y,
2539        )?;
2540
2541        // APP1 (EXIF) - if present
2542        if let Some(ref exif) = self.exif_data {
2543            marker_writer.write_app1_exif(exif)?;
2544        }
2545
2546        // ICC profile (if present)
2547        if let Some(ref icc) = self.icc_profile {
2548            marker_writer.write_icc_profile(icc)?;
2549        }
2550
2551        // Custom APP markers
2552        for (app_num, data) in &self.custom_markers {
2553            marker_writer.write_app(*app_num, data)?;
2554        }
2555
2556        // DQT (quantization tables in zigzag order) - combined into single marker
2557        let luma_qtable_zz = natural_to_zigzag(&luma_qtable.values);
2558        let chroma_qtable_zz = natural_to_zigzag(&chroma_qtable.values);
2559        marker_writer
2560            .write_dqt_multiple(&[(0, &luma_qtable_zz, false), (1, &chroma_qtable_zz, false)])?;
2561
2562        // SOF
2563        marker_writer.write_sof(
2564            self.progressive,
2565            8,
2566            height as u16,
2567            width as u16,
2568            &components,
2569        )?;
2570
2571        // DRI (restart interval) - if enabled
2572        if self.restart_interval > 0 {
2573            marker_writer.write_dri(self.restart_interval)?;
2574        }
2575
2576        // DHT (Huffman tables) - written here for non-optimized modes,
2577        // or later after frequency counting for optimized modes
2578        if !self.optimize_huffman {
2579            // Combine all tables into single DHT marker for smaller file size
2580            marker_writer.write_dht_multiple(&[
2581                (0, false, &dc_luma_huff),
2582                (1, false, &dc_chroma_huff),
2583                (0, true, &ac_luma_huff),
2584                (1, true, &ac_chroma_huff),
2585            ])?;
2586        }
2587
2588        if self.progressive {
2589            // Progressive mode: Store all blocks, then encode multiple scans
2590            let mcu_rows = mcu_height / (DCTSIZE * luma_v as usize);
2591            let mcu_cols = mcu_width / (DCTSIZE * luma_h as usize);
2592            let num_y_blocks = mcu_rows
2593                .checked_mul(mcu_cols)
2594                .and_then(|n| n.checked_mul(luma_h as usize))
2595                .and_then(|n| n.checked_mul(luma_v as usize))
2596                .ok_or(Error::AllocationFailed)?;
2597            let num_chroma_blocks = mcu_rows
2598                .checked_mul(mcu_cols)
2599                .ok_or(Error::AllocationFailed)?;
2600
2601            // Collect all quantized blocks
2602            let mut y_blocks = try_alloc_vec_array::<i16, DCTSIZE2>(num_y_blocks)?;
2603            let mut cb_blocks = try_alloc_vec_array::<i16, DCTSIZE2>(num_chroma_blocks)?;
2604            let mut cr_blocks = try_alloc_vec_array::<i16, DCTSIZE2>(num_chroma_blocks)?;
2605
2606            // Optionally collect raw DCT for DC trellis
2607            let dc_trellis_enabled = self.trellis.enabled && self.trellis.dc_enabled;
2608            let mut y_raw_dct = if dc_trellis_enabled {
2609                Some(try_alloc_vec_array::<i32, DCTSIZE2>(num_y_blocks)?)
2610            } else {
2611                None
2612            };
2613            let mut cb_raw_dct = if dc_trellis_enabled {
2614                Some(try_alloc_vec_array::<i32, DCTSIZE2>(num_chroma_blocks)?)
2615            } else {
2616                None
2617            };
2618            let mut cr_raw_dct = if dc_trellis_enabled {
2619                Some(try_alloc_vec_array::<i32, DCTSIZE2>(num_chroma_blocks)?)
2620            } else {
2621                None
2622            };
2623
2624            self.collect_blocks(
2625                y_mcu,
2626                mcu_width,
2627                mcu_height,
2628                cb_mcu,
2629                cr_mcu,
2630                mcu_chroma_w,
2631                mcu_chroma_h,
2632                &luma_qtable.values,
2633                &chroma_qtable.values,
2634                &ac_luma_derived,
2635                &ac_chroma_derived,
2636                &mut y_blocks,
2637                &mut cb_blocks,
2638                &mut cr_blocks,
2639                y_raw_dct.as_deref_mut(),
2640                cb_raw_dct.as_deref_mut(),
2641                cr_raw_dct.as_deref_mut(),
2642                luma_h,
2643                luma_v,
2644            )?;
2645
2646            // Run DC trellis optimization if enabled
2647            // C mozjpeg processes DC trellis row by row (each row is an independent chain)
2648            if dc_trellis_enabled {
2649                let h = luma_h as usize;
2650                let v = luma_v as usize;
2651                let y_block_cols = mcu_cols * h;
2652                let y_block_rows = mcu_rows * v;
2653
2654                if let Some(ref y_raw) = y_raw_dct {
2655                    run_dc_trellis_by_row(
2656                        y_raw,
2657                        &mut y_blocks,
2658                        luma_qtable.values[0],
2659                        &dc_luma_derived,
2660                        self.trellis.lambda_log_scale1,
2661                        self.trellis.lambda_log_scale2,
2662                        y_block_rows,
2663                        y_block_cols,
2664                        mcu_cols,
2665                        h,
2666                        v,
2667                        self.trellis.delta_dc_weight,
2668                    );
2669                }
2670                // Chroma has 1x1 per MCU, so MCU order = row order
2671                if let Some(ref cb_raw) = cb_raw_dct {
2672                    run_dc_trellis_by_row(
2673                        cb_raw,
2674                        &mut cb_blocks,
2675                        chroma_qtable.values[0],
2676                        &dc_chroma_derived,
2677                        self.trellis.lambda_log_scale1,
2678                        self.trellis.lambda_log_scale2,
2679                        mcu_rows,
2680                        mcu_cols,
2681                        mcu_cols,
2682                        1,
2683                        1,
2684                        self.trellis.delta_dc_weight,
2685                    );
2686                }
2687                if let Some(ref cr_raw) = cr_raw_dct {
2688                    run_dc_trellis_by_row(
2689                        cr_raw,
2690                        &mut cr_blocks,
2691                        chroma_qtable.values[0],
2692                        &dc_chroma_derived,
2693                        self.trellis.lambda_log_scale1,
2694                        self.trellis.lambda_log_scale2,
2695                        mcu_rows,
2696                        mcu_cols,
2697                        mcu_cols,
2698                        1,
2699                        1,
2700                        self.trellis.delta_dc_weight,
2701                    );
2702                }
2703            }
2704
2705            // Run EOB optimization if enabled (cross-block EOBRUN optimization)
2706            if self.trellis.enabled && self.trellis.eob_opt {
2707                use crate::trellis::{estimate_block_eob_info, optimize_eob_runs};
2708
2709                // Y component
2710                let y_eob_info: Vec<_> = y_blocks
2711                    .iter()
2712                    .map(|block| estimate_block_eob_info(block, &ac_luma_derived, 1, 63))
2713                    .collect();
2714                optimize_eob_runs(&mut y_blocks, &y_eob_info, &ac_luma_derived, 1, 63);
2715
2716                // Cb component
2717                let cb_eob_info: Vec<_> = cb_blocks
2718                    .iter()
2719                    .map(|block| estimate_block_eob_info(block, &ac_chroma_derived, 1, 63))
2720                    .collect();
2721                optimize_eob_runs(&mut cb_blocks, &cb_eob_info, &ac_chroma_derived, 1, 63);
2722
2723                // Cr component
2724                let cr_eob_info: Vec<_> = cr_blocks
2725                    .iter()
2726                    .map(|block| estimate_block_eob_info(block, &ac_chroma_derived, 1, 63))
2727                    .collect();
2728                optimize_eob_runs(&mut cr_blocks, &cr_eob_info, &ac_chroma_derived, 1, 63);
2729            }
2730
2731            // Generate progressive scan script
2732            let scans = if self.optimize_scans {
2733                // When optimize_scans is enabled, use the scan optimizer to find
2734                // the best frequency split and Al levels, including SA refinement.
2735                self.optimize_progressive_scans(
2736                    3, // num_components
2737                    &y_blocks,
2738                    &cb_blocks,
2739                    &cr_blocks,
2740                    mcu_rows,
2741                    mcu_cols,
2742                    luma_h,
2743                    luma_v,
2744                    width,
2745                    height,
2746                    chroma_width,
2747                    chroma_height,
2748                    &dc_luma_derived,
2749                    &dc_chroma_derived,
2750                    &ac_luma_derived,
2751                    &ac_chroma_derived,
2752                )?
2753            } else {
2754                // Use C mozjpeg's 9-scan JCP_MAX_COMPRESSION script.
2755                // This matches jcparam.c lines 932-947 (the JCP_MAX_COMPRESSION branch).
2756                // mozjpeg-sys defaults to JCP_MAX_COMPRESSION profile, which uses:
2757                // - DC with no successive approximation (Al=0)
2758                // - 8/9 frequency split for luma with successive approximation
2759                // - No successive approximation for chroma
2760                generate_mozjpeg_max_compression_scans(3)
2761            };
2762
2763            // Build Huffman tables and encode scans
2764            //
2765            // When optimize_scans=true, each AC scan gets its own optimal Huffman table
2766            // written immediately before the scan. This matches C mozjpeg behavior and
2767            // ensures the trial encoder's size estimates match actual encoded sizes.
2768            //
2769            // When optimize_huffman=true, use per-scan AC tables (matching C mozjpeg).
2770            // C automatically enables optimize_coding for progressive mode and does
2771            // 2 passes per scan: gather statistics, then output with optimal tables.
2772
2773            if self.optimize_huffman {
2774                // Per-scan AC tables mode: DC tables global, AC tables per-scan
2775                // This matches C mozjpeg's progressive behavior
2776
2777                // Count DC frequencies for first-pass DC scans only (Ah == 0)
2778                // DC refinement scans (Ah > 0) don't use Huffman coding - they output raw bits
2779                let mut dc_luma_freq = FrequencyCounter::new();
2780                let mut dc_chroma_freq = FrequencyCounter::new();
2781
2782                for scan in &scans {
2783                    let is_dc_first_scan = scan.ss == 0 && scan.se == 0 && scan.ah == 0;
2784                    if is_dc_first_scan {
2785                        self.count_dc_scan_symbols(
2786                            scan,
2787                            &y_blocks,
2788                            &cb_blocks,
2789                            &cr_blocks,
2790                            mcu_rows,
2791                            mcu_cols,
2792                            luma_h,
2793                            luma_v,
2794                            &mut dc_luma_freq,
2795                            &mut dc_chroma_freq,
2796                        );
2797                    }
2798                }
2799
2800                // Generate and write DC tables upfront
2801                let opt_dc_luma_huff = dc_luma_freq.generate_table()?;
2802                let opt_dc_chroma_huff = dc_chroma_freq.generate_table()?;
2803                marker_writer.write_dht_multiple(&[
2804                    (0, false, &opt_dc_luma_huff),
2805                    (1, false, &opt_dc_chroma_huff),
2806                ])?;
2807
2808                let opt_dc_luma = DerivedTable::from_huff_table(&opt_dc_luma_huff, true)?;
2809                let opt_dc_chroma = DerivedTable::from_huff_table(&opt_dc_chroma_huff, true)?;
2810
2811                // Get output writer from marker_writer
2812                let output = marker_writer.into_inner();
2813                let mut bit_writer = BitWriter::new(output);
2814
2815                // Encode each scan with per-scan AC tables
2816                for scan in &scans {
2817                    bit_writer.flush()?;
2818                    let mut inner = bit_writer.into_inner();
2819
2820                    let is_dc_scan = scan.ss == 0 && scan.se == 0;
2821
2822                    if !is_dc_scan {
2823                        // AC scan: build per-scan optimal Huffman table
2824                        let comp_idx = scan.component_index[0] as usize;
2825                        let blocks = match comp_idx {
2826                            0 => &y_blocks,
2827                            1 => &cb_blocks,
2828                            2 => &cr_blocks,
2829                            _ => &y_blocks,
2830                        };
2831                        let (block_cols, block_rows) = if comp_idx == 0 {
2832                            (width.div_ceil(DCTSIZE), height.div_ceil(DCTSIZE))
2833                        } else {
2834                            (
2835                                chroma_width.div_ceil(DCTSIZE),
2836                                chroma_height.div_ceil(DCTSIZE),
2837                            )
2838                        };
2839
2840                        // Count frequencies for this scan only
2841                        let mut ac_freq = FrequencyCounter::new();
2842                        self.count_ac_scan_symbols(
2843                            scan,
2844                            blocks,
2845                            mcu_rows,
2846                            mcu_cols,
2847                            luma_h,
2848                            luma_v,
2849                            comp_idx,
2850                            block_cols,
2851                            block_rows,
2852                            &mut ac_freq,
2853                        );
2854
2855                        // Build optimal table and write DHT
2856                        let ac_huff = ac_freq.generate_table()?;
2857                        let table_idx = if comp_idx == 0 { 0 } else { 1 };
2858                        write_dht_marker(&mut inner, table_idx, true, &ac_huff)?;
2859
2860                        // Write SOS and encode
2861                        write_sos_marker(&mut inner, scan, &components)?;
2862                        bit_writer = BitWriter::new(inner);
2863
2864                        let ac_derived = DerivedTable::from_huff_table(&ac_huff, false)?;
2865                        let mut prog_encoder = ProgressiveEncoder::new(&mut bit_writer);
2866
2867                        self.encode_progressive_scan(
2868                            scan,
2869                            &y_blocks,
2870                            &cb_blocks,
2871                            &cr_blocks,
2872                            mcu_rows,
2873                            mcu_cols,
2874                            luma_h,
2875                            luma_v,
2876                            width,
2877                            height,
2878                            chroma_width,
2879                            chroma_height,
2880                            &opt_dc_luma,
2881                            &opt_dc_chroma,
2882                            &ac_derived,
2883                            &ac_derived, // Not used for AC scans, but needed for signature
2884                            &mut prog_encoder,
2885                        )?;
2886                        prog_encoder.finish_scan(Some(&ac_derived))?;
2887                    } else {
2888                        // DC scan: use global DC tables
2889                        write_sos_marker(&mut inner, scan, &components)?;
2890                        bit_writer = BitWriter::new(inner);
2891
2892                        let mut prog_encoder = ProgressiveEncoder::new(&mut bit_writer);
2893                        self.encode_progressive_scan(
2894                            scan,
2895                            &y_blocks,
2896                            &cb_blocks,
2897                            &cr_blocks,
2898                            mcu_rows,
2899                            mcu_cols,
2900                            luma_h,
2901                            luma_v,
2902                            width,
2903                            height,
2904                            chroma_width,
2905                            chroma_height,
2906                            &opt_dc_luma,
2907                            &opt_dc_chroma,
2908                            &ac_luma_derived, // Not used for DC scans
2909                            &ac_chroma_derived,
2910                            &mut prog_encoder,
2911                        )?;
2912                        prog_encoder.finish_scan(None)?;
2913                    }
2914                }
2915
2916                // Flush and write EOI
2917                bit_writer.flush()?;
2918                let mut output = bit_writer.into_inner();
2919                output.write_all(&[0xFF, 0xD9])?;
2920            } else {
2921                // Standard tables mode (no optimization)
2922                let output = marker_writer.into_inner();
2923                let mut bit_writer = BitWriter::new(output);
2924
2925                for scan in &scans {
2926                    bit_writer.flush()?;
2927                    let mut inner = bit_writer.into_inner();
2928                    write_sos_marker(&mut inner, scan, &components)?;
2929
2930                    bit_writer = BitWriter::new(inner);
2931                    let mut prog_encoder = ProgressiveEncoder::new_standard_tables(&mut bit_writer);
2932
2933                    self.encode_progressive_scan(
2934                        scan,
2935                        &y_blocks,
2936                        &cb_blocks,
2937                        &cr_blocks,
2938                        mcu_rows,
2939                        mcu_cols,
2940                        luma_h,
2941                        luma_v,
2942                        width,
2943                        height,
2944                        chroma_width,
2945                        chroma_height,
2946                        &dc_luma_derived,
2947                        &dc_chroma_derived,
2948                        &ac_luma_derived,
2949                        &ac_chroma_derived,
2950                        &mut prog_encoder,
2951                    )?;
2952
2953                    let ac_table = if scan.ss > 0 {
2954                        if scan.component_index[0] == 0 {
2955                            Some(&ac_luma_derived)
2956                        } else {
2957                            Some(&ac_chroma_derived)
2958                        }
2959                    } else {
2960                        None
2961                    };
2962                    prog_encoder.finish_scan(ac_table)?;
2963                }
2964
2965                bit_writer.flush()?;
2966                let mut output = bit_writer.into_inner();
2967                output.write_all(&[0xFF, 0xD9])?;
2968            }
2969        } else if self.optimize_huffman {
2970            // Baseline mode with Huffman optimization (2-pass)
2971            // Pass 1: Collect blocks and count frequencies
2972            let mcu_rows = mcu_height / (DCTSIZE * luma_v as usize);
2973            let mcu_cols = mcu_width / (DCTSIZE * luma_h as usize);
2974            let num_y_blocks = mcu_rows
2975                .checked_mul(mcu_cols)
2976                .and_then(|n| n.checked_mul(luma_h as usize))
2977                .and_then(|n| n.checked_mul(luma_v as usize))
2978                .ok_or(Error::AllocationFailed)?;
2979            let num_chroma_blocks = mcu_rows
2980                .checked_mul(mcu_cols)
2981                .ok_or(Error::AllocationFailed)?;
2982
2983            let mut y_blocks = try_alloc_vec_array::<i16, DCTSIZE2>(num_y_blocks)?;
2984            let mut cb_blocks = try_alloc_vec_array::<i16, DCTSIZE2>(num_chroma_blocks)?;
2985            let mut cr_blocks = try_alloc_vec_array::<i16, DCTSIZE2>(num_chroma_blocks)?;
2986
2987            // Optionally collect raw DCT for DC trellis
2988            let dc_trellis_enabled = self.trellis.enabled && self.trellis.dc_enabled;
2989            let mut y_raw_dct = if dc_trellis_enabled {
2990                Some(try_alloc_vec_array::<i32, DCTSIZE2>(num_y_blocks)?)
2991            } else {
2992                None
2993            };
2994            let mut cb_raw_dct = if dc_trellis_enabled {
2995                Some(try_alloc_vec_array::<i32, DCTSIZE2>(num_chroma_blocks)?)
2996            } else {
2997                None
2998            };
2999            let mut cr_raw_dct = if dc_trellis_enabled {
3000                Some(try_alloc_vec_array::<i32, DCTSIZE2>(num_chroma_blocks)?)
3001            } else {
3002                None
3003            };
3004
3005            self.collect_blocks(
3006                y_mcu,
3007                mcu_width,
3008                mcu_height,
3009                cb_mcu,
3010                cr_mcu,
3011                mcu_chroma_w,
3012                mcu_chroma_h,
3013                &luma_qtable.values,
3014                &chroma_qtable.values,
3015                &ac_luma_derived,
3016                &ac_chroma_derived,
3017                &mut y_blocks,
3018                &mut cb_blocks,
3019                &mut cr_blocks,
3020                y_raw_dct.as_deref_mut(),
3021                cb_raw_dct.as_deref_mut(),
3022                cr_raw_dct.as_deref_mut(),
3023                luma_h,
3024                luma_v,
3025            )?;
3026
3027            // Run DC trellis optimization if enabled
3028            // C mozjpeg processes DC trellis row by row (each row is an independent chain)
3029            if dc_trellis_enabled {
3030                let h = luma_h as usize;
3031                let v = luma_v as usize;
3032                let y_block_cols = mcu_cols * h;
3033                let y_block_rows = mcu_rows * v;
3034
3035                if let Some(ref y_raw) = y_raw_dct {
3036                    run_dc_trellis_by_row(
3037                        y_raw,
3038                        &mut y_blocks,
3039                        luma_qtable.values[0],
3040                        &dc_luma_derived,
3041                        self.trellis.lambda_log_scale1,
3042                        self.trellis.lambda_log_scale2,
3043                        y_block_rows,
3044                        y_block_cols,
3045                        mcu_cols,
3046                        h,
3047                        v,
3048                        self.trellis.delta_dc_weight,
3049                    );
3050                }
3051                // Chroma has 1x1 per MCU, so MCU order = row order
3052                if let Some(ref cb_raw) = cb_raw_dct {
3053                    run_dc_trellis_by_row(
3054                        cb_raw,
3055                        &mut cb_blocks,
3056                        chroma_qtable.values[0],
3057                        &dc_chroma_derived,
3058                        self.trellis.lambda_log_scale1,
3059                        self.trellis.lambda_log_scale2,
3060                        mcu_rows,
3061                        mcu_cols,
3062                        mcu_cols,
3063                        1,
3064                        1,
3065                        self.trellis.delta_dc_weight,
3066                    );
3067                }
3068                if let Some(ref cr_raw) = cr_raw_dct {
3069                    run_dc_trellis_by_row(
3070                        cr_raw,
3071                        &mut cr_blocks,
3072                        chroma_qtable.values[0],
3073                        &dc_chroma_derived,
3074                        self.trellis.lambda_log_scale1,
3075                        self.trellis.lambda_log_scale2,
3076                        mcu_rows,
3077                        mcu_cols,
3078                        mcu_cols,
3079                        1,
3080                        1,
3081                        self.trellis.delta_dc_weight,
3082                    );
3083                }
3084            }
3085
3086            // Count symbol frequencies
3087            let mut dc_luma_freq = FrequencyCounter::new();
3088            let mut dc_chroma_freq = FrequencyCounter::new();
3089            let mut ac_luma_freq = FrequencyCounter::new();
3090            let mut ac_chroma_freq = FrequencyCounter::new();
3091
3092            let mut counter = SymbolCounter::new();
3093            let blocks_per_mcu_y = (luma_h * luma_v) as usize;
3094            let mut y_idx = 0;
3095            let mut c_idx = 0;
3096
3097            for _mcu_row in 0..mcu_rows {
3098                for _mcu_col in 0..mcu_cols {
3099                    // Y blocks
3100                    for _ in 0..blocks_per_mcu_y {
3101                        counter.count_block(
3102                            &y_blocks[y_idx],
3103                            0,
3104                            &mut dc_luma_freq,
3105                            &mut ac_luma_freq,
3106                        );
3107                        y_idx += 1;
3108                    }
3109                    // Cb block
3110                    counter.count_block(
3111                        &cb_blocks[c_idx],
3112                        1,
3113                        &mut dc_chroma_freq,
3114                        &mut ac_chroma_freq,
3115                    );
3116                    // Cr block
3117                    counter.count_block(
3118                        &cr_blocks[c_idx],
3119                        2,
3120                        &mut dc_chroma_freq,
3121                        &mut ac_chroma_freq,
3122                    );
3123                    c_idx += 1;
3124                }
3125            }
3126
3127            // Generate optimized Huffman tables
3128            let opt_dc_luma_huff = dc_luma_freq.generate_table()?;
3129            let opt_dc_chroma_huff = dc_chroma_freq.generate_table()?;
3130            let opt_ac_luma_huff = ac_luma_freq.generate_table()?;
3131            let opt_ac_chroma_huff = ac_chroma_freq.generate_table()?;
3132
3133            let opt_dc_luma = DerivedTable::from_huff_table(&opt_dc_luma_huff, true)?;
3134            let opt_dc_chroma = DerivedTable::from_huff_table(&opt_dc_chroma_huff, true)?;
3135            let opt_ac_luma = DerivedTable::from_huff_table(&opt_ac_luma_huff, false)?;
3136            let opt_ac_chroma = DerivedTable::from_huff_table(&opt_ac_chroma_huff, false)?;
3137
3138            // Write DHT with optimized tables - combined into single marker
3139            marker_writer.write_dht_multiple(&[
3140                (0, false, &opt_dc_luma_huff),
3141                (1, false, &opt_dc_chroma_huff),
3142                (0, true, &opt_ac_luma_huff),
3143                (1, true, &opt_ac_chroma_huff),
3144            ])?;
3145
3146            // Write SOS and encode
3147            let scans = generate_baseline_scan(3);
3148            let scan = &scans[0];
3149            marker_writer.write_sos(scan, &components)?;
3150
3151            let mut output = marker_writer.into_inner();
3152
3153            // Use SIMD entropy encoder on x86_64 for ~2x faster encoding
3154            #[cfg(target_arch = "x86_64")]
3155            {
3156                let mut simd_entropy = SimdEntropyEncoder::new();
3157
3158                // Encode from stored blocks with restart marker support
3159                y_idx = 0;
3160                c_idx = 0;
3161                let restart_interval = self.restart_interval as usize;
3162                let mut mcu_count = 0usize;
3163                let mut restart_num = 0u8;
3164
3165                for _mcu_row in 0..mcu_rows {
3166                    for _mcu_col in 0..mcu_cols {
3167                        // Emit restart marker if needed (before this MCU, not first)
3168                        if restart_interval > 0
3169                            && mcu_count > 0
3170                            && mcu_count.is_multiple_of(restart_interval)
3171                        {
3172                            simd_entropy.emit_restart(restart_num);
3173                            restart_num = restart_num.wrapping_add(1) & 0x07;
3174                        }
3175
3176                        // Y blocks
3177                        for _ in 0..blocks_per_mcu_y {
3178                            simd_entropy.encode_block(
3179                                &y_blocks[y_idx],
3180                                0,
3181                                &opt_dc_luma,
3182                                &opt_ac_luma,
3183                            );
3184                            y_idx += 1;
3185                        }
3186                        // Cb block
3187                        simd_entropy.encode_block(
3188                            &cb_blocks[c_idx],
3189                            1,
3190                            &opt_dc_chroma,
3191                            &opt_ac_chroma,
3192                        );
3193                        // Cr block
3194                        simd_entropy.encode_block(
3195                            &cr_blocks[c_idx],
3196                            2,
3197                            &opt_dc_chroma,
3198                            &opt_ac_chroma,
3199                        );
3200                        c_idx += 1;
3201                        mcu_count += 1;
3202                    }
3203                }
3204
3205                simd_entropy.flush();
3206                output.write_all(simd_entropy.get_buffer())?;
3207            }
3208
3209            // Fallback for non-x86_64 platforms
3210            #[cfg(not(target_arch = "x86_64"))]
3211            {
3212                let mut bit_writer = BitWriter::new(output);
3213                let mut entropy = EntropyEncoder::new(&mut bit_writer);
3214
3215                // Encode from stored blocks with restart marker support
3216                y_idx = 0;
3217                c_idx = 0;
3218                let restart_interval = self.restart_interval as usize;
3219                let mut mcu_count = 0usize;
3220                let mut restart_num = 0u8;
3221
3222                for _mcu_row in 0..mcu_rows {
3223                    for _mcu_col in 0..mcu_cols {
3224                        // Emit restart marker if needed (before this MCU, not first)
3225                        if restart_interval > 0
3226                            && mcu_count > 0
3227                            && mcu_count.is_multiple_of(restart_interval)
3228                        {
3229                            entropy.emit_restart(restart_num)?;
3230                            restart_num = restart_num.wrapping_add(1) & 0x07;
3231                        }
3232
3233                        // Y blocks
3234                        for _ in 0..blocks_per_mcu_y {
3235                            entropy.encode_block(
3236                                &y_blocks[y_idx],
3237                                0,
3238                                &opt_dc_luma,
3239                                &opt_ac_luma,
3240                            )?;
3241                            y_idx += 1;
3242                        }
3243                        // Cb block
3244                        entropy.encode_block(
3245                            &cb_blocks[c_idx],
3246                            1,
3247                            &opt_dc_chroma,
3248                            &opt_ac_chroma,
3249                        )?;
3250                        // Cr block
3251                        entropy.encode_block(
3252                            &cr_blocks[c_idx],
3253                            2,
3254                            &opt_dc_chroma,
3255                            &opt_ac_chroma,
3256                        )?;
3257                        c_idx += 1;
3258                        mcu_count += 1;
3259                    }
3260                }
3261
3262                bit_writer.flush()?;
3263                output = bit_writer.into_inner();
3264            }
3265
3266            output.write_all(&[0xFF, 0xD9])?;
3267        } else {
3268            // Baseline mode: Encode directly (streaming)
3269            let scans = generate_baseline_scan(3);
3270            let scan = &scans[0]; // Baseline has only one scan
3271            marker_writer.write_sos(scan, &components)?;
3272
3273            // Encode MCU data
3274            let output = marker_writer.into_inner();
3275            let mut bit_writer = BitWriter::new(output);
3276            let mut entropy = EntropyEncoder::new(&mut bit_writer);
3277
3278            self.encode_mcus(
3279                y_mcu,
3280                mcu_width,
3281                mcu_height,
3282                cb_mcu,
3283                cr_mcu,
3284                mcu_chroma_w,
3285                mcu_chroma_h,
3286                &luma_qtable.values,
3287                &chroma_qtable.values,
3288                &dc_luma_derived,
3289                &dc_chroma_derived,
3290                &ac_luma_derived,
3291                &ac_chroma_derived,
3292                &mut entropy,
3293                luma_h,
3294                luma_v,
3295            )?;
3296
3297            // Flush bits and get output back
3298            bit_writer.flush()?;
3299            let mut output = bit_writer.into_inner();
3300
3301            // EOI
3302            output.write_all(&[0xFF, 0xD9])?;
3303        }
3304
3305        Ok(())
3306    }
3307
3308    /// Encode all MCUs (Minimum Coded Units).
3309    #[allow(clippy::too_many_arguments)]
3310    fn encode_mcus<W: Write>(
3311        &self,
3312        y_plane: &[u8],
3313        y_width: usize,
3314        y_height: usize,
3315        cb_plane: &[u8],
3316        cr_plane: &[u8],
3317        chroma_width: usize,
3318        _chroma_height: usize,
3319        luma_qtable: &[u16; DCTSIZE2],
3320        chroma_qtable: &[u16; DCTSIZE2],
3321        dc_luma: &DerivedTable,
3322        dc_chroma: &DerivedTable,
3323        ac_luma: &DerivedTable,
3324        ac_chroma: &DerivedTable,
3325        entropy: &mut EntropyEncoder<W>,
3326        h_samp: u8,
3327        v_samp: u8,
3328    ) -> Result<()> {
3329        let mcu_rows = y_height / (DCTSIZE * v_samp as usize);
3330        let mcu_cols = y_width / (DCTSIZE * h_samp as usize);
3331        let total_mcus = mcu_rows * mcu_cols;
3332
3333        let mut dct_block = [0i16; DCTSIZE2];
3334        let mut quant_block = [0i16; DCTSIZE2];
3335
3336        // Restart marker tracking
3337        let restart_interval = self.restart_interval as usize;
3338        let mut mcu_count = 0usize;
3339        let mut restart_num = 0u8;
3340
3341        for mcu_row in 0..mcu_rows {
3342            for mcu_col in 0..mcu_cols {
3343                // Check if we need to emit a restart marker BEFORE this MCU
3344                // (except for the first MCU)
3345                if restart_interval > 0
3346                    && mcu_count > 0
3347                    && mcu_count.is_multiple_of(restart_interval)
3348                {
3349                    entropy.emit_restart(restart_num)?;
3350                    restart_num = restart_num.wrapping_add(1) & 0x07;
3351                }
3352
3353                // Encode Y blocks (may be multiple per MCU for subsampling)
3354                for v in 0..v_samp as usize {
3355                    for h in 0..h_samp as usize {
3356                        let block_row = mcu_row * v_samp as usize + v;
3357                        let block_col = mcu_col * h_samp as usize + h;
3358
3359                        self.encode_block(
3360                            y_plane,
3361                            y_width,
3362                            block_row,
3363                            block_col,
3364                            luma_qtable,
3365                            dc_luma,
3366                            ac_luma,
3367                            0, // Y component
3368                            entropy,
3369                            &mut dct_block,
3370                            &mut quant_block,
3371                        )?;
3372                    }
3373                }
3374
3375                // Encode Cb block
3376                self.encode_block(
3377                    cb_plane,
3378                    chroma_width,
3379                    mcu_row,
3380                    mcu_col,
3381                    chroma_qtable,
3382                    dc_chroma,
3383                    ac_chroma,
3384                    1, // Cb component
3385                    entropy,
3386                    &mut dct_block,
3387                    &mut quant_block,
3388                )?;
3389
3390                // Encode Cr block
3391                self.encode_block(
3392                    cr_plane,
3393                    chroma_width,
3394                    mcu_row,
3395                    mcu_col,
3396                    chroma_qtable,
3397                    dc_chroma,
3398                    ac_chroma,
3399                    2, // Cr component
3400                    entropy,
3401                    &mut dct_block,
3402                    &mut quant_block,
3403                )?;
3404
3405                mcu_count += 1;
3406            }
3407        }
3408
3409        // Suppress unused variable warning
3410        let _ = total_mcus;
3411
3412        Ok(())
3413    }
3414
3415    /// Encode a single 8x8 block.
3416    #[allow(clippy::too_many_arguments)]
3417    fn encode_block<W: Write>(
3418        &self,
3419        plane: &[u8],
3420        plane_width: usize,
3421        block_row: usize,
3422        block_col: usize,
3423        qtable: &[u16; DCTSIZE2],
3424        dc_table: &DerivedTable,
3425        ac_table: &DerivedTable,
3426        component: usize,
3427        entropy: &mut EntropyEncoder<W>,
3428        dct_block: &mut [i16; DCTSIZE2],
3429        quant_block: &mut [i16; DCTSIZE2],
3430    ) -> Result<()> {
3431        // Extract 8x8 block from plane
3432        let mut samples = [0u8; DCTSIZE2];
3433        let base_y = block_row * DCTSIZE;
3434        let base_x = block_col * DCTSIZE;
3435
3436        for row in 0..DCTSIZE {
3437            let src_offset = (base_y + row) * plane_width + base_x;
3438            let dst_offset = row * DCTSIZE;
3439            samples[dst_offset..dst_offset + DCTSIZE]
3440                .copy_from_slice(&plane[src_offset..src_offset + DCTSIZE]);
3441        }
3442
3443        // Level shift (center around 0 for DCT)
3444        let mut shifted = [0i16; DCTSIZE2];
3445        for i in 0..DCTSIZE2 {
3446            shifted[i] = (samples[i] as i16) - 128;
3447        }
3448
3449        // Apply overshoot deringing if enabled (reduces ringing on white backgrounds)
3450        if self.overshoot_deringing {
3451            preprocess_deringing(&mut shifted, qtable[0]);
3452        }
3453
3454        // Forward DCT (output scaled by factor of 8)
3455        self.simd.do_forward_dct(&shifted, dct_block);
3456
3457        // Convert to i32 for quantization
3458        let mut dct_i32 = [0i32; DCTSIZE2];
3459        for i in 0..DCTSIZE2 {
3460            dct_i32[i] = dct_block[i] as i32;
3461        }
3462
3463        // Use trellis quantization if enabled
3464        // Both paths expect raw DCT (scaled by 8) and handle the scaling internally
3465        if self.trellis.enabled {
3466            trellis_quantize_block(&dct_i32, quant_block, qtable, ac_table, &self.trellis);
3467        } else {
3468            // Non-trellis path: use single-step quantization matching C mozjpeg
3469            // This takes raw DCT (scaled by 8) and uses q_scaled = 8 * qtable[i]
3470            quantize_block_raw(&dct_i32, qtable, quant_block);
3471        }
3472
3473        // Entropy encode
3474        entropy.encode_block(quant_block, component, dc_table, ac_table)?;
3475
3476        Ok(())
3477    }
3478
3479    /// Collect all quantized DCT blocks for progressive encoding.
3480    /// Also collects raw DCT blocks if DC trellis is enabled.
3481    #[allow(clippy::too_many_arguments)]
3482    fn collect_blocks(
3483        &self,
3484        y_plane: &[u8],
3485        y_width: usize,
3486        y_height: usize,
3487        cb_plane: &[u8],
3488        cr_plane: &[u8],
3489        chroma_width: usize,
3490        _chroma_height: usize,
3491        luma_qtable: &[u16; DCTSIZE2],
3492        chroma_qtable: &[u16; DCTSIZE2],
3493        ac_luma: &DerivedTable,
3494        ac_chroma: &DerivedTable,
3495        y_blocks: &mut [[i16; DCTSIZE2]],
3496        cb_blocks: &mut [[i16; DCTSIZE2]],
3497        cr_blocks: &mut [[i16; DCTSIZE2]],
3498        mut y_raw_dct: Option<&mut [[i32; DCTSIZE2]]>,
3499        mut cb_raw_dct: Option<&mut [[i32; DCTSIZE2]]>,
3500        mut cr_raw_dct: Option<&mut [[i32; DCTSIZE2]]>,
3501        h_samp: u8,
3502        v_samp: u8,
3503    ) -> Result<()> {
3504        let mcu_rows = y_height / (DCTSIZE * v_samp as usize);
3505        let mcu_cols = y_width / (DCTSIZE * h_samp as usize);
3506
3507        let mut y_idx = 0;
3508        let mut c_idx = 0;
3509        let mut dct_block = [0i16; DCTSIZE2];
3510
3511        for mcu_row in 0..mcu_rows {
3512            for mcu_col in 0..mcu_cols {
3513                // Collect Y blocks (may be multiple per MCU for subsampling)
3514                for v in 0..v_samp as usize {
3515                    for h in 0..h_samp as usize {
3516                        let block_row = mcu_row * v_samp as usize + v;
3517                        let block_col = mcu_col * h_samp as usize + h;
3518
3519                        // Get mutable reference to raw DCT output if collecting
3520                        let raw_dct_out = y_raw_dct.as_mut().map(|arr| &mut arr[y_idx][..]);
3521                        self.process_block_to_storage_with_raw(
3522                            y_plane,
3523                            y_width,
3524                            block_row,
3525                            block_col,
3526                            luma_qtable,
3527                            ac_luma,
3528                            &mut y_blocks[y_idx],
3529                            &mut dct_block,
3530                            raw_dct_out,
3531                        )?;
3532                        y_idx += 1;
3533                    }
3534                }
3535
3536                // Collect Cb block
3537                let raw_dct_out = cb_raw_dct.as_mut().map(|arr| &mut arr[c_idx][..]);
3538                self.process_block_to_storage_with_raw(
3539                    cb_plane,
3540                    chroma_width,
3541                    mcu_row,
3542                    mcu_col,
3543                    chroma_qtable,
3544                    ac_chroma,
3545                    &mut cb_blocks[c_idx],
3546                    &mut dct_block,
3547                    raw_dct_out,
3548                )?;
3549
3550                // Collect Cr block
3551                let raw_dct_out = cr_raw_dct.as_mut().map(|arr| &mut arr[c_idx][..]);
3552                self.process_block_to_storage_with_raw(
3553                    cr_plane,
3554                    chroma_width,
3555                    mcu_row,
3556                    mcu_col,
3557                    chroma_qtable,
3558                    ac_chroma,
3559                    &mut cr_blocks[c_idx],
3560                    &mut dct_block,
3561                    raw_dct_out,
3562                )?;
3563
3564                c_idx += 1;
3565            }
3566        }
3567
3568        Ok(())
3569    }
3570
3571    /// Process a block: DCT + quantize, storing the result.
3572    /// Optionally stores raw DCT coefficients for DC trellis.
3573    #[allow(clippy::too_many_arguments)]
3574    fn process_block_to_storage_with_raw(
3575        &self,
3576        plane: &[u8],
3577        plane_width: usize,
3578        block_row: usize,
3579        block_col: usize,
3580        qtable: &[u16; DCTSIZE2],
3581        ac_table: &DerivedTable,
3582        out_block: &mut [i16; DCTSIZE2],
3583        dct_block: &mut [i16; DCTSIZE2],
3584        raw_dct_out: Option<&mut [i32]>,
3585    ) -> Result<()> {
3586        // Extract 8x8 block from plane
3587        let mut samples = [0u8; DCTSIZE2];
3588        let base_y = block_row * DCTSIZE;
3589        let base_x = block_col * DCTSIZE;
3590
3591        for row in 0..DCTSIZE {
3592            let src_offset = (base_y + row) * plane_width + base_x;
3593            let dst_offset = row * DCTSIZE;
3594            samples[dst_offset..dst_offset + DCTSIZE]
3595                .copy_from_slice(&plane[src_offset..src_offset + DCTSIZE]);
3596        }
3597
3598        // Level shift (center around 0 for DCT)
3599        let mut shifted = [0i16; DCTSIZE2];
3600        for i in 0..DCTSIZE2 {
3601            shifted[i] = (samples[i] as i16) - 128;
3602        }
3603
3604        // Apply overshoot deringing if enabled (reduces ringing on white backgrounds)
3605        if self.overshoot_deringing {
3606            preprocess_deringing(&mut shifted, qtable[0]);
3607        }
3608
3609        // Forward DCT (output scaled by factor of 8)
3610        self.simd.do_forward_dct(&shifted, dct_block);
3611
3612        // Convert to i32 for quantization
3613        let mut dct_i32 = [0i32; DCTSIZE2];
3614        for i in 0..DCTSIZE2 {
3615            dct_i32[i] = dct_block[i] as i32;
3616        }
3617
3618        // Store raw DCT if requested (for DC trellis)
3619        if let Some(raw_out) = raw_dct_out {
3620            raw_out.copy_from_slice(&dct_i32);
3621        }
3622
3623        // Use trellis quantization if enabled
3624        // Both paths expect raw DCT (scaled by 8) and handle the scaling internally
3625        if self.trellis.enabled {
3626            trellis_quantize_block(&dct_i32, out_block, qtable, ac_table, &self.trellis);
3627        } else {
3628            // Non-trellis path: use single-step quantization matching C mozjpeg
3629            // This takes raw DCT (scaled by 8) and uses q_scaled = 8 * qtable[i]
3630            quantize_block_raw(&dct_i32, qtable, out_block);
3631        }
3632
3633        Ok(())
3634    }
3635
3636    /// Optimize progressive scan configuration (C mozjpeg-compatible).
3637    ///
3638    /// This implements the optimize_scans feature from C mozjpeg:
3639    /// 1. Generate 64 individual candidate scans
3640    /// 2. Trial-encode scans SEQUENTIALLY to get accurate sizes
3641    /// 3. Use ScanSelector to find optimal Al levels and frequency splits
3642    /// 4. Build the final scan script from the selection
3643    ///
3644    /// IMPORTANT: Scans must be encoded sequentially (not independently) because
3645    /// refinement scans (Ah > 0) need context from previous scans to produce
3646    /// correct output sizes.
3647    #[allow(clippy::too_many_arguments)]
3648    fn optimize_progressive_scans(
3649        &self,
3650        num_components: u8,
3651        y_blocks: &[[i16; DCTSIZE2]],
3652        cb_blocks: &[[i16; DCTSIZE2]],
3653        cr_blocks: &[[i16; DCTSIZE2]],
3654        mcu_rows: usize,
3655        mcu_cols: usize,
3656        h_samp: u8,
3657        v_samp: u8,
3658        actual_width: usize,
3659        actual_height: usize,
3660        chroma_width: usize,
3661        chroma_height: usize,
3662        dc_luma: &DerivedTable,
3663        dc_chroma: &DerivedTable,
3664        ac_luma: &DerivedTable,
3665        ac_chroma: &DerivedTable,
3666    ) -> Result<Vec<crate::types::ScanInfo>> {
3667        let config = ScanSearchConfig::default();
3668        let candidate_scans = generate_search_scans(num_components, &config);
3669
3670        // Use ScanTrialEncoder for sequential trial encoding with proper state tracking
3671        let mut trial_encoder = ScanTrialEncoder::new(
3672            y_blocks,
3673            cb_blocks,
3674            cr_blocks,
3675            dc_luma,
3676            dc_chroma,
3677            ac_luma,
3678            ac_chroma,
3679            mcu_rows,
3680            mcu_cols,
3681            h_samp,
3682            v_samp,
3683            actual_width,
3684            actual_height,
3685            chroma_width,
3686            chroma_height,
3687        );
3688
3689        // Trial-encode all scans sequentially to get accurate sizes
3690        let scan_sizes = trial_encoder.encode_all_scans(&candidate_scans)?;
3691
3692        // Use ScanSelector to find the optimal configuration
3693        let selector = ScanSelector::new(num_components, config.clone());
3694        let result = selector.select_best(&scan_sizes);
3695
3696        // Build the final scan script from the selection
3697        Ok(result.build_final_scans(num_components, &config))
3698    }
3699
3700    /// Encode a single progressive scan.
3701    #[allow(clippy::too_many_arguments)]
3702    fn encode_progressive_scan<W: Write>(
3703        &self,
3704        scan: &crate::types::ScanInfo,
3705        y_blocks: &[[i16; DCTSIZE2]],
3706        cb_blocks: &[[i16; DCTSIZE2]],
3707        cr_blocks: &[[i16; DCTSIZE2]],
3708        mcu_rows: usize,
3709        mcu_cols: usize,
3710        h_samp: u8,
3711        v_samp: u8,
3712        actual_width: usize,
3713        actual_height: usize,
3714        chroma_width: usize,
3715        chroma_height: usize,
3716        dc_luma: &DerivedTable,
3717        dc_chroma: &DerivedTable,
3718        ac_luma: &DerivedTable,
3719        ac_chroma: &DerivedTable,
3720        encoder: &mut ProgressiveEncoder<W>,
3721    ) -> Result<()> {
3722        let is_dc_scan = scan.ss == 0 && scan.se == 0;
3723        let is_refinement = scan.ah != 0;
3724
3725        if is_dc_scan {
3726            // DC scan - can be interleaved (multiple components)
3727            self.encode_dc_scan(
3728                scan,
3729                y_blocks,
3730                cb_blocks,
3731                cr_blocks,
3732                mcu_rows,
3733                mcu_cols,
3734                h_samp,
3735                v_samp,
3736                dc_luma,
3737                dc_chroma,
3738                is_refinement,
3739                encoder,
3740            )?;
3741        } else {
3742            // AC scan - single component only (non-interleaved)
3743            // For non-interleaved scans, use actual component block dimensions
3744            let comp_idx = scan.component_index[0] as usize;
3745            let blocks = match comp_idx {
3746                0 => y_blocks,
3747                1 => cb_blocks,
3748                2 => cr_blocks,
3749                _ => return Err(Error::InvalidComponentIndex(comp_idx)),
3750            };
3751            let ac_table = if comp_idx == 0 { ac_luma } else { ac_chroma };
3752
3753            // Calculate actual block dimensions for this component.
3754            // Non-interleaved AC scans encode only the actual image blocks, not MCU padding.
3755            // This differs from interleaved DC scans which encode all MCU blocks.
3756            // Reference: ITU-T T.81 Section F.2.3
3757            let (block_cols, block_rows) = if comp_idx == 0 {
3758                // Y component: full resolution
3759                (
3760                    actual_width.div_ceil(DCTSIZE),
3761                    actual_height.div_ceil(DCTSIZE),
3762                )
3763            } else {
3764                // Chroma components: subsampled resolution
3765                (
3766                    chroma_width.div_ceil(DCTSIZE),
3767                    chroma_height.div_ceil(DCTSIZE),
3768                )
3769            };
3770
3771            self.encode_ac_scan(
3772                scan,
3773                blocks,
3774                mcu_rows,
3775                mcu_cols,
3776                h_samp,
3777                v_samp,
3778                comp_idx,
3779                block_cols,
3780                block_rows,
3781                ac_table,
3782                is_refinement,
3783                encoder,
3784            )?;
3785        }
3786
3787        Ok(())
3788    }
3789
3790    /// Encode a DC scan (Ss=Se=0).
3791    #[allow(clippy::too_many_arguments)]
3792    fn encode_dc_scan<W: Write>(
3793        &self,
3794        scan: &crate::types::ScanInfo,
3795        y_blocks: &[[i16; DCTSIZE2]],
3796        cb_blocks: &[[i16; DCTSIZE2]],
3797        cr_blocks: &[[i16; DCTSIZE2]],
3798        mcu_rows: usize,
3799        mcu_cols: usize,
3800        h_samp: u8,
3801        v_samp: u8,
3802        dc_luma: &DerivedTable,
3803        dc_chroma: &DerivedTable,
3804        is_refinement: bool,
3805        encoder: &mut ProgressiveEncoder<W>,
3806    ) -> Result<()> {
3807        let blocks_per_mcu_y = (h_samp * v_samp) as usize;
3808        let mut y_idx = 0;
3809        let mut c_idx = 0;
3810
3811        for _mcu_row in 0..mcu_rows {
3812            for _mcu_col in 0..mcu_cols {
3813                // Encode Y blocks
3814                for _ in 0..blocks_per_mcu_y {
3815                    if is_refinement {
3816                        encoder.encode_dc_refine(&y_blocks[y_idx], scan.al)?;
3817                    } else {
3818                        encoder.encode_dc_first(&y_blocks[y_idx], 0, dc_luma, scan.al)?;
3819                    }
3820                    y_idx += 1;
3821                }
3822
3823                // Encode Cb
3824                if is_refinement {
3825                    encoder.encode_dc_refine(&cb_blocks[c_idx], scan.al)?;
3826                } else {
3827                    encoder.encode_dc_first(&cb_blocks[c_idx], 1, dc_chroma, scan.al)?;
3828                }
3829
3830                // Encode Cr
3831                if is_refinement {
3832                    encoder.encode_dc_refine(&cr_blocks[c_idx], scan.al)?;
3833                } else {
3834                    encoder.encode_dc_first(&cr_blocks[c_idx], 2, dc_chroma, scan.al)?;
3835                }
3836
3837                c_idx += 1;
3838            }
3839        }
3840
3841        Ok(())
3842    }
3843
3844    /// Encode an AC scan (Ss > 0).
3845    ///
3846    /// **IMPORTANT**: Progressive AC scans are always non-interleaved, meaning blocks
3847    /// must be encoded in component raster order (row-major within the component's
3848    /// block grid), NOT in MCU-interleaved order.
3849    ///
3850    /// For non-interleaved scans, the number of blocks is determined by the actual
3851    /// component dimensions (ceil(width/8) × ceil(height/8)), NOT the MCU-padded
3852    /// dimensions. This is different from interleaved DC scans which use MCU order.
3853    /// The padding blocks (beyond actual image dimensions) have DC coefficients but
3854    /// no AC coefficients - the decoder only outputs the actual image dimensions.
3855    ///
3856    /// Reference: ITU-T T.81 Section F.2.3 - "The scan data for a non-interleaved
3857    /// scan shall consist of a sequence of entropy-coded segments... The data units
3858    /// are processed in the order defined by the scan component."
3859    #[allow(clippy::too_many_arguments)]
3860    fn encode_ac_scan<W: Write>(
3861        &self,
3862        scan: &crate::types::ScanInfo,
3863        blocks: &[[i16; DCTSIZE2]],
3864        _mcu_rows: usize,
3865        mcu_cols: usize,
3866        h_samp: u8,
3867        v_samp: u8,
3868        comp_idx: usize,
3869        block_cols: usize,
3870        block_rows: usize,
3871        ac_table: &DerivedTable,
3872        is_refinement: bool,
3873        encoder: &mut ProgressiveEncoder<W>,
3874    ) -> Result<()> {
3875        // For Y component with subsampling, blocks are stored in MCU-interleaved order
3876        // but AC scans must encode them in component raster order.
3877        // For chroma components (1 block per MCU), the orders are identical.
3878        //
3879        // For non-interleaved scans, encode only the actual image blocks (block_rows × block_cols),
3880        // not all MCU-padded blocks. Padding blocks have DC coefficients but no AC coefficients.
3881
3882        let blocks_per_mcu = if comp_idx == 0 {
3883            (h_samp * v_samp) as usize
3884        } else {
3885            1
3886        };
3887
3888        if blocks_per_mcu == 1 {
3889            // Chroma or 4:4:4 Y: storage order = raster order
3890            let total_blocks = block_rows * block_cols;
3891            for block in blocks.iter().take(total_blocks) {
3892                if is_refinement {
3893                    encoder
3894                        .encode_ac_refine(block, scan.ss, scan.se, scan.ah, scan.al, ac_table)?;
3895                } else {
3896                    encoder.encode_ac_first(block, scan.ss, scan.se, scan.al, ac_table)?;
3897                }
3898            }
3899        } else {
3900            // Y component with subsampling (h_samp > 1 or v_samp > 1)
3901            // Convert from MCU-interleaved storage to component raster order
3902            let h = h_samp as usize;
3903            let v = v_samp as usize;
3904
3905            for block_row in 0..block_rows {
3906                for block_col in 0..block_cols {
3907                    // Convert raster position to MCU-interleaved storage index
3908                    let mcu_row = block_row / v;
3909                    let mcu_col = block_col / h;
3910                    let v_idx = block_row % v;
3911                    let h_idx = block_col % h;
3912                    let storage_idx = mcu_row * (mcu_cols * blocks_per_mcu)
3913                        + mcu_col * blocks_per_mcu
3914                        + v_idx * h
3915                        + h_idx;
3916
3917                    if is_refinement {
3918                        encoder.encode_ac_refine(
3919                            &blocks[storage_idx],
3920                            scan.ss,
3921                            scan.se,
3922                            scan.ah,
3923                            scan.al,
3924                            ac_table,
3925                        )?;
3926                    } else {
3927                        encoder.encode_ac_first(
3928                            &blocks[storage_idx],
3929                            scan.ss,
3930                            scan.se,
3931                            scan.al,
3932                            ac_table,
3933                        )?;
3934                    }
3935                }
3936            }
3937        }
3938
3939        Ok(())
3940    }
3941
3942    /// Count DC symbols for a progressive DC scan.
3943    #[allow(clippy::too_many_arguments)]
3944    fn count_dc_scan_symbols(
3945        &self,
3946        scan: &crate::types::ScanInfo,
3947        y_blocks: &[[i16; DCTSIZE2]],
3948        cb_blocks: &[[i16; DCTSIZE2]],
3949        cr_blocks: &[[i16; DCTSIZE2]],
3950        mcu_rows: usize,
3951        mcu_cols: usize,
3952        h_samp: u8,
3953        v_samp: u8,
3954        dc_luma_freq: &mut FrequencyCounter,
3955        dc_chroma_freq: &mut FrequencyCounter,
3956    ) {
3957        let blocks_per_mcu_y = (h_samp * v_samp) as usize;
3958        let mut y_idx = 0;
3959        let mut c_idx = 0;
3960        let mut counter = ProgressiveSymbolCounter::new();
3961
3962        for _mcu_row in 0..mcu_rows {
3963            for _mcu_col in 0..mcu_cols {
3964                // Y blocks
3965                for _ in 0..blocks_per_mcu_y {
3966                    counter.count_dc_first(&y_blocks[y_idx], 0, scan.al, dc_luma_freq);
3967                    y_idx += 1;
3968                }
3969                // Cb block
3970                counter.count_dc_first(&cb_blocks[c_idx], 1, scan.al, dc_chroma_freq);
3971                // Cr block
3972                counter.count_dc_first(&cr_blocks[c_idx], 2, scan.al, dc_chroma_freq);
3973                c_idx += 1;
3974            }
3975        }
3976    }
3977
3978    /// Count AC symbols for a progressive AC scan.
3979    ///
3980    /// Must iterate blocks in the same order as `encode_ac_scan` (component raster order)
3981    /// to ensure EOBRUN counts match and Huffman tables are correct.
3982    ///
3983    /// Uses actual block dimensions (not MCU-padded) for non-interleaved scans.
3984    #[allow(clippy::too_many_arguments)]
3985    fn count_ac_scan_symbols(
3986        &self,
3987        scan: &crate::types::ScanInfo,
3988        blocks: &[[i16; DCTSIZE2]],
3989        _mcu_rows: usize,
3990        mcu_cols: usize,
3991        h_samp: u8,
3992        v_samp: u8,
3993        comp_idx: usize,
3994        block_cols: usize,
3995        block_rows: usize,
3996        ac_freq: &mut FrequencyCounter,
3997    ) {
3998        let blocks_per_mcu = if comp_idx == 0 {
3999            (h_samp * v_samp) as usize
4000        } else {
4001            1
4002        };
4003
4004        let mut counter = ProgressiveSymbolCounter::new();
4005        let is_refinement = scan.ah != 0;
4006
4007        if blocks_per_mcu == 1 {
4008            // Chroma or 4:4:4 Y: storage order = raster order
4009            let total_blocks = block_rows * block_cols;
4010            for block in blocks.iter().take(total_blocks) {
4011                if is_refinement {
4012                    counter.count_ac_refine(block, scan.ss, scan.se, scan.ah, scan.al, ac_freq);
4013                } else {
4014                    counter.count_ac_first(block, scan.ss, scan.se, scan.al, ac_freq);
4015                }
4016            }
4017        } else {
4018            // Y component with subsampling - iterate in raster order (matching encode_ac_scan)
4019            let h = h_samp as usize;
4020            let v = v_samp as usize;
4021
4022            for block_row in 0..block_rows {
4023                for block_col in 0..block_cols {
4024                    // Convert raster position to MCU-interleaved storage index
4025                    let mcu_row = block_row / v;
4026                    let mcu_col = block_col / h;
4027                    let v_idx = block_row % v;
4028                    let h_idx = block_col % h;
4029                    let storage_idx = mcu_row * (mcu_cols * blocks_per_mcu)
4030                        + mcu_col * blocks_per_mcu
4031                        + v_idx * h
4032                        + h_idx;
4033
4034                    if is_refinement {
4035                        counter.count_ac_refine(
4036                            &blocks[storage_idx],
4037                            scan.ss,
4038                            scan.se,
4039                            scan.ah,
4040                            scan.al,
4041                            ac_freq,
4042                        );
4043                    } else {
4044                        counter.count_ac_first(
4045                            &blocks[storage_idx],
4046                            scan.ss,
4047                            scan.se,
4048                            scan.al,
4049                            ac_freq,
4050                        );
4051                    }
4052                }
4053            }
4054        }
4055
4056        // Flush any pending EOBRUN
4057        counter.finish_scan(Some(ac_freq));
4058    }
4059}
4060
4061// ============================================================================
4062// Encode Trait Implementation
4063// ============================================================================
4064
4065impl Encode for Encoder {
4066    fn encode_rgb(&self, rgb_data: &[u8], width: u32, height: u32) -> Result<Vec<u8>> {
4067        self.encode_rgb(rgb_data, width, height)
4068    }
4069
4070    fn encode_gray(&self, gray_data: &[u8], width: u32, height: u32) -> Result<Vec<u8>> {
4071        self.encode_gray(gray_data, width, height)
4072    }
4073}
4074
4075// Note: StreamingEncoder and EncodingStream are in the `streaming` module.
4076
4077// Add streaming() method to Encoder
4078impl Encoder {
4079    /// Create a streaming encoder.
4080    ///
4081    /// Returns a [`StreamingEncoder`] which supports scanline-by-scanline encoding.
4082    /// Note that streaming mode does NOT support trellis quantization, progressive
4083    /// mode, or Huffman optimization (these require buffering the entire image).
4084    ///
4085    /// For full-featured encoding with all mozjpeg optimizations, use [`Encoder::new()`]
4086    /// with [`encode_rgb()`](Encoder::encode_rgb) or [`encode_gray()`](Encoder::encode_gray).
4087    ///
4088    /// # Example
4089    ///
4090    /// ```ignore
4091    /// use mozjpeg_rs::Encoder;
4092    /// use std::fs::File;
4093    ///
4094    /// let file = File::create("output.jpg")?;
4095    /// let mut stream = Encoder::streaming()
4096    ///     .quality(85)
4097    ///     .start_rgb(1920, 1080, file)?;
4098    ///
4099    /// // Write scanlines...
4100    /// stream.finish()?;
4101    /// ```
4102    pub fn streaming() -> StreamingEncoder {
4103        StreamingEncoder::baseline_fastest()
4104    }
4105}
4106
4107// ============================================================================
4108// C mozjpeg encoding (optional feature)
4109// ============================================================================
4110
4111#[cfg(feature = "mozjpeg-sys-config")]
4112impl Encoder {
4113    /// Convert this encoder to a C mozjpeg encoder.
4114    ///
4115    /// Returns a [`CMozjpeg`](crate::CMozjpeg) that can encode images using
4116    /// the C mozjpeg library with settings matching this Rust encoder.
4117    ///
4118    /// # Example
4119    ///
4120    /// ```no_run
4121    /// use mozjpeg_rs::{Encoder, Preset};
4122    ///
4123    /// let pixels: Vec<u8> = vec![128; 64 * 64 * 3];
4124    /// let encoder = Encoder::new(Preset::ProgressiveBalanced).quality(85);
4125    ///
4126    /// // Encode with C mozjpeg
4127    /// let c_jpeg = encoder.to_c_mozjpeg().encode_rgb(&pixels, 64, 64)?;
4128    ///
4129    /// // Compare with Rust encoder
4130    /// let rust_jpeg = encoder.encode_rgb(&pixels, 64, 64)?;
4131    /// # Ok::<(), mozjpeg_rs::Error>(())
4132    /// ```
4133    pub fn to_c_mozjpeg(&self) -> crate::compat::CMozjpeg {
4134        crate::compat::CMozjpeg {
4135            quality: self.quality,
4136            force_baseline: self.force_baseline,
4137            subsampling: self.subsampling,
4138            progressive: self.progressive,
4139            optimize_huffman: self.optimize_huffman,
4140            optimize_scans: self.optimize_scans,
4141            trellis: self.trellis,
4142            overshoot_deringing: self.overshoot_deringing,
4143            smoothing: self.smoothing,
4144            restart_interval: self.restart_interval,
4145            quant_table_idx: self.quant_table_idx,
4146            has_custom_qtables: self.custom_luma_qtable.is_some()
4147                || self.custom_chroma_qtable.is_some(),
4148            exif_data: self.exif_data.clone(),
4149            icc_profile: self.icc_profile.clone(),
4150            custom_markers: self.custom_markers.clone(),
4151        }
4152    }
4153}
4154
4155/// Unit tests for private encoder internals.
4156/// Public API tests are in tests/encode_tests.rs.
4157#[cfg(test)]
4158mod tests {
4159    use super::*;
4160
4161    #[test]
4162    fn test_encoder_defaults() {
4163        // Default preset is ProgressiveBalanced
4164        let enc = Encoder::default();
4165        assert_eq!(enc.quality, 75);
4166        assert!(enc.progressive); // ProgressiveBalanced is progressive
4167        assert_eq!(enc.subsampling, Subsampling::S420);
4168        assert!(enc.trellis.enabled);
4169        assert!(enc.optimize_huffman);
4170        assert!(!enc.optimize_scans); // ProgressiveBalanced does NOT include optimize_scans
4171    }
4172
4173    #[test]
4174    fn test_encoder_presets() {
4175        let fastest = Encoder::new(Preset::BaselineFastest);
4176        assert!(!fastest.progressive);
4177        assert!(!fastest.trellis.enabled);
4178        assert!(!fastest.optimize_huffman);
4179
4180        let baseline = Encoder::new(Preset::BaselineBalanced);
4181        assert!(!baseline.progressive);
4182        assert!(baseline.trellis.enabled);
4183        assert!(baseline.optimize_huffman);
4184
4185        let prog_balanced = Encoder::new(Preset::ProgressiveBalanced);
4186        assert!(prog_balanced.progressive);
4187        assert!(prog_balanced.trellis.enabled);
4188        assert!(!prog_balanced.optimize_scans);
4189
4190        let prog_smallest = Encoder::new(Preset::ProgressiveSmallest);
4191        assert!(prog_smallest.progressive);
4192        assert!(prog_smallest.optimize_scans);
4193    }
4194
4195    #[test]
4196    fn test_encoder_builder_fields() {
4197        let enc = Encoder::baseline_optimized()
4198            .quality(90)
4199            .progressive(true)
4200            .subsampling(Subsampling::S444);
4201
4202        assert_eq!(enc.quality, 90);
4203        assert!(enc.progressive);
4204        assert_eq!(enc.subsampling, Subsampling::S444);
4205    }
4206
4207    #[test]
4208    fn test_quality_clamping() {
4209        let enc = Encoder::baseline_optimized().quality(0);
4210        assert_eq!(enc.quality, 1);
4211
4212        let enc = Encoder::baseline_optimized().quality(150);
4213        assert_eq!(enc.quality, 100);
4214    }
4215
4216    #[test]
4217    fn test_natural_to_zigzag() {
4218        let mut natural = [0u16; 64];
4219        for i in 0..64 {
4220            natural[i] = i as u16;
4221        }
4222        let zigzag = natural_to_zigzag(&natural);
4223
4224        assert_eq!(zigzag[0], 0);
4225        assert_eq!(zigzag[1], 1);
4226    }
4227
4228    #[test]
4229    fn test_max_compression_uses_all_optimizations() {
4230        let encoder = Encoder::max_compression();
4231        assert!(encoder.trellis.enabled);
4232        assert!(encoder.progressive);
4233        assert!(encoder.optimize_huffman);
4234        assert!(encoder.optimize_scans);
4235    }
4236
4237    #[test]
4238    fn test_encode_ycbcr_planar_444() {
4239        let width = 32u32;
4240        let height = 32u32;
4241
4242        // Create test image with gradient pattern
4243        let y_plane: Vec<u8> = (0..width * height)
4244            .map(|i| ((i % width) * 255 / width) as u8)
4245            .collect();
4246        let cb_plane: Vec<u8> = (0..width * height)
4247            .map(|i| ((i / width) * 255 / height) as u8)
4248            .collect();
4249        let cr_plane: Vec<u8> = vec![128u8; (width * height) as usize];
4250
4251        let encoder = Encoder::new(Preset::BaselineBalanced)
4252            .quality(85)
4253            .subsampling(Subsampling::S444);
4254
4255        let jpeg_data = encoder
4256            .encode_ycbcr_planar(&y_plane, &cb_plane, &cr_plane, width, height)
4257            .expect("encode_ycbcr_planar should succeed");
4258
4259        // Verify it's a valid JPEG
4260        assert!(jpeg_data.starts_with(&[0xFF, 0xD8, 0xFF])); // SOI + marker
4261        assert!(jpeg_data.ends_with(&[0xFF, 0xD9])); // EOI
4262        assert!(jpeg_data.len() > 200); // Reasonable size for 32x32
4263    }
4264
4265    #[test]
4266    fn test_encode_ycbcr_planar_420() {
4267        let width = 32u32;
4268        let height = 32u32;
4269
4270        // For 4:2:0, chroma planes are half resolution in each dimension
4271        let chroma_w = (width + 1) / 2;
4272        let chroma_h = (height + 1) / 2;
4273
4274        let y_plane: Vec<u8> = vec![128u8; (width * height) as usize];
4275        let cb_plane: Vec<u8> = vec![100u8; (chroma_w * chroma_h) as usize];
4276        let cr_plane: Vec<u8> = vec![150u8; (chroma_w * chroma_h) as usize];
4277
4278        let encoder = Encoder::new(Preset::BaselineBalanced)
4279            .quality(85)
4280            .subsampling(Subsampling::S420);
4281
4282        let jpeg_data = encoder
4283            .encode_ycbcr_planar(&y_plane, &cb_plane, &cr_plane, width, height)
4284            .expect("encode_ycbcr_planar with 4:2:0 should succeed");
4285
4286        // Verify it's a valid JPEG
4287        assert!(jpeg_data.starts_with(&[0xFF, 0xD8, 0xFF]));
4288        assert!(jpeg_data.ends_with(&[0xFF, 0xD9]));
4289    }
4290
4291    #[test]
4292    fn test_encode_ycbcr_planar_422() {
4293        let width = 32u32;
4294        let height = 32u32;
4295
4296        // For 4:2:2, chroma is half width, full height
4297        let chroma_w = (width + 1) / 2;
4298
4299        let y_plane: Vec<u8> = vec![128u8; (width * height) as usize];
4300        let cb_plane: Vec<u8> = vec![100u8; (chroma_w * height) as usize];
4301        let cr_plane: Vec<u8> = vec![150u8; (chroma_w * height) as usize];
4302
4303        let encoder = Encoder::new(Preset::BaselineBalanced)
4304            .quality(85)
4305            .subsampling(Subsampling::S422);
4306
4307        let jpeg_data = encoder
4308            .encode_ycbcr_planar(&y_plane, &cb_plane, &cr_plane, width, height)
4309            .expect("encode_ycbcr_planar with 4:2:2 should succeed");
4310
4311        assert!(jpeg_data.starts_with(&[0xFF, 0xD8, 0xFF]));
4312        assert!(jpeg_data.ends_with(&[0xFF, 0xD9]));
4313    }
4314
4315    #[test]
4316    fn test_encode_ycbcr_planar_wrong_size() {
4317        let width = 32u32;
4318        let height = 32u32;
4319
4320        // Correct Y plane but wrong chroma plane sizes for 4:2:0
4321        let y_plane: Vec<u8> = vec![128u8; (width * height) as usize];
4322        let cb_plane: Vec<u8> = vec![100u8; 10]; // Too small!
4323        let cr_plane: Vec<u8> = vec![150u8; 10]; // Too small!
4324
4325        let encoder = Encoder::new(Preset::BaselineBalanced)
4326            .quality(85)
4327            .subsampling(Subsampling::S420);
4328
4329        let result = encoder.encode_ycbcr_planar(&y_plane, &cb_plane, &cr_plane, width, height);
4330
4331        assert!(result.is_err());
4332    }
4333
4334    #[test]
4335    fn test_encode_ycbcr_planar_strided() {
4336        let width = 30u32; // Not a multiple of stride
4337        let height = 20u32;
4338        let y_stride = 32usize; // Stride with 2 bytes padding per row
4339
4340        // For 4:2:0, chroma is half resolution
4341        let chroma_width = 15usize;
4342        let chroma_height = 10usize;
4343        let cb_stride = 16usize; // Stride with 1 byte padding per row
4344
4345        // Create Y plane with stride (fill with gradient, padding with zeros)
4346        let mut y_plane = vec![0u8; y_stride * height as usize];
4347        for row in 0..height as usize {
4348            for col in 0..width as usize {
4349                y_plane[row * y_stride + col] = ((col * 255) / width as usize) as u8;
4350            }
4351        }
4352
4353        // Create chroma planes with stride
4354        let mut cb_plane = vec![0u8; cb_stride * chroma_height];
4355        let mut cr_plane = vec![0u8; cb_stride * chroma_height];
4356        for row in 0..chroma_height {
4357            for col in 0..chroma_width {
4358                cb_plane[row * cb_stride + col] = 100;
4359                cr_plane[row * cb_stride + col] = 150;
4360            }
4361        }
4362
4363        let encoder = Encoder::new(Preset::BaselineBalanced)
4364            .quality(85)
4365            .subsampling(Subsampling::S420);
4366
4367        let jpeg_data = encoder
4368            .encode_ycbcr_planar_strided(
4369                &y_plane, y_stride, &cb_plane, cb_stride, &cr_plane, cb_stride, width, height,
4370            )
4371            .expect("strided encoding should succeed");
4372
4373        // Verify it's a valid JPEG
4374        assert!(jpeg_data.starts_with(&[0xFF, 0xD8, 0xFF]));
4375        assert!(jpeg_data.ends_with(&[0xFF, 0xD9]));
4376    }
4377
4378    #[test]
4379    fn test_encode_ycbcr_planar_strided_matches_packed() {
4380        let width = 32u32;
4381        let height = 32u32;
4382
4383        // Create packed plane data
4384        let y_packed: Vec<u8> = (0..width * height).map(|i| (i % 256) as u8).collect();
4385        let chroma_w = (width + 1) / 2;
4386        let chroma_h = (height + 1) / 2;
4387        let cb_packed: Vec<u8> = vec![100u8; (chroma_w * chroma_h) as usize];
4388        let cr_packed: Vec<u8> = vec![150u8; (chroma_w * chroma_h) as usize];
4389
4390        let encoder = Encoder::new(Preset::BaselineBalanced)
4391            .quality(85)
4392            .subsampling(Subsampling::S420);
4393
4394        // Encode with packed API
4395        let jpeg_packed = encoder
4396            .encode_ycbcr_planar(&y_packed, &cb_packed, &cr_packed, width, height)
4397            .expect("packed encoding should succeed");
4398
4399        // Encode with strided API (stride == width means packed)
4400        let jpeg_strided = encoder
4401            .encode_ycbcr_planar_strided(
4402                &y_packed,
4403                width as usize,
4404                &cb_packed,
4405                chroma_w as usize,
4406                &cr_packed,
4407                chroma_w as usize,
4408                width,
4409                height,
4410            )
4411            .expect("strided encoding should succeed");
4412
4413        // Both should produce identical output
4414        assert_eq!(jpeg_packed, jpeg_strided);
4415    }
4416
4417    // =========================================================================
4418    // Resource Estimation Tests
4419    // =========================================================================
4420
4421    #[test]
4422    fn test_estimate_resources_basic() {
4423        let encoder = Encoder::new(Preset::BaselineBalanced);
4424        let estimate = encoder.estimate_resources(1920, 1080);
4425
4426        // Should have reasonable memory estimate (> input size)
4427        let input_size = 1920 * 1080 * 3;
4428        assert!(
4429            estimate.peak_memory_bytes > input_size,
4430            "Peak memory {} should exceed input size {}",
4431            estimate.peak_memory_bytes,
4432            input_size
4433        );
4434
4435        // Should have reasonable CPU cost (> 1.0 due to trellis)
4436        assert!(
4437            estimate.cpu_cost_multiplier > 1.0,
4438            "CPU cost {} should be > 1.0 for BaselineBalanced",
4439            estimate.cpu_cost_multiplier
4440        );
4441
4442        // Block count should match expected
4443        assert!(estimate.block_count > 0, "Block count should be > 0");
4444    }
4445
4446    #[test]
4447    fn test_estimate_resources_fastest_has_lower_cpu() {
4448        let fastest = Encoder::new(Preset::BaselineFastest);
4449        let balanced = Encoder::new(Preset::BaselineBalanced);
4450
4451        let est_fast = fastest.estimate_resources(512, 512);
4452        let est_balanced = balanced.estimate_resources(512, 512);
4453
4454        // Fastest should have lower CPU cost (no trellis)
4455        assert!(
4456            est_fast.cpu_cost_multiplier < est_balanced.cpu_cost_multiplier,
4457            "Fastest ({:.2}) should have lower CPU cost than Balanced ({:.2})",
4458            est_fast.cpu_cost_multiplier,
4459            est_balanced.cpu_cost_multiplier
4460        );
4461    }
4462
4463    #[test]
4464    fn test_estimate_resources_progressive_has_higher_cpu() {
4465        let baseline = Encoder::new(Preset::BaselineBalanced);
4466        let progressive = Encoder::new(Preset::ProgressiveBalanced);
4467
4468        let est_baseline = baseline.estimate_resources(512, 512);
4469        let est_prog = progressive.estimate_resources(512, 512);
4470
4471        // Progressive should have higher CPU cost (multiple scans)
4472        assert!(
4473            est_prog.cpu_cost_multiplier > est_baseline.cpu_cost_multiplier,
4474            "Progressive ({:.2}) should have higher CPU cost than Baseline ({:.2})",
4475            est_prog.cpu_cost_multiplier,
4476            est_baseline.cpu_cost_multiplier
4477        );
4478    }
4479
4480    #[test]
4481    fn test_estimate_resources_gray() {
4482        let encoder = Encoder::new(Preset::BaselineBalanced);
4483        let rgb_estimate = encoder.estimate_resources(512, 512);
4484        let gray_estimate = encoder.estimate_resources_gray(512, 512);
4485
4486        // Grayscale should use less memory (1 channel vs 3)
4487        assert!(
4488            gray_estimate.peak_memory_bytes < rgb_estimate.peak_memory_bytes,
4489            "Grayscale memory {} should be less than RGB {}",
4490            gray_estimate.peak_memory_bytes,
4491            rgb_estimate.peak_memory_bytes
4492        );
4493
4494        // Grayscale should have lower CPU cost
4495        assert!(
4496            gray_estimate.cpu_cost_multiplier < rgb_estimate.cpu_cost_multiplier,
4497            "Grayscale CPU {:.2} should be less than RGB {:.2}",
4498            gray_estimate.cpu_cost_multiplier,
4499            rgb_estimate.cpu_cost_multiplier
4500        );
4501    }
4502
4503    // =========================================================================
4504    // Resource Limit Tests
4505    // =========================================================================
4506
4507    #[test]
4508    fn test_dimension_limit_width() {
4509        let limits = Limits::default().max_width(100).max_height(100);
4510        let encoder = Encoder::new(Preset::BaselineFastest).limits(limits);
4511
4512        let pixels = vec![128u8; 200 * 50 * 3];
4513        let result = encoder.encode_rgb(&pixels, 200, 50);
4514
4515        assert!(matches!(result, Err(Error::DimensionLimitExceeded { .. })));
4516    }
4517
4518    #[test]
4519    fn test_dimension_limit_height() {
4520        let limits = Limits::default().max_width(100).max_height(100);
4521        let encoder = Encoder::new(Preset::BaselineFastest).limits(limits);
4522
4523        let pixels = vec![128u8; 50 * 200 * 3];
4524        let result = encoder.encode_rgb(&pixels, 50, 200);
4525
4526        assert!(matches!(result, Err(Error::DimensionLimitExceeded { .. })));
4527    }
4528
4529    #[test]
4530    fn test_dimension_limit_passes_when_within() {
4531        let limits = Limits::default().max_width(100).max_height(100);
4532        let encoder = Encoder::new(Preset::BaselineFastest).limits(limits);
4533
4534        let pixels = vec![128u8; 64 * 64 * 3];
4535        let result = encoder.encode_rgb(&pixels, 64, 64);
4536
4537        assert!(result.is_ok());
4538    }
4539
4540    #[test]
4541    fn test_allocation_limit() {
4542        let limits = Limits::default().max_alloc_bytes(1000); // Very small limit
4543        let encoder = Encoder::new(Preset::BaselineFastest).limits(limits);
4544
4545        let pixels = vec![128u8; 256 * 256 * 3];
4546        let result = encoder.encode_rgb(&pixels, 256, 256);
4547
4548        assert!(matches!(result, Err(Error::AllocationLimitExceeded { .. })));
4549    }
4550
4551    #[test]
4552    fn test_allocation_limit_passes_when_within() {
4553        let limits = Limits::default().max_alloc_bytes(10_000_000); // 10 MB limit
4554        let encoder = Encoder::new(Preset::BaselineFastest).limits(limits);
4555
4556        let pixels = vec![128u8; 64 * 64 * 3];
4557        let result = encoder.encode_rgb(&pixels, 64, 64);
4558
4559        assert!(result.is_ok());
4560    }
4561
4562    #[test]
4563    fn test_pixel_count_limit() {
4564        let limits = Limits::default().max_pixel_count(1000); // Very small limit
4565        let encoder = Encoder::new(Preset::BaselineFastest).limits(limits);
4566
4567        let pixels = vec![128u8; 64 * 64 * 3]; // 4096 pixels
4568        let result = encoder.encode_rgb(&pixels, 64, 64);
4569
4570        assert!(matches!(result, Err(Error::PixelCountExceeded { .. })));
4571    }
4572
4573    #[test]
4574    fn test_pixel_count_limit_passes_when_within() {
4575        let limits = Limits::default().max_pixel_count(10000); // 10000 pixels
4576        let encoder = Encoder::new(Preset::BaselineFastest).limits(limits);
4577
4578        let pixels = vec![128u8; 64 * 64 * 3]; // 4096 pixels
4579        let result = encoder.encode_rgb(&pixels, 64, 64);
4580
4581        assert!(result.is_ok());
4582    }
4583
4584    #[test]
4585    fn test_icc_profile_size_limit() {
4586        let limits = Limits::default().max_icc_profile_bytes(100);
4587        let encoder = Encoder::new(Preset::BaselineFastest)
4588            .limits(limits)
4589            .icc_profile(vec![0u8; 1000]); // 1000 byte ICC profile
4590
4591        let pixels = vec![128u8; 64 * 64 * 3];
4592        let result = encoder.encode_rgb(&pixels, 64, 64);
4593
4594        assert!(matches!(result, Err(Error::IccProfileTooLarge { .. })));
4595    }
4596
4597    #[test]
4598    fn test_icc_profile_size_limit_passes_when_within() {
4599        let limits = Limits::default().max_icc_profile_bytes(2000);
4600        let encoder = Encoder::new(Preset::BaselineFastest)
4601            .limits(limits)
4602            .icc_profile(vec![0u8; 1000]); // 1000 byte ICC profile
4603
4604        let pixels = vec![128u8; 64 * 64 * 3];
4605        let result = encoder.encode_rgb(&pixels, 64, 64);
4606
4607        assert!(result.is_ok());
4608    }
4609
4610    #[test]
4611    fn test_limits_disabled_by_default() {
4612        let encoder = Encoder::new(Preset::BaselineFastest);
4613        assert_eq!(encoder.limits, Limits::none());
4614    }
4615
4616    #[test]
4617    fn test_limits_has_limits() {
4618        assert!(!Limits::none().has_limits());
4619        assert!(Limits::default().max_width(100).has_limits());
4620        assert!(Limits::default().max_height(100).has_limits());
4621        assert!(Limits::default().max_pixel_count(1000).has_limits());
4622        assert!(Limits::default().max_alloc_bytes(1000).has_limits());
4623        assert!(Limits::default().max_icc_profile_bytes(1000).has_limits());
4624    }
4625
4626    // =========================================================================
4627    // Stop-based Cancellation Tests
4628    // =========================================================================
4629
4630    #[test]
4631    fn test_encode_rgb_with_stop_unstoppable() {
4632        use enough::Unstoppable;
4633        let encoder = Encoder::new(Preset::BaselineFastest);
4634        let pixels = vec![128u8; 64 * 64 * 3];
4635
4636        let result = encoder.encode_rgb_with_stop(&pixels, 64, 64, &Unstoppable);
4637        assert!(result.is_ok());
4638    }
4639
4640    #[test]
4641    fn test_encode_rgb_with_stop_pre_cancelled() {
4642        // Use CancellationContext with an already-set flag as our Stop impl
4643        let cancel = AtomicBool::new(true);
4644        let ctx = CancellationContext::new(Some(&cancel), None);
4645
4646        let encoder = Encoder::new(Preset::BaselineFastest);
4647        let pixels = vec![128u8; 64 * 64 * 3];
4648        let result = encoder.encode_rgb_with_stop(&pixels, 64, 64, &ctx);
4649
4650        assert!(matches!(result, Err(Error::Cancelled)));
4651    }
4652
4653    #[test]
4654    fn test_encode_gray_with_stop_unstoppable() {
4655        use enough::Unstoppable;
4656        let encoder = Encoder::new(Preset::BaselineFastest);
4657        let pixels = vec![128u8; 64 * 64];
4658
4659        let result = encoder.encode_gray_with_stop(&pixels, 64, 64, &Unstoppable);
4660        assert!(result.is_ok());
4661    }
4662
4663    #[test]
4664    fn test_encode_with_stop_limits() {
4665        use enough::Unstoppable;
4666        let limits = Limits::default().max_width(32);
4667        let encoder = Encoder::new(Preset::BaselineFastest).limits(limits);
4668
4669        let pixels = vec![128u8; 64 * 64 * 3];
4670        let result = encoder.encode_rgb_with_stop(&pixels, 64, 64, &Unstoppable);
4671
4672        assert!(matches!(result, Err(Error::DimensionLimitExceeded { .. })));
4673    }
4674
4675    #[test]
4676    fn test_cancellation_context_implements_stop() {
4677        let ctx = CancellationContext::none();
4678        assert!(!enough::Stop::may_stop(&ctx));
4679        assert!(enough::Stop::check(&ctx).is_ok());
4680
4681        let cancel = AtomicBool::new(false);
4682        let ctx = CancellationContext::new(Some(&cancel), None);
4683        assert!(enough::Stop::may_stop(&ctx));
4684        assert!(enough::Stop::check(&ctx).is_ok());
4685
4686        cancel.store(true, Ordering::Relaxed);
4687        assert!(enough::Stop::check(&ctx).is_err());
4688    }
4689
4690    #[test]
4691    fn test_cancellation_context_timeout_via_stop() {
4692        let ctx = CancellationContext {
4693            cancel: None,
4694            deadline: Some(Instant::now() - Duration::from_secs(1)),
4695        };
4696        assert!(enough::Stop::may_stop(&ctx));
4697        let err = enough::Stop::check(&ctx).unwrap_err();
4698        assert!(matches!(err, enough::StopReason::TimedOut));
4699    }
4700
4701    #[test]
4702    fn test_stop_reason_to_error_cancelled() {
4703        let err: Error = enough::StopReason::Cancelled.into();
4704        assert!(matches!(err, Error::Cancelled));
4705    }
4706
4707    #[test]
4708    fn test_stop_reason_to_error_timed_out() {
4709        let err: Error = enough::StopReason::TimedOut.into();
4710        assert!(matches!(err, Error::TimedOut));
4711    }
4712
4713    // =========================================================================
4714    // Deprecated Cancellation Tests (backwards compat)
4715    // =========================================================================
4716
4717    #[test]
4718    #[allow(deprecated)]
4719    fn test_cancellable_with_no_cancellation() {
4720        let encoder = Encoder::new(Preset::BaselineFastest);
4721        let pixels = vec![128u8; 64 * 64 * 3];
4722
4723        let result = encoder.encode_rgb_cancellable(&pixels, 64, 64, None, None);
4724
4725        assert!(result.is_ok());
4726    }
4727
4728    #[test]
4729    #[allow(deprecated)]
4730    fn test_cancellable_immediate_cancel() {
4731        let encoder = Encoder::new(Preset::BaselineFastest);
4732        let pixels = vec![128u8; 64 * 64 * 3];
4733        let cancel = AtomicBool::new(true); // Already cancelled
4734
4735        let result = encoder.encode_rgb_cancellable(&pixels, 64, 64, Some(&cancel), None);
4736
4737        assert!(matches!(result, Err(Error::Cancelled)));
4738    }
4739
4740    #[test]
4741    #[allow(deprecated)]
4742    fn test_cancellable_with_timeout() {
4743        let encoder = Encoder::new(Preset::BaselineFastest);
4744        let pixels = vec![128u8; 64 * 64 * 3];
4745
4746        // 10 second timeout - should complete well within this
4747        let result =
4748            encoder.encode_rgb_cancellable(&pixels, 64, 64, None, Some(Duration::from_secs(10)));
4749
4750        assert!(result.is_ok());
4751    }
4752
4753    #[test]
4754    #[allow(deprecated)]
4755    fn test_cancellable_gray() {
4756        let encoder = Encoder::new(Preset::BaselineFastest);
4757        let pixels = vec![128u8; 64 * 64];
4758
4759        let result = encoder.encode_gray_cancellable(&pixels, 64, 64, None, None);
4760
4761        assert!(result.is_ok());
4762    }
4763
4764    #[test]
4765    #[allow(deprecated)]
4766    fn test_cancellable_with_limits() {
4767        // Test that limits work in cancellable method too
4768        let limits = Limits::default().max_width(32);
4769        let encoder = Encoder::new(Preset::BaselineFastest).limits(limits);
4770
4771        let pixels = vec![128u8; 64 * 64 * 3];
4772        let result = encoder.encode_rgb_cancellable(&pixels, 64, 64, None, None);
4773
4774        assert!(matches!(result, Err(Error::DimensionLimitExceeded { .. })));
4775    }
4776
4777    #[test]
4778    fn test_cancellation_context_none() {
4779        let ctx = CancellationContext::none();
4780        assert!(ctx.check().is_ok());
4781    }
4782
4783    #[test]
4784    fn test_cancellation_context_with_cancel_flag() {
4785        use std::sync::atomic::Ordering;
4786
4787        let cancel = AtomicBool::new(false);
4788        let ctx = CancellationContext::new(Some(&cancel), None);
4789        assert!(ctx.check().is_ok());
4790
4791        cancel.store(true, Ordering::Relaxed);
4792        assert!(matches!(ctx.check(), Err(Error::Cancelled)));
4793    }
4794
4795    #[test]
4796    fn test_cancellation_context_with_expired_deadline() {
4797        // Create a deadline that's already passed
4798        let ctx = CancellationContext {
4799            cancel: None,
4800            deadline: Some(Instant::now() - Duration::from_secs(1)),
4801        };
4802
4803        assert!(matches!(ctx.check(), Err(Error::TimedOut)));
4804    }
4805
4806    #[test]
4807    fn test_dimension_exact_at_limit_passes() {
4808        // Dimensions exactly at limit should pass
4809        let limits = Limits::default().max_width(64).max_height(64);
4810        let encoder = Encoder::new(Preset::BaselineFastest).limits(limits);
4811
4812        let pixels = vec![128u8; 64 * 64 * 3];
4813        let result = encoder.encode_rgb(&pixels, 64, 64);
4814
4815        assert!(result.is_ok());
4816    }
4817
4818    #[test]
4819    fn test_pixel_count_exact_at_limit_passes() {
4820        // Pixel count exactly at limit should pass
4821        let limits = Limits::default().max_pixel_count(4096); // Exactly 64*64
4822        let encoder = Encoder::new(Preset::BaselineFastest).limits(limits);
4823
4824        let pixels = vec![128u8; 64 * 64 * 3];
4825        let result = encoder.encode_rgb(&pixels, 64, 64);
4826
4827        assert!(result.is_ok());
4828    }
4829
4830    #[test]
4831    fn test_multiple_limits_all_checked() {
4832        // Test that all limits are checked, not just the first
4833        let limits = Limits::default()
4834            .max_width(1000)
4835            .max_height(1000)
4836            .max_pixel_count(100); // This should fail
4837
4838        let encoder = Encoder::new(Preset::BaselineFastest).limits(limits);
4839        let pixels = vec![128u8; 64 * 64 * 3]; // 4096 pixels
4840
4841        let result = encoder.encode_rgb(&pixels, 64, 64);
4842        assert!(matches!(result, Err(Error::PixelCountExceeded { .. })));
4843    }
4844
4845    #[test]
4846    fn test_limits_with_grayscale() {
4847        let limits = Limits::default().max_pixel_count(100);
4848        let encoder = Encoder::new(Preset::BaselineFastest).limits(limits);
4849
4850        let pixels = vec![128u8; 64 * 64]; // Grayscale, 4096 pixels
4851        let result = encoder.encode_gray(&pixels, 64, 64);
4852
4853        assert!(matches!(result, Err(Error::PixelCountExceeded { .. })));
4854    }
4855
4856    #[test]
4857    fn test_estimate_resources_with_subsampling() {
4858        let encoder_444 = Encoder::new(Preset::BaselineBalanced).subsampling(Subsampling::S444);
4859        let encoder_420 = Encoder::new(Preset::BaselineBalanced).subsampling(Subsampling::S420);
4860
4861        let est_444 = encoder_444.estimate_resources(512, 512);
4862        let est_420 = encoder_420.estimate_resources(512, 512);
4863
4864        // 4:4:4 should use more memory than 4:2:0 (no chroma downsampling)
4865        assert!(
4866            est_444.peak_memory_bytes > est_420.peak_memory_bytes,
4867            "4:4:4 memory {} should exceed 4:2:0 memory {}",
4868            est_444.peak_memory_bytes,
4869            est_420.peak_memory_bytes
4870        );
4871    }
4872
4873    #[test]
4874    fn test_estimate_resources_block_count() {
4875        // With 4:2:0 subsampling (default): Y gets full blocks, chroma gets 1/4
4876        let encoder = Encoder::new(Preset::BaselineFastest);
4877
4878        // 64x64 image with 4:2:0:
4879        // Y blocks: 8x8 = 64
4880        // Chroma: 32x32 pixels, 4x4 blocks each = 16 per component
4881        // Total: 64 + 16 + 16 = 96
4882        let estimate = encoder.estimate_resources(64, 64);
4883        assert_eq!(estimate.block_count, 96);
4884
4885        // With 4:4:4 subsampling: all components get full blocks
4886        let encoder_444 = Encoder::new(Preset::BaselineFastest).subsampling(Subsampling::S444);
4887        let estimate_444 = encoder_444.estimate_resources(64, 64);
4888        // 64 blocks * 3 components = 192
4889        assert_eq!(estimate_444.block_count, 192);
4890    }
4891
4892    #[test]
4893    #[allow(deprecated)]
4894    fn test_cancellable_gray_with_limits() {
4895        let limits = Limits::default().max_width(32);
4896        let encoder = Encoder::new(Preset::BaselineFastest).limits(limits);
4897
4898        let pixels = vec![128u8; 64 * 64];
4899        let result = encoder.encode_gray_cancellable(&pixels, 64, 64, None, None);
4900
4901        assert!(matches!(result, Err(Error::DimensionLimitExceeded { .. })));
4902    }
4903}