Skip to main content

av_denoise_core/frame/
mod.rs

1use std::collections::VecDeque;
2
3use crate::accelerate::Accelerator;
4use crate::{
5    Algorithm,
6    ChannelMode,
7    Denoiser,
8    DenoiserError,
9    DenoiserOptions,
10    DenoisingMode,
11    Depth,
12    Device,
13    Nl4dOptions,
14    NlmTuning,
15    NlmeansHqOptions,
16    NlmeansOptions,
17    WindowSpan,
18};
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum Subsampling {
22    Yuv420,
23    Yuv422,
24    Yuv444,
25}
26
27impl Subsampling {
28    pub fn chroma_dims(self, w: u32, h: u32) -> (u32, u32) {
29        match self {
30            Subsampling::Yuv420 => (w / 2, h / 2),
31            Subsampling::Yuv422 => (w / 2, h),
32            Subsampling::Yuv444 => (w, h),
33        }
34    }
35}
36
37#[derive(Debug, Clone, Copy)]
38pub struct FrameLayout {
39    pub width: u32,
40    pub height: u32,
41    pub subsampling: Subsampling,
42    pub depth: Depth,
43}
44
45impl FrameLayout {
46    pub fn luma_pixels(&self) -> usize {
47        (self.width as usize) * (self.height as usize)
48    }
49
50    pub fn chroma_dims(&self) -> (u32, u32) {
51        self.subsampling.chroma_dims(self.width, self.height)
52    }
53
54    pub fn chroma_pixels(&self) -> usize {
55        let (w, h) = self.chroma_dims();
56        (w as usize) * (h as usize)
57    }
58
59    /// Wire size of the luma plane.
60    pub fn luma_bytes(&self) -> usize {
61        self.luma_pixels() * self.depth.bytes_per_sample()
62    }
63
64    /// Wire size of one chroma plane.
65    pub fn chroma_bytes(&self) -> usize {
66        self.chroma_pixels() * self.depth.bytes_per_sample()
67    }
68
69    /// A full black luma plane, used when no luma source is available.
70    pub fn black_luma_plane(&self) -> Vec<u8> {
71        fill_plane(self.luma_pixels(), 0, self.depth)
72    }
73
74    /// A full neutral chroma plane, used when a source has no chroma.
75    pub fn neutral_chroma_plane(&self) -> Vec<u8> {
76        fill_plane(self.chroma_pixels(), self.depth.neutral_chroma(), self.depth)
77    }
78}
79
80/// Builds a plane of `samples` copies of `value` in wire-byte form.
81pub fn fill_plane(samples: usize, value: u16, depth: Depth) -> Vec<u8> {
82    match depth.bytes_per_sample() {
83        1 => vec![value as u8; samples],
84        _ => {
85            let word = value.to_le_bytes();
86            let mut out = Vec::with_capacity(samples * 2);
87            for _ in 0..samples {
88                out.extend_from_slice(&word);
89            }
90            out
91        },
92    }
93}
94
95/// A planar YUV frame holding little-endian wire bytes.
96///
97/// Plane lengths come from [`FrameLayout`], so `y.len()` is
98/// `layout.luma_bytes()` and both `u.len()` and `v.len()` are
99/// `layout.chroma_bytes()`.
100#[derive(Debug, Clone)]
101pub struct Planes {
102    pub y: Vec<u8>,
103    pub u: Vec<u8>,
104    pub v: Vec<u8>,
105}
106
107/// Which planes a caller wants cleaned, once `--channel-mode` (or the
108/// equivalent host option) has been resolved.
109///
110/// This is separate from the library's [`ChannelMode`] because this layer
111/// may run more than one `Denoiser` in lockstep, one for luma and one for
112/// chroma. It may also run a single fused three-channel denoiser instead.
113/// Which of those applies depends on the caller's channel selection and
114/// the source's chroma subsampling.
115#[derive(Debug, Copy, Clone, PartialEq, Eq)]
116pub enum ChannelIntent {
117    /// Denoise luma only. Chroma passes through.
118    Luma,
119    /// Denoise chroma only. Luma passes through.
120    Chroma,
121    /// Denoise both luma and chroma as two independent denoisers.
122    /// Chroma runs at the source's native subsampled resolution.
123    LumaChroma,
124    /// A single library `Denoiser` running the fused three-channel
125    /// kernel. Needs a YUV444 source, which is checked at ingest setup
126    /// time.
127    YuvFused,
128}
129
130impl ChannelIntent {
131    /// Rejects the intent if the source's subsampling cannot support it.
132    pub fn validate_for_source(self, layout: FrameLayout) -> Result<(), anyhow::Error> {
133        match self {
134            ChannelIntent::YuvFused if layout.subsampling != Subsampling::Yuv444 => {
135                anyhow::bail!(
136                    "--channel-mode yuv requires a YUV444 source, got {:?}. Convert the input first, for example with `ffmpeg -pix_fmt yuv444p`",
137                    layout.subsampling
138                );
139            },
140            _ => Ok(()),
141        }
142    }
143}
144
145/// The per-plane option set a caller resolves once and passes into
146/// [`PlanarDenoiser::create`].
147#[derive(Debug, Clone)]
148pub struct PlaneOptions {
149    pub accelerators: Vec<Accelerator>,
150    pub device: Device,
151    pub intent: ChannelIntent,
152    pub mode: DenoisingMode,
153    /// Which denoising algorithm to run, along with the settings only
154    /// that algorithm reads.
155    pub algorithm: Algorithm,
156    /// Per-plane strength override for the luma denoiser. Takes
157    /// precedence over the algorithm's own `tuning.strength` when set.
158    /// Only has an effect on the two NLM algorithms.
159    pub luma_strength: Option<f32>,
160    /// Per-plane strength override for the chroma denoiser. Takes
161    /// precedence over the algorithm's own `tuning.strength` when set.
162    /// Only has an effect on the two NLM algorithms.
163    pub chroma_strength: Option<f32>,
164    /// Per-plane override for `lambda_ht`, luma. Takes precedence over
165    /// `algorithm`'s value when set, which itself falls back to a
166    /// calibrated per-plane default when nothing at all is set. Only
167    /// has an effect when `algorithm` is `Algorithm::Nl4d`, where it
168    /// pins the temporal grouping stage's hard threshold.
169    pub luma_lambda_ht: Option<f32>,
170    /// Per-plane override for `lambda_ht`, chroma. Takes precedence over
171    /// `algorithm`'s value when set, which itself falls back to a
172    /// calibrated per-plane default when nothing at all is set. Only
173    /// has an effect when `algorithm` is `Algorithm::Nl4d`, where it
174    /// pins the temporal grouping stage's hard threshold.
175    pub chroma_lambda_ht: Option<f32>,
176    /// Per-plane override for `mismatch_scale`, luma. Takes precedence
177    /// over `algorithm`'s value when set. Only has an effect when
178    /// `algorithm` is `Algorithm::Nl4d`.
179    pub luma_mismatch_scale: Option<f32>,
180    /// Per-plane override for `mismatch_scale`, chroma. Takes precedence
181    /// over `algorithm`'s value when set. Only has an effect when
182    /// `algorithm` is `Algorithm::Nl4d`.
183    pub chroma_mismatch_scale: Option<f32>,
184}
185
186impl PlaneOptions {
187    /// Resolves `self.algorithm` for one plane, folding in the per-plane
188    /// overrides that apply to whichever algorithm `self.algorithm` is.
189    ///
190    /// For the two NLM algorithms that is `strength`. For `Nl4d` it is
191    /// `lambda_ht`, since nl4d has no NLM weighting pass for a strength
192    /// to affect.
193    ///
194    /// `Nl4d`'s `lambda_ht` stays `Option<f32>` all the way through
195    /// this method. When neither a per-plane flag nor the matching
196    /// shared flag was set, the result is `None`, deferred to
197    /// `nl4d_default_lambda_ht` at construction, once the plane being
198    /// denoised is known there too. That is what gives luma and chroma
199    /// different values when a caller passes no flags at all.
200    fn algorithm_for(&self, channels: ChannelMode) -> Algorithm {
201        let per_plane = |luma, chroma| match channels {
202            ChannelMode::Luma => luma,
203            ChannelMode::Chroma => chroma,
204            ChannelMode::Yuv => None,
205        };
206
207        match self.algorithm {
208            Algorithm::Nl4d(nl4d) => Algorithm::Nl4d(Nl4dOptions {
209                // Left unresolved when unset, since the calibrated
210                // default depends on the plane, which
211                // `nl4d_default_lambda_ht` resolves at construction.
212                lambda_ht: per_plane(self.luma_lambda_ht, self.chroma_lambda_ht).or(nl4d.lambda_ht),
213                // Unlike `lambda_ht` this has one default for both
214                // planes, so an unset override simply leaves the shared
215                // value in place rather than deferring to construction.
216                mismatch_scale: per_plane(self.luma_mismatch_scale, self.chroma_mismatch_scale)
217                    .unwrap_or(nl4d.mismatch_scale),
218                ..nl4d
219            }),
220            Algorithm::Nlmeans(nlm) => {
221                let strength = per_plane(self.luma_strength, self.chroma_strength);
222                Algorithm::Nlmeans(with_plane_strength(nlm, strength))
223            },
224            Algorithm::NlmeansHq(opts) => {
225                let strength = per_plane(self.luma_strength, self.chroma_strength);
226                Algorithm::NlmeansHq(NlmeansHqOptions {
227                    nlm: with_plane_strength(opts.nlm, strength),
228                    ..opts
229                })
230            },
231        }
232    }
233
234    fn denoiser_options(&self, channels: ChannelMode) -> DenoiserOptions {
235        DenoiserOptions::builder()
236            .channel_mode(channels)
237            .mode(self.mode)
238            .algorithm(self.algorithm_for(channels))
239            .build()
240    }
241}
242
243/// `nlm` with `strength` replaced by the per-plane override, when there
244/// is one. An unset override leaves the shared value alone.
245fn with_plane_strength(nlm: NlmeansOptions, strength: Option<f32>) -> NlmeansOptions {
246    match strength {
247        None => nlm,
248        Some(strength) => NlmeansOptions {
249            tuning: NlmTuning {
250                strength: Some(strength),
251                ..nlm.tuning
252            },
253            ..nlm
254        },
255    }
256}
257
258/// Pops up to `count` entries off the front of `queue`, discarding them.
259fn drop_leading<T>(queue: &mut VecDeque<T>, count: usize) {
260    for _ in 0..count.min(queue.len()) {
261        queue.pop_front();
262    }
263}
264
265/// Reads the result of a `PlanarDenoiser::push` call for the
266/// push-then-drain-then-retry loop that `file_mode.rs` and
267/// `stream_mode.rs` both use.
268///
269/// `Ok(false)` means the push landed. `Ok(true)` means the queue was
270/// full, so the caller should drain one output and push again.
271///
272/// Any error other than `QueueFull` is passed on rather than discarded.
273pub fn push_needs_retry(result: Result<(), DenoiserError>) -> Result<bool, anyhow::Error> {
274    match result {
275        Ok(()) => Ok(false),
276        Err(DenoiserError::QueueFull) => Ok(true),
277        Err(other) => Err(other.into()),
278    }
279}
280
281/// Wraps the luma and chroma `Denoiser` instances needed for one
282/// subsampled YUV source.
283///
284/// The caller pushes planar frames in and gets planar frames out. The
285/// luma and chroma split is invisible from the outside.
286pub struct PlanarDenoiser {
287    layout: FrameLayout,
288    luma: Option<Denoiser>,
289    chroma: Option<Denoiser>,
290    /// Set when the intent is `YuvFused`, in which case `luma` and
291    /// `chroma` are both unset.
292    yuv: Option<Denoiser>,
293    // Source planes queued for passthrough when the matching denoiser is
294    // disabled. Only the disabled side's queue is ever filled. Entries
295    // are popped one per frame the enabled side emits, so temporal
296    // delays stay aligned.
297    luma_passthrough: VecDeque<Vec<u8>>,
298    chroma_passthrough: VecDeque<(Vec<u8>, Vec<u8>)>,
299    /// The temporal radius every owned denoiser runs at, resolved from
300    /// `opts.mode` at construction.
301    temporal_radius: u32,
302}
303
304impl PlanarDenoiser {
305    pub fn create(opts: &PlaneOptions, layout: FrameLayout) -> Result<Self, anyhow::Error> {
306        let (chroma_w, chroma_h) = layout.chroma_dims();
307
308        if chroma_w == 0 || chroma_h == 0 {
309            anyhow::bail!(
310                "frame dimensions {}x{} are too small for subsampling {:?}",
311                layout.width,
312                layout.height,
313                layout.subsampling
314            );
315        }
316
317        opts.intent.validate_for_source(layout)?;
318
319        let (denoise_luma, denoise_chroma, denoise_yuv) = match opts.intent {
320            ChannelIntent::Luma => (true, false, false),
321            ChannelIntent::Chroma => (false, true, false),
322            ChannelIntent::LumaChroma => (true, true, false),
323            ChannelIntent::YuvFused => (false, false, true),
324        };
325
326        let luma = denoise_luma
327            .then(|| {
328                Denoiser::create(
329                    &opts.accelerators,
330                    &opts.device,
331                    layout.width,
332                    layout.height,
333                    opts.denoiser_options(ChannelMode::Luma),
334                )
335            })
336            .transpose()?;
337
338        let chroma = denoise_chroma
339            .then(|| {
340                Denoiser::create(
341                    &opts.accelerators,
342                    &opts.device,
343                    chroma_w,
344                    chroma_h,
345                    opts.denoiser_options(ChannelMode::Chroma),
346                )
347            })
348            .transpose()?;
349
350        let yuv = denoise_yuv
351            .then(|| {
352                Denoiser::create(
353                    &opts.accelerators,
354                    &opts.device,
355                    layout.width,
356                    layout.height,
357                    opts.denoiser_options(ChannelMode::Yuv),
358                )
359            })
360            .transpose()?;
361
362        let temporal_radius = match opts.mode {
363            DenoisingMode::Spacial => 0,
364            DenoisingMode::Temporal { radius } => radius,
365        };
366
367        Ok(Self {
368            layout,
369            luma,
370            chroma,
371            yuv,
372            luma_passthrough: VecDeque::new(),
373            chroma_passthrough: VecDeque::new(),
374            temporal_radius,
375        })
376    }
377
378    /// The temporal radius the underlying denoisers run at.
379    pub fn temporal_radius(&self) -> u32 {
380        self.temporal_radius
381    }
382
383    /// Pushes one planar frame.
384    ///
385    /// On `QueueFull` the caller should receive one frame and then retry
386    /// the whole call. Any other error is passed on unchanged.
387    ///
388    /// The denoiser push runs before either passthrough queue is
389    /// touched, so a retry replays the whole frame cleanly instead of
390    /// queueing the disabled side's plane twice.
391    ///
392    /// # Why a retry cannot duplicate a frame
393    ///
394    /// In `LumaChroma` mode `luma` and `chroma` are both real
395    /// `Denoiser`s with their own queues. A retry pushes again into
396    /// whichever half already succeeded, which would duplicate that
397    /// half's frame if the two could ever sit at different fill levels.
398    ///
399    /// They cannot. Both are built from the same `opts.mode`, so they
400    /// share a temporal radius and a `MAX_PENDING` ceiling. Every
401    /// successful push or receive moves both on by exactly one frame,
402    /// and a failed push moves neither, because the `QueueFull` check
403    /// runs before anything changes.
404    ///
405    /// So the two halves always enter this function with the same frame
406    /// count and the same pending depth, and the `QueueFull` check
407    /// inside `push_frame` answers the same way for each. If the luma
408    /// push succeeds then the chroma push succeeds too, which makes the
409    /// duplicate unreachable.
410    pub fn push(&mut self, planes: &Planes) -> Result<(), DenoiserError> {
411        self.push_with(planes, Denoiser::push_frame)
412    }
413
414    /// Uploads one planar frame into the temporal window without starting
415    /// a denoise.
416    ///
417    /// Mirrors [`Self::push`], down to queueing the disabled side's
418    /// passthrough plane, but no output is ever produced for this call.
419    /// This is how [`Self::reseed`] fills the window from an explicit
420    /// window of frames before the one real push that starts a denoise.
421    fn push_priming(&mut self, planes: &Planes) -> Result<(), DenoiserError> {
422        self.push_with(planes, Denoiser::push_frame_priming)
423    }
424
425    /// Shared body of [`Self::push`] and [`Self::push_priming`].
426    ///
427    /// `push_frame` is [`Denoiser::push_frame`] for a real push or
428    /// [`Denoiser::push_frame_priming`] for a priming one, run against
429    /// whichever of `yuv`, `luma`, and `chroma` is enabled.
430    fn push_with(
431        &mut self,
432        planes: &Planes,
433        push_frame: fn(&mut Denoiser, &[f32]) -> Result<(), DenoiserError>,
434    ) -> Result<(), DenoiserError> {
435        if let Some(d) = self.yuv.as_mut() {
436            let buf = interleave_yuv_to_f32(&planes.y, &planes.u, &planes.v, self.layout.depth);
437            push_frame(d, &buf)?;
438            return Ok(());
439        }
440
441        if let Some(d) = self.luma.as_mut() {
442            let buf = plane_to_f32(&planes.y, self.layout.depth);
443            push_frame(d, &buf)?;
444        }
445
446        if let Some(d) = self.chroma.as_mut() {
447            let buf = interleave_uv_to_f32(&planes.u, &planes.v, self.layout.depth);
448            push_frame(d, &buf)?;
449        }
450
451        if self.luma.is_none() {
452            self.luma_passthrough.push_back(planes.y.clone());
453        }
454
455        if self.chroma.is_none() {
456            self.chroma_passthrough
457                .push_back((planes.u.clone(), planes.v.clone()));
458        }
459
460        Ok(())
461    }
462
463    /// Blocks until each enabled half emits one frame, then reassembles
464    /// them into a planar frame.
465    ///
466    /// Returns `Ok(None)` if neither half had pending output.
467    pub fn recv(&mut self) -> Result<Option<Planes>, anyhow::Error> {
468        if let Some(d) = self.yuv.as_mut() {
469            return match d.recv_frame()? {
470                Some(packed) => Ok(Some(unpack_yuv_from_f32(
471                    &packed,
472                    self.layout.luma_pixels(),
473                    self.layout.depth,
474                ))),
475                None => Ok(None),
476            };
477        }
478
479        let luma_out = self.luma.as_mut().map(|d| d.recv_frame()).transpose()?.flatten();
480
481        let chroma_out = self
482            .chroma
483            .as_mut()
484            .map(|d| d.recv_frame())
485            .transpose()?
486            .flatten();
487
488        // A disabled side has no Denoiser to query. When the enabled side
489        // produced output, pop the matching source plane from the
490        // disabled side's passthrough queue instead.
491        let luma_passthrough = if self.luma.is_none() && chroma_out.is_some() {
492            self.luma_passthrough.pop_front()
493        } else {
494            None
495        };
496
497        let chroma_passthrough = if self.chroma.is_none() && luma_out.is_some() {
498            self.chroma_passthrough.pop_front()
499        } else {
500            None
501        };
502
503        if luma_out.is_none() && chroma_out.is_none() {
504            return Ok(None);
505        }
506
507        let planes = self.assemble(luma_out, chroma_out, luma_passthrough, chroma_passthrough);
508
509        Ok(Some(planes))
510    }
511
512    /// Drains the temporal tail of both halves.
513    ///
514    /// `sink` is called once per emitted planar frame.
515    pub fn flush(&mut self, mut sink: impl FnMut(Planes)) -> Result<(), anyhow::Error> {
516        if let Some(d) = self.yuv.as_mut() {
517            let pixels = self.layout.luma_pixels();
518            let depth = self.layout.depth;
519            d.flush(|packed| sink(unpack_yuv_from_f32(&packed, pixels, depth)))?;
520            return Ok(());
521        }
522
523        let chroma_pixels = self.layout.chroma_pixels();
524
525        let mut luma_buf: Vec<Vec<f32>> = Vec::new();
526        let mut chroma_buf: Vec<Vec<f32>> = Vec::new();
527
528        if let Some(d) = self.luma.as_mut() {
529            d.flush(|v| luma_buf.push(v))?;
530        }
531
532        if let Some(d) = self.chroma.as_mut() {
533            d.flush(|v| chroma_buf.push(v))?;
534        }
535
536        // The two halves run in lockstep, so they flush the same number
537        // of frames. For each emitted frame the disabled side, if there
538        // is one, pops the matching source plane from its passthrough
539        // queue.
540        let count = luma_buf.len().max(chroma_buf.len());
541
542        for i in 0..count {
543            let y = if let Some(buf) = luma_buf.get(i) {
544                f32_to_plane(buf, self.layout.depth)
545            } else if let Some(src) = self.luma_passthrough.pop_front() {
546                src
547            } else {
548                self.layout.black_luma_plane()
549            };
550
551            let (u, v) = if let Some(packed) = chroma_buf.get(i) {
552                unpack_uv_from_f32(packed, chroma_pixels, self.layout.depth)
553            } else if let Some((src_u, src_v)) = self.chroma_passthrough.pop_front() {
554                (src_u, src_v)
555            } else {
556                (
557                    self.layout.neutral_chroma_plane(),
558                    self.layout.neutral_chroma_plane(),
559                )
560            };
561
562            sink(Planes { y, u, v });
563        }
564
565        if !self.luma_passthrough.is_empty() || !self.chroma_passthrough.is_empty() {
566            tracing::warn!(
567                luma_remaining = self.luma_passthrough.len(),
568                chroma_remaining = self.chroma_passthrough.len(),
569                "passthrough queue not fully drained after flush",
570            );
571            self.luma_passthrough.clear();
572            self.chroma_passthrough.clear();
573        }
574
575        Ok(())
576    }
577
578    /// The number of frames behind and ahead of a target frame a
579    /// [`Self::reseed`] window must supply, for whichever algorithm this
580    /// `PlanarDenoiser` runs.
581    ///
582    /// Every owned `Denoiser` was built from the same algorithm, so any
583    /// one of them answers for all of them.
584    pub fn window_span(&self) -> WindowSpan {
585        self.yuv
586            .as_ref()
587            .or(self.luma.as_ref())
588            .or(self.chroma.as_ref())
589            .expect("PlanarDenoiser always keeps at least one Denoiser")
590            .window_span()
591    }
592
593    /// Denoises the target frame of an explicit window, sized and
594    /// shaped exactly as [`Self::window_span`] reports for whichever
595    /// algorithm this `PlanarDenoiser` runs.
596    ///
597    /// This abandons whatever stream was running and starts a new one
598    /// from the window, keeping every GPU allocation. When it returns,
599    /// the stream sits exactly where it would be had the window been
600    /// pushed frame by frame, so the caller can carry on with
601    /// [`Self::push`] and [`Self::recv`] for the frame after the target.
602    ///
603    /// Callers clamp the window's indices at the clip's ends, matching
604    /// how the streaming path repeats the first and last frames.
605    ///
606    /// # Why the window is wider than `2r+1` for some algorithms
607    ///
608    /// The two NLM algorithms produce one output per submit once their
609    /// own `2r+1`-frame window is full, so a symmetric window centred
610    /// on the target frame is enough.
611    ///
612    /// nl4d scatters every pass's contribution across the `2r+1`
613    /// frames that pass reaches, and a target frame's own region only
614    /// starts collecting contributions once the earliest pass able to
615    /// reach it, the one centred `r` frames behind the target, has
616    /// actually run, which itself needs the front end's own window
617    /// full at that earlier centre. Both of those requirements push
618    /// the target's own `r`-wide neighbourhood back by another `r`, on
619    /// both sides, which is exactly what [`Self::window_span`] reports
620    /// through nl4d's doubled `behind` and `ahead`. This is bit-exact
621    /// with the streaming path because every frame the window supplies
622    /// is real, distinct content, run through the same sequence of
623    /// passes streaming would have run to reach the target frame.
624    pub fn reseed(&mut self, window: &[Planes]) -> Result<Planes, anyhow::Error> {
625        let span = self.window_span();
626        let expected = span.frame_count();
627        if window.len() != expected {
628            anyhow::bail!("reseed needs a window of {expected} frames, got {}", window.len());
629        }
630
631        self.luma_passthrough.clear();
632        self.chroma_passthrough.clear();
633
634        for d in [self.yuv.as_mut(), self.luma.as_mut(), self.chroma.as_mut()]
635            .into_iter()
636            .flatten()
637        {
638            d.reset_stream();
639        }
640
641        // Prime the first `2 * temporal_radius` frames, filling the
642        // underlying denoiser's own window without submitting anything,
643        // exactly as streaming would have primed it. This count comes
644        // from the front end's own window size, not from `span`, so it
645        // stays the same for every algorithm. Every remaining frame is
646        // then a real push, one submit per frame.
647        let radius = self.temporal_radius as usize;
648        let priming_count = 2 * radius;
649        let (head, tail) = window.split_at(priming_count);
650        for planes in head {
651            self.push_priming(planes)?;
652        }
653
654        // Priming queues one passthrough entry per frame, just as a
655        // real push does. `nlmeans`'s single real push, below, always
656        // emits and pairs with the target's own entry once `radius` of
657        // these leading ones are out of the way, exactly as before.
658        //
659        // nl4d's real pushes below emit more than once: nl4d's own
660        // gate gives every push once its own window is full a real
661        // output, but only the last `ahead - behind + 1` of them
662        // complete a region as new as the target's, the earlier ones
663        // complete regions further behind it that this call has no use
664        // for. Draining after every real push, not only the last,
665        // keeps the pending queue from ever holding more than one
666        // frame at a time, and it walks the passthrough queue forward
667        // by exactly one entry per region completed, so by the time
668        // the target's own region completes, its entry is the one at
669        // the front to pop. The same `radius` leading drop lines that
670        // front up correctly beforehand for both algorithms, because
671        // nl4d's own gate width is `radius` regardless of how wide
672        // `span` is.
673        drop_leading(&mut self.luma_passthrough, radius);
674        drop_leading(&mut self.chroma_passthrough, radius);
675
676        let mut result = None;
677        for planes in tail {
678            self.push(planes)?;
679            if let Some(out) = self.recv()? {
680                result = Some(out);
681            }
682        }
683
684        result.ok_or_else(|| anyhow::anyhow!("a full window produced no frame, this is a bug"))
685    }
686
687    fn assemble(
688        &self,
689        luma: Option<Vec<f32>>,
690        chroma: Option<Vec<f32>>,
691        luma_passthrough: Option<Vec<u8>>,
692        chroma_passthrough: Option<(Vec<u8>, Vec<u8>)>,
693    ) -> Planes {
694        let chroma_pixels = self.layout.chroma_pixels();
695
696        let y = match (luma, luma_passthrough) {
697            (Some(v), _) => f32_to_plane(&v, self.layout.depth),
698            (None, Some(src)) => src,
699            (None, None) => self.layout.black_luma_plane(),
700        };
701
702        let (u, v) = match (chroma, chroma_passthrough) {
703            (Some(packed), _) => unpack_uv_from_f32(&packed, chroma_pixels, self.layout.depth),
704            (None, Some(src)) => src,
705            (None, None) => (
706                self.layout.neutral_chroma_plane(),
707                self.layout.neutral_chroma_plane(),
708            ),
709        };
710
711        Planes { y, u, v }
712    }
713}
714
715/// Reads and writes samples in one wire format.
716///
717/// The implementor is chosen once per conversion, which keeps the
718/// per-sample path free of depth branches.
719trait SampleCodec {
720    const BYTES: usize;
721
722    fn read(plane: &[u8], i: usize) -> u16;
723    fn write(plane: &mut [u8], i: usize, value: u16);
724}
725
726/// One byte per sample.
727struct Narrow;
728
729impl SampleCodec for Narrow {
730    const BYTES: usize = 1;
731
732    #[inline(always)]
733    fn read(plane: &[u8], i: usize) -> u16 {
734        plane[i] as u16
735    }
736
737    #[inline(always)]
738    fn write(plane: &mut [u8], i: usize, value: u16) {
739        plane[i] = value as u8;
740    }
741}
742
743/// Two bytes per sample, little-endian.
744struct Wide;
745
746impl SampleCodec for Wide {
747    const BYTES: usize = 2;
748
749    #[inline(always)]
750    fn read(plane: &[u8], i: usize) -> u16 {
751        u16::from_le_bytes([plane[2 * i], plane[2 * i + 1]])
752    }
753
754    #[inline(always)]
755    fn write(plane: &mut [u8], i: usize, value: u16) {
756        plane[2 * i..2 * i + 2].copy_from_slice(&value.to_le_bytes());
757    }
758}
759
760/// Quantises a normalised value to a native-depth sample.
761#[inline(always)]
762fn quantise(v: f32, max: f32) -> u16 {
763    (v.clamp(0.0, 1.0) * max + 0.5) as u16
764}
765
766/// Converts one wire-byte plane to normalised f32.
767pub fn plane_to_f32(plane: &[u8], depth: Depth) -> Vec<f32> {
768    let max = depth.max_value();
769
770    fn run<C: SampleCodec>(plane: &[u8], max: f32) -> Vec<f32> {
771        let samples = plane.len() / C::BYTES;
772        (0..samples).map(|i| C::read(plane, i) as f32 / max).collect()
773    }
774
775    match depth.bytes_per_sample() {
776        1 => run::<Narrow>(plane, max),
777        _ => run::<Wide>(plane, max),
778    }
779}
780
781/// Reverse of [`plane_to_f32`].
782pub fn f32_to_plane(plane: &[f32], depth: Depth) -> Vec<u8> {
783    let max = depth.max_value();
784
785    fn run<C: SampleCodec>(plane: &[f32], max: f32) -> Vec<u8> {
786        let mut out = vec![0u8; plane.len() * C::BYTES];
787        for (i, &v) in plane.iter().enumerate() {
788            C::write(&mut out, i, quantise(v, max));
789        }
790        out
791    }
792
793    match depth.bytes_per_sample() {
794        1 => run::<Narrow>(plane, max),
795        _ => run::<Wide>(plane, max),
796    }
797}
798
799/// Interleaves equal-length Y, U, and V planes from a YUV444 source into
800/// `[Y0, U0, V0, Y1, U1, V1, ...]` as f32 in `[0, 1]`.
801///
802/// This is the layout the library's fused three-channel kernel expects.
803pub fn interleave_yuv_to_f32(y: &[u8], u: &[u8], v: &[u8], depth: Depth) -> Vec<f32> {
804    debug_assert_eq!(y.len(), u.len());
805    debug_assert_eq!(u.len(), v.len());
806
807    let max = depth.max_value();
808
809    fn run<C: SampleCodec>(y: &[u8], u: &[u8], v: &[u8], max: f32) -> Vec<f32> {
810        let pixels = y.len() / C::BYTES;
811        let mut out = Vec::with_capacity(pixels * 3);
812
813        for i in 0..pixels {
814            out.push(C::read(y, i) as f32 / max);
815            out.push(C::read(u, i) as f32 / max);
816            out.push(C::read(v, i) as f32 / max);
817        }
818
819        out
820    }
821
822    match depth.bytes_per_sample() {
823        1 => run::<Narrow>(y, u, v, max),
824        _ => run::<Wide>(y, u, v, max),
825    }
826}
827
828/// Reverse of [`interleave_yuv_to_f32`].
829pub fn unpack_yuv_from_f32(packed: &[f32], pixels: usize, depth: Depth) -> Planes {
830    debug_assert_eq!(packed.len(), 3 * pixels);
831
832    let max = depth.max_value();
833
834    fn run<C: SampleCodec>(packed: &[f32], pixels: usize, max: f32) -> Planes {
835        let mut y = vec![0u8; pixels * C::BYTES];
836        let mut u = vec![0u8; pixels * C::BYTES];
837        let mut v = vec![0u8; pixels * C::BYTES];
838
839        for (i, chunk) in packed.as_chunks::<3>().0.iter().enumerate() {
840            C::write(&mut y, i, quantise(chunk[0], max));
841            C::write(&mut u, i, quantise(chunk[1], max));
842            C::write(&mut v, i, quantise(chunk[2], max));
843        }
844
845        Planes { y, u, v }
846    }
847
848    match depth.bytes_per_sample() {
849        1 => run::<Narrow>(packed, pixels, max),
850        _ => run::<Wide>(packed, pixels, max),
851    }
852}
853
854/// Interleaves separate U and V planes into `[U, V, U, V, ...]` as f32
855/// in `[0, 1]`.
856pub fn interleave_uv_to_f32(u: &[u8], v: &[u8], depth: Depth) -> Vec<f32> {
857    debug_assert_eq!(u.len(), v.len());
858
859    let max = depth.max_value();
860
861    fn run<C: SampleCodec>(u: &[u8], v: &[u8], max: f32) -> Vec<f32> {
862        let pixels = u.len() / C::BYTES;
863        let mut out = Vec::with_capacity(pixels * 2);
864
865        for i in 0..pixels {
866            out.push(C::read(u, i) as f32 / max);
867            out.push(C::read(v, i) as f32 / max);
868        }
869
870        out
871    }
872
873    match depth.bytes_per_sample() {
874        1 => run::<Narrow>(u, v, max),
875        _ => run::<Wide>(u, v, max),
876    }
877}
878
879/// Reverse of [`interleave_uv_to_f32`].
880pub fn unpack_uv_from_f32(packed: &[f32], chroma_pixels: usize, depth: Depth) -> (Vec<u8>, Vec<u8>) {
881    debug_assert_eq!(packed.len(), 2 * chroma_pixels);
882
883    let max = depth.max_value();
884
885    fn run<C: SampleCodec>(packed: &[f32], chroma_pixels: usize, max: f32) -> (Vec<u8>, Vec<u8>) {
886        let mut u = vec![0u8; chroma_pixels * C::BYTES];
887        let mut v = vec![0u8; chroma_pixels * C::BYTES];
888
889        for (i, chunk) in packed.as_chunks::<2>().0.iter().enumerate() {
890            C::write(&mut u, i, quantise(chunk[0], max));
891            C::write(&mut v, i, quantise(chunk[1], max));
892        }
893
894        (u, v)
895    }
896
897    match depth.bytes_per_sample() {
898        1 => run::<Narrow>(packed, chroma_pixels, max),
899        _ => run::<Wide>(packed, chroma_pixels, max),
900    }
901}
902
903#[cfg(test)]
904mod converter_tests {
905    use super::*;
906
907    /// Encodes native-depth samples into wire bytes, the inverse of what
908    /// the converters read.
909    fn wire(samples: &[u16], depth: Depth) -> Vec<u8> {
910        match depth.bytes_per_sample() {
911            1 => samples.iter().map(|&s| s as u8).collect(),
912            _ => samples.iter().flat_map(|&s| s.to_le_bytes()).collect(),
913        }
914    }
915
916    #[test]
917    fn plane_round_trips_boundary_codes_at_every_depth() {
918        for depth in [Depth::Eight, Depth::Ten, Depth::Twelve] {
919            let max = depth.max_value() as u16;
920            let samples: Vec<u16> = vec![0, 1, 16, 64, 235, max / 2, max - 1, max]
921                .into_iter()
922                .filter(|&s| s <= max)
923                .collect();
924
925            let bytes = wire(&samples, depth);
926            let restored = f32_to_plane(&plane_to_f32(&bytes, depth), depth);
927
928            assert_eq!(restored, bytes, "plane round trip failed at {depth:?}");
929        }
930    }
931
932    /// Samples above 8 bits are little-endian on the wire regardless of
933    /// host endianness.
934    #[test]
935    fn high_depth_samples_are_little_endian() {
936        // 1023 = 0x03FF -> [0xFF, 0x03]
937        let bytes = wire(&[1023, 0, 512], Depth::Ten);
938        assert_eq!(bytes, vec![0xFF, 0x03, 0x00, 0x00, 0x00, 0x02]);
939
940        let f = plane_to_f32(&bytes, Depth::Ten);
941        assert!(
942            (f[0] - 1.0).abs() < 1e-6,
943            "0x03FF should normalize to 1.0, got {}",
944            f[0]
945        );
946        assert_eq!(f[1], 0.0);
947    }
948
949    #[test]
950    fn uv_interleave_round_trips_at_every_depth() {
951        for depth in [Depth::Eight, Depth::Ten, Depth::Twelve] {
952            let max = depth.max_value() as u16;
953            let u_samples = vec![0, max / 4, max];
954            let v_samples = vec![max, max / 2, 1];
955
956            let u_bytes = wire(&u_samples, depth);
957            let v_bytes = wire(&v_samples, depth);
958
959            let packed = interleave_uv_to_f32(&u_bytes, &v_bytes, depth);
960            assert_eq!(packed.len(), 6, "packed UV length wrong at {depth:?}");
961
962            let (ru, rv) = unpack_uv_from_f32(&packed, 3, depth);
963            assert_eq!(ru, u_bytes, "U round trip failed at {depth:?}");
964            assert_eq!(rv, v_bytes, "V round trip failed at {depth:?}");
965        }
966    }
967
968    #[test]
969    fn yuv_interleave_round_trips_at_every_depth() {
970        for depth in [Depth::Eight, Depth::Ten, Depth::Twelve] {
971            let max = depth.max_value() as u16;
972            let y_samples = vec![0, max / 3, max];
973            let u_samples = vec![max, 0, max / 2];
974            let v_samples = vec![max / 4, max, 0];
975
976            let y_bytes = wire(&y_samples, depth);
977            let u_bytes = wire(&u_samples, depth);
978            let v_bytes = wire(&v_samples, depth);
979
980            let packed = interleave_yuv_to_f32(&y_bytes, &u_bytes, &v_bytes, depth);
981            assert_eq!(packed.len(), 9, "packed YUV length wrong at {depth:?}");
982
983            let out = unpack_yuv_from_f32(&packed, 3, depth);
984            assert_eq!(out.y, y_bytes, "Y round trip failed at {depth:?}");
985            assert_eq!(out.u, u_bytes, "U round trip failed at {depth:?}");
986            assert_eq!(out.v, v_bytes, "V round trip failed at {depth:?}");
987        }
988    }
989
990    #[test]
991    fn quantise_matches_the_clamping_form_including_nan() {
992        fn reference(v: f32, max: f32) -> u16 {
993            (v.clamp(0.0, 1.0) * max + 0.5) as u16
994        }
995
996        let max = 1023.0;
997        let cases = [
998            -1.0,
999            -0.001,
1000            0.0,
1001            0.5,
1002            0.999,
1003            1.0,
1004            1.001,
1005            2.0,
1006            f32::NAN,
1007            f32::INFINITY,
1008            f32::NEG_INFINITY,
1009        ];
1010
1011        for v in cases {
1012            assert_eq!(quantise(v, max), reference(v, max), "mismatch at {v}");
1013        }
1014    }
1015
1016    /// Limited-range codes normalise to matching values at every depth,
1017    /// which is the property the whole design rests on.
1018    ///
1019    /// The match is within one 8-bit code level rather than exact. ITU
1020    /// defines the limited-range endpoints as exact multiples, so 16
1021    /// becomes 64 and 235 becomes 940, but full scale is not a multiple,
1022    /// because 255 becomes 1023. That leaves 235/255 and 940/1023
1023    /// differing by 0.0027, roughly 0.69 of an 8-bit step.
1024    ///
1025    /// Agreement below one step is the real property here.
1026    ///
1027    /// `normalized_scale_is_identical_across_depths` in
1028    /// `src/nlmeans/mod.rs` pins the same property on the library's own
1029    /// normalise helper.
1030    #[test]
1031    fn limited_range_codes_agree_across_depths() {
1032        /// One 8-bit code level, the precision the endpoints agree to.
1033        const TOL: f32 = 1.0 / 255.0;
1034
1035        let eight = plane_to_f32(&wire(&[16, 235], Depth::Eight), Depth::Eight);
1036        let ten = plane_to_f32(&wire(&[64, 940], Depth::Ten), Depth::Ten);
1037
1038        for (a, b) in eight.iter().zip(ten.iter()) {
1039            assert!((a - b).abs() < TOL, "8-bit {a} vs 10-bit {b}");
1040        }
1041    }
1042}
1043
1044#[cfg(test)]
1045mod cli_options_tests {
1046    use super::*;
1047    use crate::nlmeans::NlmParams;
1048
1049    /// A `PlaneOptions` with every field at a neutral default, so each test
1050    /// only overrides what it cares about.
1051    ///
1052    /// `mode` and `algorithm` are the two fields every test below sets
1053    /// for itself.
1054    fn base_opts(
1055        mode: DenoisingMode,
1056        algorithm: Algorithm,
1057        luma_strength: Option<f32>,
1058        chroma_strength: Option<f32>,
1059    ) -> PlaneOptions {
1060        PlaneOptions {
1061            accelerators: vec![],
1062            device: Device::Default,
1063            intent: ChannelIntent::LumaChroma,
1064            mode,
1065            algorithm,
1066            luma_strength,
1067            chroma_strength,
1068            luma_lambda_ht: None,
1069            chroma_lambda_ht: None,
1070            luma_mismatch_scale: None,
1071            chroma_mismatch_scale: None,
1072        }
1073    }
1074
1075    #[test]
1076    fn luma_strength_alone_overrides_only_the_luma_plane() {
1077        let opts = base_opts(DenoisingMode::Spacial, Algorithm::default(), Some(0.7), None);
1078
1079        let luma = expect_nlmeans(opts.denoiser_options(ChannelMode::Luma).algorithm);
1080        let chroma = expect_nlmeans(opts.denoiser_options(ChannelMode::Chroma).algorithm);
1081
1082        assert!(
1083            matches!(luma.tuning.strength, Some(s) if (s - 0.7).abs() < f32::EPSILON),
1084            "expected luma tuning.strength = Some(0.7), got {:?}",
1085            luma.tuning.strength
1086        );
1087        assert_eq!(
1088            chroma.tuning.strength, None,
1089            "chroma plane should carry no override so the table default applies"
1090        );
1091    }
1092
1093    #[test]
1094    fn both_per_plane_strengths_set_independently() {
1095        let opts = base_opts(DenoisingMode::Spacial, Algorithm::default(), Some(0.7), Some(0.3));
1096
1097        let luma = expect_nlmeans(opts.denoiser_options(ChannelMode::Luma).algorithm);
1098        let chroma = expect_nlmeans(opts.denoiser_options(ChannelMode::Chroma).algorithm);
1099
1100        assert!(
1101            matches!(luma.tuning.strength, Some(s) if (s - 0.7).abs() < f32::EPSILON),
1102            "expected luma tuning.strength = Some(0.7), got {:?}",
1103            luma.tuning.strength
1104        );
1105        assert!(
1106            matches!(chroma.tuning.strength, Some(s) if (s - 0.3).abs() < f32::EPSILON),
1107            "expected chroma tuning.strength = Some(0.3), got {:?}",
1108            chroma.tuning.strength
1109        );
1110    }
1111
1112    #[test]
1113    fn no_overrides_hq_resolves_through_to_nlm_params_to_the_measured_tables() {
1114        // Radius 4 in the measured tables is luma 0.35 and chroma
1115        // 0.70 (see the table docs in `src/nlmeans/params.rs`).
1116        let opts = base_opts(
1117            DenoisingMode::Temporal { radius: 4 },
1118            Algorithm::NlmeansHq(NlmeansHqOptions::default()),
1119            None,
1120            None,
1121        );
1122
1123        let luma_params: NlmParams = opts.denoiser_options(ChannelMode::Luma).to_nlm_params();
1124        let chroma_params: NlmParams = opts.denoiser_options(ChannelMode::Chroma).to_nlm_params();
1125
1126        assert!(
1127            (luma_params.strength - 0.35).abs() < f32::EPSILON,
1128            "expected luma strength 0.35 at r4, got {}",
1129            luma_params.strength
1130        );
1131        assert!(
1132            (chroma_params.strength - 0.70).abs() < f32::EPSILON,
1133            "expected chroma strength 0.70 at r4, got {}",
1134            chroma_params.strength
1135        );
1136    }
1137
1138    /// A `PlaneOptions` running `Algorithm::Nl4d`, with every field at a
1139    /// neutral default except the two per-plane `lambda_ht` overrides
1140    /// under test.
1141    fn nl4d_opts(luma_lambda_ht: Option<f32>, chroma_lambda_ht: Option<f32>) -> PlaneOptions {
1142        PlaneOptions {
1143            accelerators: vec![],
1144            device: Device::Default,
1145            intent: ChannelIntent::LumaChroma,
1146            mode: DenoisingMode::Temporal { radius: 2 },
1147            algorithm: Algorithm::Nl4d(Nl4dOptions::default()),
1148            luma_strength: None,
1149            chroma_strength: None,
1150            luma_lambda_ht,
1151            chroma_lambda_ht,
1152            luma_mismatch_scale: None,
1153            chroma_mismatch_scale: None,
1154        }
1155    }
1156
1157    /// Unwraps an `Algorithm::Nlmeans`, panicking with the whole value
1158    /// on any other variant.
1159    fn expect_nlmeans(algorithm: Algorithm) -> NlmeansOptions {
1160        match algorithm {
1161            Algorithm::Nlmeans(n) => n,
1162            other => panic!("expected Algorithm::Nlmeans, got {other:?}"),
1163        }
1164    }
1165
1166    /// Unwraps an `Algorithm::Nl4d`, panicking with the whole value on
1167    /// any other variant.
1168    fn expect_nl4d(algorithm: Algorithm) -> Nl4dOptions {
1169        match algorithm {
1170            Algorithm::Nl4d(n) => n,
1171            other => panic!("expected Algorithm::Nl4d, got {other:?}"),
1172        }
1173    }
1174
1175    /// A `PlaneOptions` running `Algorithm::Nl4d` with a shared
1176    /// `mismatch_scale` and the two per-plane overrides under test.
1177    fn nl4d_mismatch_opts(
1178        shared: f32,
1179        luma_mismatch_scale: Option<f32>,
1180        chroma_mismatch_scale: Option<f32>,
1181    ) -> PlaneOptions {
1182        PlaneOptions {
1183            algorithm: Algorithm::Nl4d(Nl4dOptions {
1184                mismatch_scale: shared,
1185                ..Nl4dOptions::default()
1186            }),
1187            luma_mismatch_scale,
1188            chroma_mismatch_scale,
1189            ..nl4d_opts(None, None)
1190        }
1191    }
1192
1193    /// The same routing property the `lambda_ht` pair is checked for,
1194    /// applied to `mismatch_scale`. An override aimed at one plane must
1195    /// leave the other on the shared value, which for this field is a
1196    /// resolved number rather than a deferred `None`.
1197    #[test]
1198    fn a_per_plane_mismatch_scale_overrides_only_its_own_instance_for_nl4d() {
1199        let luma_only = nl4d_mismatch_opts(2.0, Some(8.0), None);
1200        let luma = expect_nl4d(luma_only.algorithm_for(ChannelMode::Luma));
1201        let chroma = expect_nl4d(luma_only.algorithm_for(ChannelMode::Chroma));
1202        assert!((luma.mismatch_scale - 8.0).abs() < f32::EPSILON);
1203        assert!(
1204            (chroma.mismatch_scale - 2.0).abs() < f32::EPSILON,
1205            "chroma should keep the shared value, got {}",
1206            chroma.mismatch_scale
1207        );
1208
1209        let chroma_only = nl4d_mismatch_opts(2.0, None, Some(8.0));
1210        let luma = expect_nl4d(chroma_only.algorithm_for(ChannelMode::Luma));
1211        let chroma = expect_nl4d(chroma_only.algorithm_for(ChannelMode::Chroma));
1212        assert!((chroma.mismatch_scale - 8.0).abs() < f32::EPSILON);
1213        assert!(
1214            (luma.mismatch_scale - 2.0).abs() < f32::EPSILON,
1215            "luma should keep the shared value, got {}",
1216            luma.mismatch_scale
1217        );
1218    }
1219
1220    /// A fused Yuv pass has no plane to pick, so neither override
1221    /// applies and the shared value stands.
1222    #[test]
1223    fn a_yuv_instance_ignores_both_per_plane_mismatch_scales() {
1224        let opts = nl4d_mismatch_opts(2.0, Some(8.0), Some(4.0));
1225        let yuv = expect_nl4d(opts.algorithm_for(ChannelMode::Yuv));
1226
1227        assert!((yuv.mismatch_scale - 2.0).abs() < f32::EPSILON);
1228    }
1229
1230    /// The routing property that matters most for a shared field: an
1231    /// override aimed at one plane must never leak into the other
1232    /// instance. `luma_lambda_ht` set alone must change nothing about
1233    /// the chroma instance, and vice versa in the sibling test below.
1234    #[test]
1235    fn luma_lambda_ht_alone_overrides_only_the_luma_instance_for_nl4d() {
1236        let opts = nl4d_opts(Some(4.0), None);
1237
1238        let luma = expect_nl4d(opts.algorithm_for(ChannelMode::Luma));
1239        let chroma = expect_nl4d(opts.algorithm_for(ChannelMode::Chroma));
1240
1241        assert!((luma.lambda_ht.unwrap() - 4.0).abs() < f32::EPSILON);
1242        assert_eq!(
1243            chroma.lambda_ht,
1244            Nl4dOptions::default().lambda_ht,
1245            "chroma should stay unresolved here (None), deferred to its own per-plane \
1246             default at construction, got {:?}",
1247            chroma.lambda_ht
1248        );
1249    }
1250
1251    #[test]
1252    fn chroma_lambda_ht_alone_overrides_only_the_chroma_instance_for_nl4d() {
1253        let opts = nl4d_opts(None, Some(4.0));
1254
1255        let luma = expect_nl4d(opts.algorithm_for(ChannelMode::Luma));
1256        let chroma = expect_nl4d(opts.algorithm_for(ChannelMode::Chroma));
1257
1258        assert_eq!(
1259            luma.lambda_ht,
1260            Nl4dOptions::default().lambda_ht,
1261            "luma should stay unresolved here (None), deferred to its own per-plane \
1262             default at construction, got {:?}",
1263            luma.lambda_ht
1264        );
1265        assert!((chroma.lambda_ht.unwrap() - 4.0).abs() < f32::EPSILON);
1266    }
1267
1268    #[test]
1269    fn both_planes_lambda_ht_set_independently_for_nl4d() {
1270        let opts = nl4d_opts(Some(2.0), Some(3.5));
1271
1272        let luma = expect_nl4d(opts.algorithm_for(ChannelMode::Luma));
1273        let chroma = expect_nl4d(opts.algorithm_for(ChannelMode::Chroma));
1274
1275        assert!((luma.lambda_ht.unwrap() - 2.0).abs() < f32::EPSILON);
1276        assert!((chroma.lambda_ht.unwrap() - 3.5).abs() < f32::EPSILON);
1277
1278        // Every other field stays shared between the two instances even
1279        // though lambda_ht diverges.
1280        assert_eq!(luma.refine, chroma.refine);
1281        assert_eq!(luma.spatial_radius, chroma.spatial_radius);
1282        assert!((luma.c_min - chroma.c_min).abs() < f32::EPSILON);
1283    }
1284
1285    #[test]
1286    fn unset_nl4d_overrides_resolve_to_different_lambda_ht_per_plane_end_to_end() {
1287        let opts = nl4d_opts(None, None);
1288
1289        let luma = expect_nl4d(opts.algorithm_for(ChannelMode::Luma));
1290        let chroma = expect_nl4d(opts.algorithm_for(ChannelMode::Chroma));
1291
1292        // Neither plane has anything set anywhere, so both stay
1293        // unresolved at this layer...
1294        assert_eq!(luma.lambda_ht, None);
1295        assert_eq!(chroma.lambda_ht, None);
1296
1297        // ...but resolving each through the same function construction
1298        // uses (`nl4d_default_lambda_ht`, see `src/denoiser.rs`) gives
1299        // luma and chroma different values, which is the whole point of
1300        // a caller passing no flags at all getting both calibrated
1301        // defaults.
1302        let luma_default = crate::nl4d_default_lambda_ht(ChannelMode::Luma);
1303        let chroma_default = crate::nl4d_default_lambda_ht(ChannelMode::Chroma);
1304        assert!((luma_default - 5.3).abs() < f32::EPSILON);
1305        assert!((chroma_default - 4.2).abs() < f32::EPSILON);
1306        assert!((chroma_default - luma_default).abs() > f32::EPSILON);
1307    }
1308}
1309
1310// Feature-gated because every test here builds its `PlaneOptions` from
1311// `chroma_only_opts`, which names the `Vulkan` accelerator variant. That
1312// variant only exists when the `vulkan` feature is enabled.
1313#[cfg(feature = "vulkan")]
1314#[cfg(test)]
1315mod passthrough_retry_tests {
1316    use super::*;
1317    use crate::accelerate::Accelerator;
1318    use crate::{Algorithm, DenoisingMode};
1319
1320    /// Chroma-only intent, so `luma` is the disabled passthrough half and
1321    /// `chroma` is the one that can report `QueueFull`.
1322    ///
1323    /// That is what drives the retry loop in `push_with_drain` and in
1324    /// `stream_mode.rs`.
1325    fn chroma_only_opts() -> PlaneOptions {
1326        PlaneOptions {
1327            accelerators: vec![Accelerator::Vulkan],
1328            device: Device::Default,
1329            intent: ChannelIntent::Chroma,
1330            mode: DenoisingMode::Spacial,
1331            algorithm: Algorithm::default(),
1332            luma_strength: None,
1333            chroma_strength: None,
1334            luma_lambda_ht: None,
1335            chroma_lambda_ht: None,
1336            luma_mismatch_scale: None,
1337            chroma_mismatch_scale: None,
1338        }
1339    }
1340
1341    fn fake_planes(layout: FrameLayout) -> Planes {
1342        Planes {
1343            y: fill_plane(layout.luma_pixels(), layout.depth.neutral_chroma(), layout.depth),
1344            u: layout.neutral_chroma_plane(),
1345            v: layout.neutral_chroma_plane(),
1346        }
1347    }
1348
1349    #[test]
1350    fn queue_full_retry_does_not_double_queue_the_passthrough_plane() {
1351        let layout = FrameLayout {
1352            width: 16,
1353            height: 16,
1354            subsampling: Subsampling::Yuv420,
1355            depth: Depth::Eight,
1356        };
1357        let mut wd =
1358            PlanarDenoiser::create(&chroma_only_opts(), layout).expect("denoiser construction failed");
1359        let planes = fake_planes(layout);
1360
1361        // Spatial mode runs a depth-2 pipeline, so the first two pushes
1362        // land directly. See `push_after_pending_returns_queue_full` in
1363        // `src/denoiser.rs`.
1364        wd.push(&planes).expect("first push should land");
1365        wd.push(&planes).expect("second push should land");
1366
1367        // Third push hits QueueFull on the chroma half.
1368        let err = wd.push(&planes).expect_err("expected QueueFull");
1369        assert!(
1370            matches!(err, DenoiserError::QueueFull),
1371            "expected QueueFull, got {err:?}"
1372        );
1373
1374        // Mirror the retry loop in `push_with_drain`. Drain one output,
1375        // then retry the whole `push()` call for the same frame.
1376        wd.recv().expect("recv after drain failed");
1377        wd.push(&planes).expect("retry push should land after drain");
1378
1379        // The chroma denoiser accepted three frames, two directly and
1380        // one on the retry, and `recv` popped one back off. The disabled
1381        // luma half's passthrough queue must track that one for one, and
1382        // must not count the frame whose first attempt hit `QueueFull`
1383        // twice.
1384        assert_eq!(
1385            wd.luma_passthrough.len(),
1386            2,
1387            "expected exactly one passthrough entry per chroma frame actually accepted, got {}",
1388            wd.luma_passthrough.len()
1389        );
1390    }
1391}
1392
1393// Feature-gated because every test here builds its `PlaneOptions` from
1394// `luma_chroma_opts`, which names the `Vulkan` accelerator variant. That
1395// variant only exists when the `vulkan` feature is enabled.
1396#[cfg(feature = "vulkan")]
1397#[cfg(test)]
1398mod lumachroma_lockstep_tests {
1399    use super::*;
1400    use crate::accelerate::Accelerator;
1401    use crate::{Algorithm, DenoisingMode};
1402
1403    /// Runs `luma` and `chroma` as two real `Denoiser`s in spatial mode.
1404    ///
1405    /// Spatial mode passes a uniform-valued plane through unchanged, as
1406    /// the `uniform_*_passthrough` tests in `src/nlmeans/tests` show. The
1407    /// test can therefore give each plane its own marker value and spot
1408    /// the two halves drifting apart.
1409    fn luma_chroma_opts() -> PlaneOptions {
1410        PlaneOptions {
1411            accelerators: vec![Accelerator::Vulkan],
1412            device: Device::Default,
1413            intent: ChannelIntent::LumaChroma,
1414            mode: DenoisingMode::Spacial,
1415            algorithm: Algorithm::default(),
1416            luma_strength: None,
1417            chroma_strength: None,
1418            luma_lambda_ht: None,
1419            chroma_lambda_ht: None,
1420            luma_mismatch_scale: None,
1421            chroma_mismatch_scale: None,
1422        }
1423    }
1424
1425    /// A uniform-valued frame whose luma and chroma planes each encode
1426    /// `idx` with a different formula.
1427    ///
1428    /// If the round trip ever pairs luma from one push with chroma from
1429    /// another, the two encodings disagree and the test catches it.
1430    fn marked_planes(layout: FrameLayout, idx: u8) -> Planes {
1431        let chroma_pixels = layout.chroma_pixels();
1432        let y_val = 10 + idx;
1433        let uv_val = 200 - idx;
1434
1435        Planes {
1436            y: fill_plane(layout.luma_pixels(), y_val as u16, layout.depth),
1437            u: fill_plane(chroma_pixels, uv_val as u16, layout.depth),
1438            v: fill_plane(chroma_pixels, uv_val as u16, layout.depth),
1439        }
1440    }
1441
1442    #[test]
1443    fn queue_full_retries_never_desync_luma_and_chroma() {
1444        let layout = FrameLayout {
1445            width: 16,
1446            height: 16,
1447            subsampling: Subsampling::Yuv420,
1448            depth: Depth::Eight,
1449        };
1450        let mut wd =
1451            PlanarDenoiser::create(&luma_chroma_opts(), layout).expect("denoiser construction failed");
1452
1453        // More pushes than the depth-2 pipeline holds, so this drives
1454        // several `QueueFull`-then-retry cycles.
1455        const N: u8 = 6;
1456        let mut outputs: Vec<Planes> = Vec::new();
1457
1458        for idx in 0..N {
1459            let planes = marked_planes(layout, idx);
1460
1461            // Mirror the retry loop in `push_with_drain` exactly, which
1462            // is the sequence `file_mode.rs` and `stream_mode.rs` run.
1463            if push_needs_retry(wd.push(&planes)).expect("push_needs_retry") {
1464                if let Some(out) = wd.recv().expect("recv failed") {
1465                    outputs.push(out);
1466                }
1467
1468                wd.push(&planes).expect("retry push should land after drain");
1469            }
1470        }
1471
1472        wd.flush(|out| outputs.push(out)).expect("flush failed");
1473
1474        assert_eq!(
1475            outputs.len(),
1476            N as usize,
1477            "expected exactly one output frame per input frame, got {}",
1478            outputs.len()
1479        );
1480
1481        for out in &outputs {
1482            let y_val = out.y[0];
1483            let uv_val = out.u[0];
1484            let idx_from_y = y_val - 10;
1485            let idx_from_uv = 200 - uv_val;
1486
1487            assert_eq!(
1488                idx_from_y, idx_from_uv,
1489                "luma marker {y_val} (frame {idx_from_y}) and chroma marker {uv_val} \
1490                 (frame {idx_from_uv}) disagree, so the luma and chroma pushes have drifted apart"
1491            );
1492        }
1493    }
1494}
1495
1496#[cfg(test)]
1497mod push_needs_retry_tests {
1498    use super::*;
1499
1500    #[test]
1501    fn ok_means_no_retry() {
1502        let outcome = push_needs_retry(Ok(())).expect("Ok(()) must not itself error");
1503        assert!(!outcome, "a landed push must not ask the caller to retry");
1504    }
1505
1506    #[test]
1507    fn queue_full_signals_retry() {
1508        let outcome =
1509            push_needs_retry(Err(DenoiserError::QueueFull)).expect("QueueFull must not itself error");
1510        assert!(outcome, "QueueFull must still trigger the retry-after-drain path");
1511    }
1512
1513    #[test]
1514    fn non_queue_full_errors_propagate_instead_of_being_swallowed() {
1515        let synthetic = DenoiserError::Other(anyhow::anyhow!("synthetic readback failure"));
1516
1517        let outcome = push_needs_retry(Err(synthetic));
1518
1519        assert!(
1520            outcome.is_err(),
1521            "a non-QueueFull push error must propagate instead of being silently treated as success"
1522        );
1523    }
1524}
1525
1526#[cfg(test)]
1527mod layout_tests {
1528    use super::*;
1529
1530    fn layout(depth: Depth) -> FrameLayout {
1531        FrameLayout {
1532            width: 4,
1533            height: 4,
1534            subsampling: Subsampling::Yuv420,
1535            depth,
1536        }
1537    }
1538
1539    #[test]
1540    fn byte_lengths_scale_with_depth() {
1541        assert_eq!(layout(Depth::Eight).luma_bytes(), 16);
1542        assert_eq!(layout(Depth::Ten).luma_bytes(), 32);
1543        assert_eq!(layout(Depth::Eight).chroma_bytes(), 4);
1544        assert_eq!(layout(Depth::Ten).chroma_bytes(), 8);
1545    }
1546
1547    #[test]
1548    fn neutral_chroma_fill_is_correct_at_each_depth() {
1549        let eight = layout(Depth::Eight).neutral_chroma_plane();
1550        assert_eq!(eight, vec![128u8; 4]);
1551
1552        // 512 little-endian is [0x00, 0x02], repeated per sample.
1553        let ten = layout(Depth::Ten).neutral_chroma_plane();
1554        assert_eq!(ten, vec![0x00, 0x02, 0x00, 0x02, 0x00, 0x02, 0x00, 0x02]);
1555
1556        // 2048 little-endian is [0x00, 0x08].
1557        let twelve = layout(Depth::Twelve).neutral_chroma_plane();
1558        assert_eq!(twelve.len(), 8);
1559        assert_eq!(&twelve[0..2], &[0x00, 0x08]);
1560    }
1561
1562    #[test]
1563    fn black_luma_fill_is_zero_at_the_right_length() {
1564        assert_eq!(layout(Depth::Eight).black_luma_plane(), vec![0u8; 16]);
1565        assert_eq!(layout(Depth::Ten).black_luma_plane(), vec![0u8; 32]);
1566    }
1567}
1568
1569#[cfg(test)]
1570mod tests;