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