Skip to main content

ff_render/nodes/
upload.rs

1use super::RenderNodeCpu;
2
3/// YUV sub-sampling format for [`YuvUploadNode`].
4#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
5pub enum YuvFormat {
6    /// Planar 4:2:0 — Y at full resolution; Cb/Cr at half width and height.
7    #[default]
8    Yuv420p,
9    /// Planar 4:2:2 — Y at full resolution; Cb/Cr at half width.
10    Yuv422p,
11    /// Planar 4:4:4 — all planes at full resolution.
12    Yuv444p,
13}
14
15/// How the sample data for a [`YuvUploadNode`] is laid out in memory.
16///
17/// One enum rather than a pair of `bool`s (`high_bit_depth` + `semi_planar`)
18/// because the fourth combination — 8-bit semi-planar, i.e. NV12 — is not a
19/// layout this node supports, and a pair of flags would let it be constructed.
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21enum PlaneLayout {
22    /// Three planes, one byte per sample.
23    Planar8,
24    /// Three planes of little-endian `u16` samples holding `0..=1023` in the
25    /// low bits (`Yuv420p10le` and friends).
26    Planar10,
27    /// A luma plane plus one plane of interleaved Cb/Cr, little-endian `u16`
28    /// samples with the 10 significant bits in the *high* bits (`P010le`).
29    SemiPlanar10,
30}
31
32// Pipeline cache
33
34#[cfg(feature = "wgpu")]
35struct YuvPipeline {
36    render_pipeline: wgpu::RenderPipeline,
37    bind_group_layout: wgpu::BindGroupLayout,
38    y_tex: wgpu::Texture,
39    /// Cb for a planar layout; the interleaved Cb/Cr plane for a semi-planar one.
40    chroma_tex: wgpu::Texture,
41    /// Cr — planar layouts only; semi-planar carries it in `chroma_tex`.
42    cr_tex: Option<wgpu::Texture>,
43    uniform_buf: wgpu::Buffer,
44}
45
46// YuvUploadNode
47
48/// Upload raw YUV plane buffers to the GPU and convert to RGBA in a fragment
49/// shader, bypassing CPU-side `sws_scale`.
50///
51/// The node has `input_count() = 0`; it sources all pixel data from the plane
52/// buffers set via [`YuvUploadNode::set_planes`]. Call `set_planes` once per
53/// frame before the graph processes it.
54pub struct YuvUploadNode {
55    /// Pixel sub-sampling format.
56    pub format: YuvFormat,
57    /// Frame width in pixels.
58    pub width: u32,
59    /// Frame height in pixels.
60    pub height: u32,
61    /// Memory layout of the stored planes, fixed by the constructor used. The
62    /// 10-bit layouts render to an `Rgba16Float` target so precision survives.
63    layout: PlaneLayout,
64    y_plane: Vec<u8>,
65    /// Cb — planar layouts only.
66    cb_plane: Vec<u8>,
67    /// Cr — planar layouts only.
68    cr_plane: Vec<u8>,
69    /// Interleaved Cb/Cr — semi-planar layout only; empty otherwise.
70    uv_plane: Vec<u8>,
71    #[cfg(feature = "wgpu")]
72    pipeline: std::sync::OnceLock<YuvPipeline>,
73}
74
75/// Neutral chroma value for 10-bit YUV (mid of the `0..=1023` range).
76const TEN_BIT_NEUTRAL_CHROMA: u16 = 512;
77/// Maximum 10-bit sample value.
78const TEN_BIT_MAX: f32 = 1023.0;
79/// Bits P010 leaves zeroed at the bottom of each 16-bit sample. Its 10
80/// significant bits are MSB-aligned, matching `FFmpeg`'s own `P010LE` pixel
81/// descriptor (`depth = 10, shift = 6`).
82const P010_SHIFT: u32 = 6;
83
84impl YuvUploadNode {
85    /// Create a new 8-bit node. Plane buffers are initialised to neutral values (Y = 0, Cb = Cr = 128).
86    #[must_use]
87    pub fn new(format: YuvFormat, width: u32, height: u32) -> Self {
88        let (cw, ch) = chroma_dims(format, width, height);
89        Self {
90            format,
91            width,
92            height,
93            layout: PlaneLayout::Planar8,
94            y_plane: vec![0u8; (width * height) as usize],
95            cb_plane: vec![128u8; (cw * ch) as usize],
96            cr_plane: vec![128u8; (cw * ch) as usize],
97            uv_plane: Vec::new(),
98            #[cfg(feature = "wgpu")]
99            pipeline: std::sync::OnceLock::new(),
100        }
101    }
102
103    /// Create a new 10-bit planar node. Plane buffers hold little-endian `u16`
104    /// samples (values `0..=1023`) and render to an `Rgba16Float` target, so
105    /// 10-bit precision survives the upload. Neutral init: Y = 0, Cb = Cr = 512.
106    #[must_use]
107    pub fn new_high_bit_depth(format: YuvFormat, width: u32, height: u32) -> Self {
108        let (cw, ch) = chroma_dims(format, width, height);
109        Self {
110            format,
111            width,
112            height,
113            layout: PlaneLayout::Planar10,
114            y_plane: vec![0u8; (width * height * 2) as usize],
115            cb_plane: u16_le_plane(TEN_BIT_NEUTRAL_CHROMA, (cw * ch) as usize),
116            cr_plane: u16_le_plane(TEN_BIT_NEUTRAL_CHROMA, (cw * ch) as usize),
117            uv_plane: Vec::new(),
118            #[cfg(feature = "wgpu")]
119            pipeline: std::sync::OnceLock::new(),
120        }
121    }
122
123    /// Create a new 10-bit semi-planar (`P010le`) node: a full-resolution luma
124    /// plane plus one plane of interleaved Cb/Cr, both little-endian `u16`.
125    ///
126    /// P010 is always 4:2:0, so there is no [`YuvFormat`] to choose. Its samples
127    /// are MSB-aligned — the 10 significant bits sit in the *high* bits of each
128    /// 16-bit sample, with the low 6 zeroed — unlike
129    /// [`new_high_bit_depth`](Self::new_high_bit_depth), whose planes hold
130    /// `0..=1023` in the low bits. Renders to an `Rgba16Float` target.
131    ///
132    /// Neutral init: Y = 0, Cb = Cr = 512 (MSB-aligned).
133    #[must_use]
134    pub fn new_p010(width: u32, height: u32) -> Self {
135        let format = YuvFormat::Yuv420p;
136        let (cw, ch) = chroma_dims(format, width, height);
137        Self {
138            format,
139            width,
140            height,
141            layout: PlaneLayout::SemiPlanar10,
142            y_plane: vec![0u8; (width * height * 2) as usize],
143            cb_plane: Vec::new(),
144            cr_plane: Vec::new(),
145            // Two samples (Cb, Cr) per chroma pixel.
146            uv_plane: u16_le_plane(TEN_BIT_NEUTRAL_CHROMA << P010_SHIFT, (cw * ch * 2) as usize),
147            #[cfg(feature = "wgpu")]
148            pipeline: std::sync::OnceLock::new(),
149        }
150    }
151
152    /// Replace the stored plane buffers of a planar node ([`new`](Self::new) or
153    /// [`new_high_bit_depth`](Self::new_high_bit_depth)).
154    ///
155    /// A semi-planar node ([`new_p010`](Self::new_p010)) reads neither `cb` nor
156    /// `cr`, and its interleaved plane keeps the neutral chroma the constructor
157    /// gave it — which is correctly sized, so no length check catches the
158    /// mistake and the frame renders **greyscale** rather than failing. Use
159    /// [`set_planes_semi_planar`](Self::set_planes_semi_planar) there.
160    ///
161    /// Expected sizes for `width × height` at `format`, per sample:
162    /// - 8-bit: 1 byte; 10-bit ([`new_high_bit_depth`](Self::new_high_bit_depth)): 2 bytes (little-endian `u16`)
163    /// - `y`:       `width × height` samples
164    /// - `cb`, `cr`: `chroma_w × chroma_h` samples (sub-sampled per [`YuvFormat`])
165    pub fn set_planes(&mut self, y: Vec<u8>, cb: Vec<u8>, cr: Vec<u8>) {
166        self.y_plane = y;
167        self.cb_plane = cb;
168        self.cr_plane = cr;
169    }
170
171    /// Replace the stored planes of a semi-planar node
172    /// ([`new_p010`](Self::new_p010)).
173    ///
174    /// Both planes hold little-endian `u16` samples (2 bytes each):
175    /// - `y`: `width × height` samples
176    /// - `uv`: `chroma_w × chroma_h × 2` samples — Cb and Cr interleaved, one
177    ///   pair per chroma pixel, so `uv` holds twice as many samples as a single
178    ///   planar chroma plane would. This is the layout `ff-format` gives
179    ///   `P010le` (`uv_stride = width × 2` bytes over `height / 2` rows).
180    ///
181    /// Planes are expected dense (no row padding). A plane too short for the
182    /// node's dimensions is refused with a warning rather than read past its end.
183    ///
184    /// The mirror of the trap in [`set_planes`](Self::set_planes): a planar node
185    /// never reads `uv`, so calling this on one updates only the luma and leaves
186    /// the chroma at its neutral init, silently.
187    pub fn set_planes_semi_planar(&mut self, y: Vec<u8>, uv: Vec<u8>) {
188        self.y_plane = y;
189        self.uv_plane = uv;
190    }
191
192    /// `true` when the stored semi-planar planes are large enough for the node's
193    /// dimensions.
194    ///
195    /// Checked up front because a short plane would index past its end on the
196    /// CPU path and hand `write_texture` an undersized slice on the GPU one.
197    fn semi_planar_planes_are_complete(&self) -> bool {
198        let (cw, ch) = chroma_dims(self.format, self.width, self.height);
199        let luma_bytes = (self.width as usize) * (self.height as usize) * 2;
200        // Cb and Cr interleaved: two `u16` samples, so 4 bytes, per chroma pixel.
201        let uv_bytes = (cw as usize) * (ch as usize) * 4;
202        self.y_plane.len() >= luma_bytes && self.uv_plane.len() >= uv_bytes
203    }
204}
205
206/// Build a plane of `count` little-endian `u16` samples all equal to `value`.
207fn u16_le_plane(value: u16, count: usize) -> Vec<u8> {
208    value
209        .to_le_bytes()
210        .iter()
211        .copied()
212        .cycle()
213        .take(count * 2)
214        .collect()
215}
216
217impl Default for YuvUploadNode {
218    fn default() -> Self {
219        Self::new(YuvFormat::Yuv420p, 0, 0)
220    }
221}
222
223/// Returns `(chroma_width, chroma_height)` for a given format and luma dimensions.
224pub(crate) fn chroma_dims(format: YuvFormat, w: u32, h: u32) -> (u32, u32) {
225    match format {
226        YuvFormat::Yuv420p => (w.div_ceil(2), h.div_ceil(2)),
227        YuvFormat::Yuv422p => (w.div_ceil(2), h),
228        YuvFormat::Yuv444p => (w, h),
229    }
230}
231
232fn chroma_divs(format: YuvFormat) -> (u32, u32) {
233    match format {
234        YuvFormat::Yuv420p => (2, 2),
235        YuvFormat::Yuv422p => (2, 1),
236        YuvFormat::Yuv444p => (1, 1),
237    }
238}
239
240// CPU path
241
242impl RenderNodeCpu for YuvUploadNode {
243    fn process_cpu(&self, rgba: &mut [u8], w: u32, h: u32) {
244        if self.y_plane.is_empty() || self.width == 0 || self.height == 0 {
245            return;
246        }
247        match self.layout {
248            PlaneLayout::Planar8 => self.process_cpu_8bit(rgba, w, h),
249            PlaneLayout::Planar10 => self.process_cpu_10bit(rgba, w, h),
250            PlaneLayout::SemiPlanar10 => self.process_cpu_p010(rgba, w, h),
251        }
252    }
253}
254
255impl YuvUploadNode {
256    /// CPU YCbCr→RGBA for 8-bit planar input (1 byte per sample).
257    #[allow(
258        clippy::cast_possible_truncation,
259        clippy::cast_sign_loss,
260        clippy::many_single_char_names
261    )]
262    fn process_cpu_8bit(&self, rgba: &mut [u8], w: u32, h: u32) {
263        let (cw, _) = chroma_dims(self.format, self.width, self.height);
264        let (x_div, y_div) = chroma_divs(self.format);
265        let rows = h.min(self.height) as usize;
266        let cols = w.min(self.width) as usize;
267        for row in 0..rows {
268            for col in 0..cols {
269                let y_val = f32::from(self.y_plane[row * self.width as usize + col]) / 255.0;
270                let cx = col / x_div as usize;
271                let cy = row / y_div as usize;
272                let ci = cy * cw as usize + cx;
273                let cb = f32::from(self.cb_plane[ci]) / 255.0 - 0.5;
274                let cr = f32::from(self.cr_plane[ci]) / 255.0 - 0.5;
275                write_ycbcr_rgba(rgba, (row * w as usize + col) * 4, y_val, cb, cr);
276            }
277        }
278    }
279
280    /// CPU YCbCr→RGBA for 10-bit planar input (little-endian `u16` samples,
281    /// values `0..=1023`). The CPU fallback still writes 8-bit RGBA, so it loses
282    /// precision the GPU `Rgba16Float` path preserves.
283    #[allow(
284        clippy::cast_possible_truncation,
285        clippy::cast_sign_loss,
286        clippy::many_single_char_names
287    )]
288    fn process_cpu_10bit(&self, rgba: &mut [u8], w: u32, h: u32) {
289        let (cw, _) = chroma_dims(self.format, self.width, self.height);
290        let (x_div, y_div) = chroma_divs(self.format);
291        let rows = h.min(self.height) as usize;
292        let cols = w.min(self.width) as usize;
293        for row in 0..rows {
294            for col in 0..cols {
295                let y_val =
296                    sample_u16_le(&self.y_plane, row * self.width as usize + col) / TEN_BIT_MAX;
297                let cx = col / x_div as usize;
298                let cy = row / y_div as usize;
299                let ci = cy * cw as usize + cx;
300                let cb = sample_u16_le(&self.cb_plane, ci) / TEN_BIT_MAX - 0.5;
301                let cr = sample_u16_le(&self.cr_plane, ci) / TEN_BIT_MAX - 0.5;
302                write_ycbcr_rgba(rgba, (row * w as usize + col) * 4, y_val, cb, cr);
303            }
304        }
305    }
306
307    /// CPU YCbCr→RGBA for 10-bit semi-planar (`P010le`) input: a luma plane plus
308    /// one plane of interleaved Cb/Cr, both little-endian `u16` with the 10
309    /// significant bits MSB-aligned. Like the planar 10-bit leg, the CPU
310    /// fallback writes 8-bit RGBA and so loses precision the GPU `Rgba16Float`
311    /// path preserves.
312    #[allow(
313        clippy::cast_possible_truncation,
314        clippy::cast_sign_loss,
315        clippy::many_single_char_names
316    )]
317    fn process_cpu_p010(&self, rgba: &mut [u8], w: u32, h: u32) {
318        if !self.semi_planar_planes_are_complete() {
319            log::warn!(
320                "YuvUploadNode P010 planes too small for the frame: width={} height={} y_len={} uv_len={}",
321                self.width,
322                self.height,
323                self.y_plane.len(),
324                self.uv_plane.len()
325            );
326            return;
327        }
328        let (cw, _) = chroma_dims(self.format, self.width, self.height);
329        let (x_div, y_div) = chroma_divs(self.format);
330        let rows = h.min(self.height) as usize;
331        let cols = w.min(self.width) as usize;
332        for row in 0..rows {
333            for col in 0..cols {
334                let y_val = p010_norm(&self.y_plane, row * self.width as usize + col);
335                let cx = col / x_div as usize;
336                let cy = row / y_div as usize;
337                // Cb and Cr are adjacent samples of the same chroma pixel.
338                let ci = (cy * cw as usize + cx) * 2;
339                let cb = p010_norm(&self.uv_plane, ci) - 0.5;
340                let cr = p010_norm(&self.uv_plane, ci + 1) - 0.5;
341                write_ycbcr_rgba(rgba, (row * w as usize + col) * 4, y_val, cb, cr);
342            }
343        }
344    }
345}
346
347/// Read the `i`-th little-endian `u16` sample of a plane.
348fn raw_u16_le(plane: &[u8], i: usize) -> u16 {
349    u16::from_le_bytes([plane[i * 2], plane[i * 2 + 1]])
350}
351
352/// Read the `i`-th little-endian `u16` sample of a plane as `f32`.
353fn sample_u16_le(plane: &[u8], i: usize) -> f32 {
354    f32::from(raw_u16_le(plane, i))
355}
356
357/// Read the `i`-th MSB-aligned P010 sample of a plane, normalised to `[0, 1]`.
358///
359/// Dropping the zeroed low [`P010_SHIFT`] bits before dividing by
360/// [`TEN_BIT_MAX`] is what makes a P010 sample agree with the planar 10-bit
361/// path: both then divide the same `0..=1023` value.
362///
363/// Named for the normalisation, not for the read: [`sample_u16_le`] returns the
364/// raw value and leaves the divide to its caller, so a `_sample` twin here would
365/// invite dividing by [`TEN_BIT_MAX`] a second time. `p010_norm` is also what the
366/// shader calls its own copy of this.
367fn p010_norm(plane: &[u8], i: usize) -> f32 {
368    f32::from(raw_u16_le(plane, i) >> P010_SHIFT) / TEN_BIT_MAX
369}
370
371/// BT.601 full-range YCbCr → RGBA, writing 4 bytes at `idx` (alpha = 255).
372#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
373fn write_ycbcr_rgba(rgba: &mut [u8], idx: usize, y_val: f32, cb: f32, cr: f32) {
374    let r = (y_val + 1.402 * cr).clamp(0.0, 1.0);
375    let g = (y_val - 0.344 * cb - 0.714 * cr).clamp(0.0, 1.0);
376    let b = (y_val + 1.772 * cb).clamp(0.0, 1.0);
377    rgba[idx] = (r * 255.0 + 0.5) as u8;
378    rgba[idx + 1] = (g * 255.0 + 0.5) as u8;
379    rgba[idx + 2] = (b * 255.0 + 0.5) as u8;
380    rgba[idx + 3] = 255;
381}
382
383// GPU path
384
385#[cfg(feature = "wgpu")]
386impl YuvUploadNode {
387    #[allow(clippy::too_many_lines, clippy::similar_names)]
388    fn get_or_create_pipeline(&self, ctx: &crate::context::RenderContext) -> &YuvPipeline {
389        self.pipeline.get_or_init(|| {
390            let device = &ctx.device;
391            let (cw, ch) = chroma_dims(self.format, self.width, self.height);
392
393            // 10-bit planes are *Uint (raw samples) divided by max_value in the
394            // shader, rendering to Rgba16Float so precision survives. R16Uint and
395            // Rg16Uint are core formats (unlike R16Unorm, which needs an optional
396            // device feature). 8-bit planes stay R8Unorm sampled as normalised
397            // floats, rendering to Rgba8Unorm. The target format is threaded
398            // through rather than hardcoded, so an Rgba16Float graph gets a
399            // matching attachment.
400            let (luma_format, chroma_format, target_format, sample_type, shader_src) =
401                match self.layout {
402                    PlaneLayout::Planar8 => (
403                        wgpu::TextureFormat::R8Unorm,
404                        wgpu::TextureFormat::R8Unorm,
405                        wgpu::TextureFormat::Rgba8Unorm,
406                        wgpu::TextureSampleType::Float { filterable: false },
407                        include_str!("../shaders/yuv_upload.wgsl"),
408                    ),
409                    PlaneLayout::Planar10 => (
410                        wgpu::TextureFormat::R16Uint,
411                        wgpu::TextureFormat::R16Uint,
412                        wgpu::TextureFormat::Rgba16Float,
413                        wgpu::TextureSampleType::Uint,
414                        include_str!("../shaders/yuv_upload_10bit.wgsl"),
415                    ),
416                    PlaneLayout::SemiPlanar10 => (
417                        wgpu::TextureFormat::R16Uint,
418                        // Cb and Cr are the two channels of one texel.
419                        wgpu::TextureFormat::Rg16Uint,
420                        wgpu::TextureFormat::Rgba16Float,
421                        wgpu::TextureSampleType::Uint,
422                        include_str!("../shaders/p010_upload.wgsl"),
423                    ),
424                };
425            let planar = self.layout != PlaneLayout::SemiPlanar10;
426
427            let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
428                label: Some("YuvUpload shader"),
429                source: wgpu::ShaderSource::Wgsl(shader_src.into()),
430            });
431
432            let texture_entry = |binding: u32| wgpu::BindGroupLayoutEntry {
433                binding,
434                visibility: wgpu::ShaderStages::FRAGMENT,
435                ty: wgpu::BindingType::Texture {
436                    sample_type,
437                    view_dimension: wgpu::TextureViewDimension::D2,
438                    multisampled: false,
439                },
440                count: None,
441            };
442            // Semi-planar input carries Cb and Cr in one texture, so binding 2 is
443            // absent there. The uniform stays at binding 3 in every layout so all
444            // three upload shaders agree on where it is.
445            let mut entries = vec![texture_entry(0), texture_entry(1)];
446            if planar {
447                entries.push(texture_entry(2));
448            }
449            entries.push(wgpu::BindGroupLayoutEntry {
450                binding: 3,
451                visibility: wgpu::ShaderStages::FRAGMENT,
452                ty: wgpu::BindingType::Buffer {
453                    ty: wgpu::BufferBindingType::Uniform,
454                    has_dynamic_offset: false,
455                    min_binding_size: None,
456                },
457                count: None,
458            });
459
460            let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
461                label: Some("YuvUpload BGL"),
462                entries: &entries,
463            });
464
465            let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
466                label: Some("YuvUpload layout"),
467                bind_group_layouts: &[Some(&bgl)],
468                immediate_size: 0,
469            });
470
471            let render_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
472                label: Some("YuvUpload pipeline"),
473                layout: Some(&pipeline_layout),
474                vertex: wgpu::VertexState {
475                    module: &shader,
476                    entry_point: Some("vs_main"),
477                    buffers: &[],
478                    compilation_options: wgpu::PipelineCompilationOptions::default(),
479                },
480                fragment: Some(wgpu::FragmentState {
481                    module: &shader,
482                    entry_point: Some("fs_main"),
483                    targets: &[Some(wgpu::ColorTargetState {
484                        format: target_format,
485                        blend: None,
486                        write_mask: wgpu::ColorWrites::ALL,
487                    })],
488                    compilation_options: wgpu::PipelineCompilationOptions::default(),
489                }),
490                primitive: wgpu::PrimitiveState::default(),
491                depth_stencil: None,
492                multisample: wgpu::MultisampleState::default(),
493                multiview_mask: None,
494                cache: None,
495            });
496
497            let plane_tex = |label: &str, format: wgpu::TextureFormat, w: u32, h: u32| {
498                device.create_texture(&wgpu::TextureDescriptor {
499                    label: Some(label),
500                    size: wgpu::Extent3d {
501                        width: w,
502                        height: h,
503                        depth_or_array_layers: 1,
504                    },
505                    mip_level_count: 1,
506                    sample_count: 1,
507                    dimension: wgpu::TextureDimension::D2,
508                    format,
509                    usage: wgpu::TextureUsages::COPY_DST | wgpu::TextureUsages::TEXTURE_BINDING,
510                    view_formats: &[],
511                })
512            };
513
514            // Y luma plane (full resolution).
515            let y_tex = plane_tex("YuvUpload Y", luma_format, self.width, self.height);
516            // Cb, or the interleaved Cb/Cr plane for a semi-planar layout
517            // (sub-sampled either way).
518            let chroma_tex = plane_tex(
519                if planar {
520                    "YuvUpload Cb"
521                } else {
522                    "YuvUpload UV"
523                },
524                chroma_format,
525                cw,
526                ch,
527            );
528            // Cr chroma plane (sub-sampled) — planar layouts only.
529            let cr_tex = planar.then(|| plane_tex("YuvUpload Cr", chroma_format, cw, ch));
530
531            // Uniform buffer: [chroma_x_div, chroma_y_div, pad, pad] = 16 bytes.
532            let uniform_buf = device.create_buffer(&wgpu::BufferDescriptor {
533                label: Some("YuvUpload uniforms"),
534                size: 16,
535                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
536                mapped_at_creation: false,
537            });
538
539            YuvPipeline {
540                render_pipeline,
541                bind_group_layout: bgl,
542                y_tex,
543                chroma_tex,
544                cr_tex,
545                uniform_buf,
546            }
547        })
548    }
549}
550
551#[cfg(feature = "wgpu")]
552impl super::RenderNode for YuvUploadNode {
553    fn input_count(&self) -> usize {
554        0
555    }
556
557    #[allow(clippy::too_many_lines, clippy::similar_names)]
558    fn process(
559        &self,
560        _inputs: &[&wgpu::Texture],
561        outputs: &[&wgpu::Texture],
562        ctx: &crate::context::RenderContext,
563    ) {
564        if self.width == 0 || self.height == 0 || self.y_plane.is_empty() {
565            log::warn!("YuvUploadNode::process called with empty frame data");
566            return;
567        }
568        let Some(output) = outputs.first() else {
569            log::warn!("YuvUploadNode::process called with no outputs");
570            return;
571        };
572        if self.layout == PlaneLayout::SemiPlanar10 && !self.semi_planar_planes_are_complete() {
573            log::warn!(
574                "YuvUploadNode::process P010 planes too small for the frame: width={} height={} y_len={} uv_len={}",
575                self.width,
576                self.height,
577                self.y_plane.len(),
578                self.uv_plane.len()
579            );
580            return;
581        }
582
583        let pd = self.get_or_create_pipeline(ctx);
584        let (cw, ch) = chroma_dims(self.format, self.width, self.height);
585        let (x_div, y_div) = chroma_divs(self.format);
586        // Bytes per texel: R8Unorm = 1, R16Uint = 2, Rg16Uint = 4 (Cb and Cr in
587        // one texel).
588        let (luma_bpt, chroma_bpt) = match self.layout {
589            PlaneLayout::Planar8 => (1, 1),
590            PlaneLayout::Planar10 => (2, 2),
591            PlaneLayout::SemiPlanar10 => (2, 4),
592        };
593        let chroma_plane = if self.layout == PlaneLayout::SemiPlanar10 {
594            &self.uv_plane
595        } else {
596            &self.cb_plane
597        };
598
599        let upload = |tex: &wgpu::Texture, data: &[u8], w: u32, h: u32, bpt: u32| {
600            ctx.queue.write_texture(
601                wgpu::TexelCopyTextureInfo {
602                    texture: tex,
603                    mip_level: 0,
604                    origin: wgpu::Origin3d::ZERO,
605                    aspect: wgpu::TextureAspect::All,
606                },
607                data,
608                wgpu::TexelCopyBufferLayout {
609                    offset: 0,
610                    bytes_per_row: Some(w * bpt),
611                    rows_per_image: None,
612                },
613                wgpu::Extent3d {
614                    width: w,
615                    height: h,
616                    depth_or_array_layers: 1,
617                },
618            );
619        };
620
621        upload(&pd.y_tex, &self.y_plane, self.width, self.height, luma_bpt);
622        upload(&pd.chroma_tex, chroma_plane, cw, ch, chroma_bpt);
623        // Semi-planar input has no separate Cr plane; its texture is not created.
624        if let Some(cr_tex) = pd.cr_tex.as_ref() {
625            upload(cr_tex, &self.cr_plane, cw, ch, chroma_bpt);
626        }
627
628        // Uniforms: [chroma_x_div: u32, chroma_y_div: u32, max_value: f32, pad].
629        // max_value is the 10-bit shaders' normalisation divisor (1023) — P010
630        // shifts its MSB-aligned samples down first, so the divisor is the same;
631        // the 8-bit shader ignores this slot (its samples are pre-normalised).
632        let mut uniforms = [0u8; 16];
633        uniforms[0..4].copy_from_slice(&x_div.to_le_bytes());
634        uniforms[4..8].copy_from_slice(&y_div.to_le_bytes());
635        uniforms[8..12].copy_from_slice(&TEN_BIT_MAX.to_le_bytes());
636        ctx.queue.write_buffer(&pd.uniform_buf, 0, &uniforms);
637
638        let y_view = pd
639            .y_tex
640            .create_view(&wgpu::TextureViewDescriptor::default());
641        let chroma_view = pd
642            .chroma_tex
643            .create_view(&wgpu::TextureViewDescriptor::default());
644        let cr_view = pd
645            .cr_tex
646            .as_ref()
647            .map(|tex| tex.create_view(&wgpu::TextureViewDescriptor::default()));
648        let out_view = output.create_view(&wgpu::TextureViewDescriptor::default());
649
650        // Mirrors the layout built in `get_or_create_pipeline`: binding 2 exists
651        // only for a planar layout, the uniform is always at binding 3.
652        let mut bg_entries = vec![
653            wgpu::BindGroupEntry {
654                binding: 0,
655                resource: wgpu::BindingResource::TextureView(&y_view),
656            },
657            wgpu::BindGroupEntry {
658                binding: 1,
659                resource: wgpu::BindingResource::TextureView(&chroma_view),
660            },
661        ];
662        if let Some(cr_view) = cr_view.as_ref() {
663            bg_entries.push(wgpu::BindGroupEntry {
664                binding: 2,
665                resource: wgpu::BindingResource::TextureView(cr_view),
666            });
667        }
668        bg_entries.push(wgpu::BindGroupEntry {
669            binding: 3,
670            resource: pd.uniform_buf.as_entire_binding(),
671        });
672
673        let bind_group = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
674            label: Some("YuvUpload BG"),
675            layout: &pd.bind_group_layout,
676            entries: &bg_entries,
677        });
678
679        let mut encoder = ctx
680            .device
681            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
682                label: Some("YuvUpload pass"),
683            });
684        {
685            let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
686                label: Some("YuvUpload pass"),
687                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
688                    view: &out_view,
689                    resolve_target: None,
690                    depth_slice: None,
691                    ops: wgpu::Operations {
692                        load: wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),
693                        store: wgpu::StoreOp::Store,
694                    },
695                })],
696                depth_stencil_attachment: None,
697                timestamp_writes: None,
698                occlusion_query_set: None,
699                multiview_mask: None,
700            });
701            pass.set_pipeline(&pd.render_pipeline);
702            pass.set_bind_group(0, &bind_group, &[]);
703            pass.draw(0..6, 0..1);
704        }
705        ctx.queue.submit(std::iter::once(encoder.finish()));
706    }
707}
708
709#[cfg(test)]
710mod tests {
711    use super::*;
712
713    #[test]
714    fn yuv_format_default_should_be_yuv420p() {
715        assert_eq!(YuvFormat::default(), YuvFormat::Yuv420p);
716    }
717
718    #[test]
719    fn chroma_dims_420p_should_halve_both_dimensions() {
720        assert_eq!(chroma_dims(YuvFormat::Yuv420p, 4, 4), (2, 2));
721        // Odd dimensions: ceiling division.
722        assert_eq!(chroma_dims(YuvFormat::Yuv420p, 3, 3), (2, 2));
723    }
724
725    #[test]
726    fn chroma_dims_422p_should_halve_width_only() {
727        assert_eq!(chroma_dims(YuvFormat::Yuv422p, 4, 4), (2, 4));
728        assert_eq!(chroma_dims(YuvFormat::Yuv422p, 3, 5), (2, 5));
729    }
730
731    #[test]
732    fn chroma_dims_444p_should_be_full_resolution() {
733        assert_eq!(chroma_dims(YuvFormat::Yuv444p, 4, 6), (4, 6));
734    }
735
736    #[test]
737    fn yuv_upload_node_cpu_black_frame_should_produce_black() {
738        let mut node = YuvUploadNode::new(YuvFormat::Yuv420p, 2, 2);
739        node.set_planes(
740            vec![0u8; 4],   // Y = 0
741            vec![128u8; 1], // Cb = neutral
742            vec![128u8; 1], // Cr = neutral
743        );
744        let mut rgba = vec![0u8; 16];
745        node.process_cpu(&mut rgba, 2, 2);
746        for pixel in rgba.chunks_exact(4) {
747            assert!(pixel[0] <= 1, "R should be ~0 for Y=0; got {}", pixel[0]);
748            assert!(pixel[1] <= 1, "G should be ~0 for Y=0; got {}", pixel[1]);
749            assert!(pixel[2] <= 1, "B should be ~0 for Y=0; got {}", pixel[2]);
750            assert_eq!(pixel[3], 255, "alpha must be opaque");
751        }
752    }
753
754    #[test]
755    fn yuv_upload_node_cpu_white_frame_should_produce_white() {
756        let mut node = YuvUploadNode::new(YuvFormat::Yuv420p, 2, 2);
757        node.set_planes(
758            vec![255u8; 4], // Y = 255
759            vec![128u8; 1], // Cb = neutral
760            vec![128u8; 1], // Cr = neutral
761        );
762        let mut rgba = vec![0u8; 16];
763        node.process_cpu(&mut rgba, 2, 2);
764        for pixel in rgba.chunks_exact(4) {
765            assert!(
766                pixel[0] >= 254,
767                "R should be ~255 for Y=255, neutral chroma; got {}",
768                pixel[0]
769            );
770            assert!(
771                pixel[1] >= 254,
772                "G should be ~255 for Y=255, neutral chroma; got {}",
773                pixel[1]
774            );
775            assert!(
776                pixel[2] >= 254,
777                "B should be ~255 for Y=255, neutral chroma; got {}",
778                pixel[2]
779            );
780        }
781    }
782
783    #[test]
784    fn yuv_upload_node_cpu_neutral_chroma_should_produce_grey() {
785        let mut node = YuvUploadNode::new(YuvFormat::Yuv420p, 2, 2);
786        // Y=128 → y_val ≈ 0.502, Cb=Cr=128 → cb=cr=0 → R=G=B ≈ 128.
787        node.set_planes(vec![128u8; 4], vec![128u8; 1], vec![128u8; 1]);
788        let mut rgba = vec![0u8; 16];
789        node.process_cpu(&mut rgba, 2, 2);
790        for pixel in rgba.chunks_exact(4) {
791            let r = pixel[0] as i32;
792            let g = pixel[1] as i32;
793            let b = pixel[2] as i32;
794            assert!(
795                (r - 128).abs() <= 2,
796                "R should be ~128 for neutral YUV; got {r}"
797            );
798            assert!(
799                (g - 128).abs() <= 2,
800                "G should be ~128 for neutral YUV; got {g}"
801            );
802            assert!(
803                (b - 128).abs() <= 2,
804                "B should be ~128 for neutral YUV; got {b}"
805            );
806        }
807    }
808
809    #[test]
810    fn yuv_upload_node_cpu_422p_should_use_half_width_chroma() {
811        // 4×2 frame, 422p: chroma planes are 2×2.
812        let mut node = YuvUploadNode::new(YuvFormat::Yuv422p, 4, 2);
813        node.set_planes(
814            vec![128u8; 8], // 4×2 luma — neutral grey
815            vec![128u8; 4], // 2×2 Cb
816            vec![128u8; 4], // 2×2 Cr
817        );
818        let mut rgba = vec![0u8; 32];
819        node.process_cpu(&mut rgba, 4, 2);
820        for pixel in rgba.chunks_exact(4) {
821            let r = pixel[0] as i32;
822            assert!(
823                (r - 128).abs() <= 2,
824                "422p neutral: R should be ~128; got {r}"
825            );
826        }
827    }
828
829    #[test]
830    fn yuv_upload_node_set_planes_should_update_stored_data() {
831        let mut node = YuvUploadNode::new(YuvFormat::Yuv444p, 1, 1);
832        // Default: Y=0, Cb=Cr=128 → near-black (128/255 ≈ 0.502, not exact 0.5).
833        let mut rgba = vec![0u8; 4];
834        node.process_cpu(&mut rgba, 1, 1);
835        assert!(
836            rgba[0] <= 2,
837            "default Y=0 must produce near-black; got {}",
838            rgba[0]
839        );
840        // After set_planes: Y=200, Cb=Cr=128 → bright grey.
841        node.set_planes(vec![200], vec![128], vec![128]);
842        node.process_cpu(&mut rgba, 1, 1);
843        assert!(
844            rgba[0] > 150,
845            "Y=200 must produce bright output; got {}",
846            rgba[0]
847        );
848    }
849
850    #[test]
851    fn yuv_upload_cpu_10bit_should_decode_u16_planes() {
852        // Y = 768 (10-bit) → 768/1023 ≈ 0.751 → grey ≈ 191. Neutral chroma 512.
853        // A byte-truncating misread of the little-endian u16 (low byte 0x00)
854        // would yield 0, so asserting ~191 proves the u16 decode (non-vacuous).
855        let mut node = YuvUploadNode::new_high_bit_depth(YuvFormat::Yuv420p, 2, 2);
856        node.set_planes(
857            u16_le_plane(768, 4),
858            u16_le_plane(512, 1),
859            u16_le_plane(512, 1),
860        );
861        let mut rgba = vec![0u8; 16];
862        node.process_cpu(&mut rgba, 2, 2);
863        for pixel in rgba.chunks_exact(4) {
864            let r = i32::from(pixel[0]);
865            assert!(
866                (r - 191).abs() <= 3,
867                "10-bit Y=768 must decode to ~191; got {r}"
868            );
869            assert_eq!(pixel[3], 255, "alpha must be opaque");
870        }
871    }
872
873    /// A plane of the given little-endian `u16` samples.
874    fn u16_le_samples(values: &[u16]) -> Vec<u8> {
875        values.iter().flat_map(|v| v.to_le_bytes()).collect()
876    }
877
878    /// The same samples MSB-aligned the way P010 stores them.
879    fn p010_samples(values: &[u16]) -> Vec<u8> {
880        let shifted: Vec<u16> = values.iter().map(|v| v << P010_SHIFT).collect();
881        u16_le_samples(&shifted)
882    }
883
884    #[test]
885    fn yuv_upload_cpu_p010_should_decode_msb_aligned_samples() {
886        // Y = 768 once the 6-bit shift is undone → 768/1023 ≈ 0.751 → grey ≈ 191.
887        // Non-vacuous: reading the sample without shifting divides 49152 by 1023
888        // and clamps to white (255), which is nowhere near 191.
889        let mut node = YuvUploadNode::new_p010(2, 2);
890        node.set_planes_semi_planar(p010_samples(&[768; 4]), p010_samples(&[512; 2]));
891        let mut rgba = vec![0u8; 16];
892        node.process_cpu(&mut rgba, 2, 2);
893        for pixel in rgba.chunks_exact(4) {
894            let r = i32::from(pixel[0]);
895            assert!(
896                (r - 191).abs() <= 3,
897                "P010 Y=768 must decode to ~191; got {r} (255 means the shift was skipped)"
898            );
899            assert_eq!(pixel[3], 255, "alpha must be opaque");
900        }
901    }
902
903    #[test]
904    fn yuv_upload_cpu_p010_should_deinterleave_cb_and_cr() {
905        // 4×2 at 4:2:0 → a 2×1 chroma plane, so the UV plane holds two pairs:
906        // [Cb0, Cr0, Cb1, Cr1]. Giving the two chroma columns opposite Cb/Cr
907        // makes a swapped or mis-strided read visible: the channel that moves
908        // changes. A single chroma column, or w == h, would hide both.
909        let mut node = YuvUploadNode::new_p010(4, 2);
910        node.set_planes_semi_planar(p010_samples(&[512; 8]), p010_samples(&[512, 800, 800, 512]));
911        let mut rgba = vec![0u8; 32];
912        node.process_cpu(&mut rgba, 4, 2);
913
914        let red_and_blue = |x: usize, y: usize| -> (i32, i32) {
915            let i = (y * 4 + x) * 4;
916            (i32::from(rgba[i]), i32::from(rgba[i + 2]))
917        };
918        // Chroma column 0 covers x = 0..2: Cr is high, so red rises and blue
919        // stays neutral.
920        for x in [0, 1] {
921            for y in [0, 1] {
922                let (r, b) = red_and_blue(x, y);
923                assert!(r > 200, "Cr=800 must push R high at ({x},{y}); got {r}");
924                assert!(b < 160, "Cb=512 must leave B neutral at ({x},{y}); got {b}");
925            }
926        }
927        // Chroma column 1 covers x = 2..4, with the roles swapped.
928        for x in [2, 3] {
929            for y in [0, 1] {
930                let (r, b) = red_and_blue(x, y);
931                assert!(r < 160, "Cr=512 must leave R neutral at ({x},{y}); got {r}");
932                assert!(b > 200, "Cb=800 must push B high at ({x},{y}); got {b}");
933            }
934        }
935    }
936
937    #[test]
938    fn yuv_upload_cpu_p010_should_match_planar_10bit_for_the_same_samples() {
939        // The strongest pin on the shift, and unlike the GPU tests it runs
940        // everywhere: P010 fed `v << 6` must land on exactly the pixels the
941        // already-verified planar 10-bit path produces from `v`. Samples vary per
942        // pixel and the two chroma columns differ, so a transposed or constant
943        // read cannot pass.
944        const Y: [u16; 8] = [100, 300, 500, 700, 900, 200, 400, 600];
945        const CB: [u16; 2] = [300, 700];
946        const CR: [u16; 2] = [800, 200];
947
948        let mut planar = YuvUploadNode::new_high_bit_depth(YuvFormat::Yuv420p, 4, 2);
949        planar.set_planes(u16_le_samples(&Y), u16_le_samples(&CB), u16_le_samples(&CR));
950        let mut expected = vec![0u8; 32];
951        planar.process_cpu(&mut expected, 4, 2);
952
953        let mut p010 = YuvUploadNode::new_p010(4, 2);
954        p010.set_planes_semi_planar(
955            p010_samples(&Y),
956            p010_samples(&[CB[0], CR[0], CB[1], CR[1]]),
957        );
958        let mut got = vec![0u8; 32];
959        p010.process_cpu(&mut got, 4, 2);
960
961        assert_eq!(
962            got, expected,
963            "P010 must decode to the same pixels as the planar 10-bit path"
964        );
965        // Non-vacuous: a flat frame on both sides would satisfy the comparison
966        // without exercising anything, so check the fixture actually varies.
967        assert!(
968            expected.chunks_exact(4).any(|p| p[0] != expected[0]),
969            "the fixture must produce varying pixels"
970        );
971    }
972
973    #[test]
974    fn yuv_upload_cpu_p010_should_refuse_planes_too_small_for_the_frame() {
975        // 4×2 needs 4 UV samples; two must be refused rather than read past the
976        // end of the plane.
977        let mut node = YuvUploadNode::new_p010(4, 2);
978        node.set_planes_semi_planar(p010_samples(&[512; 8]), p010_samples(&[512, 800]));
979        let mut rgba = vec![7u8; 32];
980        node.process_cpu(&mut rgba, 4, 2);
981        assert!(
982            rgba.iter().all(|&b| b == 7),
983            "an undersized UV plane must leave the output untouched"
984        );
985    }
986
987    #[test]
988    fn yuv_upload_node_variant_and_error_types_should_compile() {
989        let _ = YuvFormat::Yuv420p;
990        let _ = YuvFormat::Yuv422p;
991        let _ = YuvFormat::Yuv444p;
992        let _ = YuvUploadNode::new(YuvFormat::Yuv420p, 320, 240);
993        let _ = YuvUploadNode::default();
994    }
995}