Skip to main content

pixelflow_core/
y4m.rs

1//! Y4M stream writer for Phase 1 compatible formats.
2
3use std::io::Write;
4
5use crate::{
6    ChromaSubsampling, ErrorCategory, ErrorCode, FormatDescriptor, FormatFamily, Frame,
7    PixelFlowError, Rational, Result, SampleType,
8};
9
10/// Streaming Y4M writer for one fixed output clip.
11pub struct Y4mWriter<W: Write> {
12    writer: W,
13    format: FormatDescriptor,
14    width: usize,
15    height: usize,
16}
17
18impl<W: Write> Y4mWriter<W> {
19    /// Creates writer and emits stream header immediately.
20    pub fn new(
21        mut writer: W,
22        format: FormatDescriptor,
23        width: usize,
24        height: usize,
25        frame_rate: Rational,
26    ) -> Result<Self> {
27        let chroma = y4m_chroma_tag(&format)?;
28        writeln!(
29            writer,
30            "YUV4MPEG2 W{width} H{height} F{}:{} Ip A0:0 {chroma}",
31            frame_rate.numerator, frame_rate.denominator
32        )
33        .map_err(|error| write_error(&error))?;
34
35        Ok(Self {
36            writer,
37            format,
38            width,
39            height,
40        })
41    }
42
43    /// Writes one frame header and visible plane rows in storage order.
44    pub fn write_frame(&mut self, frame: &Frame) -> Result<()> {
45        if frame.format() != &self.format
46            || frame.width() != self.width
47            || frame.height() != self.height
48        {
49            return Err(PixelFlowError::new(
50                ErrorCategory::Format,
51                ErrorCode::new("format.y4m_frame_mismatch"),
52                "frame properties do not match Y4M stream header",
53            ));
54        }
55
56        self.writer
57            .write_all(b"FRAME\n")
58            .map_err(|error| write_error(&error))?;
59        match self.format.sample_type() {
60            SampleType::U8 => self.write_u8_planes(frame),
61            SampleType::U16 => self.write_u16_planes(frame),
62            SampleType::F32 => Err(unsupported_format(
63                "format does not support Phase 1 Y4M float output",
64            )),
65        }
66    }
67
68    /// Consumes writer and returns wrapped output sink.
69    #[must_use]
70    pub fn into_inner(self) -> W {
71        self.writer
72    }
73
74    fn write_u8_planes(&mut self, frame: &Frame) -> Result<()> {
75        for plane_index in 0..self.format.planes().len() {
76            let plane = frame.plane::<u8>(plane_index)?;
77            for row in plane.rows() {
78                self.writer
79                    .write_all(row)
80                    .map_err(|error| write_error(&error))?;
81            }
82        }
83
84        Ok(())
85    }
86
87    fn write_u16_planes(&mut self, frame: &Frame) -> Result<()> {
88        for plane_index in 0..self.format.planes().len() {
89            let plane = frame.plane::<u16>(plane_index)?;
90            for row in plane.rows() {
91                for sample in row {
92                    self.writer
93                        .write_all(&sample.to_le_bytes())
94                        .map_err(|error| write_error(&error))?;
95                }
96            }
97        }
98
99        Ok(())
100    }
101}
102
103/// Returns Y4M chroma tag for one Phase 1 compatible format.
104pub fn y4m_chroma_tag(format: &FormatDescriptor) -> Result<&'static str> {
105    match (
106        format.family(),
107        format.subsampling(),
108        format.bits_per_sample(),
109        format.sample_type(),
110    ) {
111        (FormatFamily::Gray, None, 8, SampleType::U8) => Ok("Cmono"),
112        (FormatFamily::Gray, None, 10, SampleType::U16) => Ok("Cmono10"),
113        (FormatFamily::Gray, None, 12, SampleType::U16) => Ok("Cmono12"),
114        (FormatFamily::Gray, None, 16, SampleType::U16) => Ok("Cmono16"),
115        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs420), 8, SampleType::U8) => Ok("C420"),
116        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs422), 8, SampleType::U8) => Ok("C422"),
117        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs444), 8, SampleType::U8) => Ok("C444"),
118        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs420), 10, SampleType::U16) => Ok("C420p10"),
119        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs422), 10, SampleType::U16) => Ok("C422p10"),
120        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs444), 10, SampleType::U16) => Ok("C444p10"),
121        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs420), 12, SampleType::U16) => Ok("C420p12"),
122        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs422), 12, SampleType::U16) => Ok("C422p12"),
123        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs444), 12, SampleType::U16) => Ok("C444p12"),
124        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs420), 16, SampleType::U16) => Ok("C420p16"),
125        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs422), 16, SampleType::U16) => Ok("C422p16"),
126        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs444), 16, SampleType::U16) => Ok("C444p16"),
127        _ => Err(unsupported_format(format!(
128            "format '{}' cannot be written as Phase 1 Y4M",
129            format.name()
130        ))),
131    }
132}
133
134fn unsupported_format(message: impl Into<String>) -> PixelFlowError {
135    PixelFlowError::new(
136        ErrorCategory::Format,
137        ErrorCode::new("format.unsupported_y4m"),
138        message,
139    )
140}
141
142fn write_error(error: &std::io::Error) -> PixelFlowError {
143    PixelFlowError::new(
144        ErrorCategory::Io,
145        ErrorCode::new("io.write_y4m"),
146        format!("failed to write Y4M output: {error}"),
147    )
148}
149
150#[cfg(test)]
151mod tests {
152    #![expect(clippy::indexing_slicing, reason = "allow in tests")]
153
154    use crate::{AllocatorConfig, FrameBuilder, MetadataSchema, Rational, resolve_format_alias};
155
156    use super::{Y4mWriter, y4m_chroma_tag};
157
158    #[test]
159    fn y4m_chroma_tag_maps_phase1_formats() {
160        assert_eq!(
161            y4m_chroma_tag(&resolve_format_alias("gray8").expect("gray8 format should resolve"))
162                .expect("gray8 should map to Y4M"),
163            "Cmono"
164        );
165        assert_eq!(
166            y4m_chroma_tag(&resolve_format_alias("gray10").expect("gray10 format should resolve"))
167                .expect("gray10 should map to Y4M"),
168            "Cmono10"
169        );
170        assert_eq!(
171            y4m_chroma_tag(
172                &resolve_format_alias("yuv420p8").expect("yuv420p8 format should resolve")
173            )
174            .expect("yuv420p8 should map to Y4M"),
175            "C420"
176        );
177        assert_eq!(
178            y4m_chroma_tag(
179                &resolve_format_alias("yuv422p10").expect("yuv422p10 format should resolve")
180            )
181            .expect("yuv422p10 should map to Y4M"),
182            "C422p10"
183        );
184        assert_eq!(
185            y4m_chroma_tag(
186                &resolve_format_alias("yuv444p16").expect("yuv444p16 format should resolve")
187            )
188            .expect("yuv444p16 should map to Y4M"),
189            "C444p16"
190        );
191        assert!(
192            y4m_chroma_tag(&resolve_format_alias("rgbp8").expect("rgbp8 format should resolve"))
193                .is_err()
194        );
195    }
196
197    #[test]
198    fn y4m_writer_writes_header_and_u8_planes() {
199        let format = resolve_format_alias("gray8").expect("gray8 format should resolve");
200        let mut builder = FrameBuilder::new(
201            format.clone(),
202            2,
203            2,
204            &MetadataSchema::core(),
205            AllocatorConfig::default(),
206        )
207        .expect("frame builder should construct");
208        {
209            let mut plane = builder.plane_mut::<u8>(0).expect("plane should exist");
210            plane
211                .row_mut(0)
212                .expect("row 0 should exist")
213                .copy_from_slice(&[1, 2]);
214            plane
215                .row_mut(1)
216                .expect("row 1 should exist")
217                .copy_from_slice(&[3, 4]);
218        }
219        let frame = builder.finish();
220        let mut output = Vec::new();
221        let mut writer = Y4mWriter::new(
222            &mut output,
223            format,
224            2,
225            2,
226            Rational {
227                numerator: 24,
228                denominator: 1,
229            },
230        )
231        .expect("writer should construct");
232
233        writer
234            .write_frame(&frame)
235            .expect("frame should write successfully");
236
237        assert_eq!(
238            output,
239            b"YUV4MPEG2 W2 H2 F24:1 Ip A0:0 Cmono\nFRAME\n\x01\x02\x03\x04".to_vec()
240        );
241    }
242
243    #[test]
244    fn y4m_writer_writes_u16_samples_little_endian() {
245        let format = resolve_format_alias("gray10").expect("gray10 format should resolve");
246        let mut builder = FrameBuilder::new(
247            format.clone(),
248            1,
249            1,
250            &MetadataSchema::core(),
251            AllocatorConfig::default(),
252        )
253        .expect("frame builder should construct");
254        builder
255            .plane_mut::<u16>(0)
256            .expect("plane should exist")
257            .row_mut(0)
258            .expect("row 0 should exist")[0] = 0x0201;
259        let frame = builder.finish();
260        let mut output = Vec::new();
261        let mut writer = Y4mWriter::new(
262            &mut output,
263            format,
264            1,
265            1,
266            Rational {
267                numerator: 1,
268                denominator: 1,
269            },
270        )
271        .expect("writer should construct");
272
273        writer
274            .write_frame(&frame)
275            .expect("frame should write successfully");
276
277        assert_eq!(
278            output,
279            b"YUV4MPEG2 W1 H1 F1:1 Ip A0:0 Cmono10\nFRAME\n\x01\x02".to_vec()
280        );
281    }
282}