1use std::io::Write;
4
5use crate::{
6 ChromaSubsampling, ErrorCategory, ErrorCode, FormatDescriptor, FormatFamily, Frame,
7 PixelFlowError, Rational, Result, SampleType,
8};
9
10pub struct Y4mWriter<W: Write> {
12 writer: W,
13 format: FormatDescriptor,
14 width: usize,
15 height: usize,
16}
17
18impl<W: Write> Y4mWriter<W> {
19 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 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 #[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
103pub 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}