zencodec_io/
sinks.rs

1//! Sinks shared by drivers and tests.
2
3use alloc::format;
4use alloc::vec::Vec;
5
6use zencodec::decode::{DecodeRowSink, SinkError};
7use zenpixels::{PixelBuffer, PixelDescriptor, PixelSlice, PixelSliceMut};
8
9/// A sink that collects every final row into one tightly packed image.
10///
11/// `begin` sizes the buffer; `provide_next_buffer` hands out the rows in
12/// place, so the decoder writes straight into the final image. Final-row
13/// coverage is tracked separately: repeated or overlapping strips cannot hide
14/// missing rows at finish. The coverage bitmap uses eight bytes per 64 rows
15/// (rounded up), independent of image width. Supplying a buffer does not prove
16/// that the codec actually wrote every byte; native correctness remains required.
17#[derive(Default)]
18pub struct CollectSink {
19    data: Option<PixelBuffer>,
20    width: u32,
21    height: u32,
22    descriptor: Option<PixelDescriptor>,
23    rows_written: u64,
24    covered_rows: Vec<u64>,
25    unique_rows: u32,
26    finished: bool,
27}
28
29impl CollectSink {
30    /// An empty sink.
31    pub fn new() -> Self {
32        Self::default()
33    }
34
35    /// Output width, once `begin` has been called.
36    pub fn width(&self) -> u32 {
37        self.width
38    }
39
40    /// Output height, once `begin` has been called.
41    pub fn height(&self) -> u32 {
42        self.height
43    }
44
45    /// Output descriptor, once `begin` has been called.
46    pub fn descriptor(&self) -> Option<PixelDescriptor> {
47        self.descriptor
48    }
49
50    /// Rows written so far (a row written twice counts twice).
51    pub fn rows_written(&self) -> u64 {
52        self.rows_written
53    }
54
55    /// Whether `finish` has been called.
56    pub fn finished(&self) -> bool {
57        self.finished
58    }
59
60    /// The collected pixels, tightly packed.
61    pub fn bytes(&self) -> &[u8] {
62        self.data
63            .as_ref()
64            .map_or(&[], |buffer| buffer.as_slice().as_strided_bytes())
65    }
66
67    /// A zero-copy view of the collected image.
68    pub fn as_pixel_slice(&self) -> Option<PixelSlice<'_>> {
69        self.data.as_ref().map(PixelBuffer::as_slice)
70    }
71
72    /// Take the finished image without copying its pixels or allocating.
73    ///
74    /// Returns `None` before [`DecodeRowSink::finish`] succeeds. Call this only
75    /// after the driver returns success: a sink's finish is not proof that the
76    /// decoder's later validation succeeded. This does not include metadata;
77    /// retain the decoder's completion details separately.
78    pub fn into_buffer(self) -> Option<PixelBuffer> {
79        self.finished.then_some(self.data).flatten()
80    }
81
82    /// Copy the collected image into an owned [`PixelBuffer`].
83    pub fn to_buffer(&self) -> Option<PixelBuffer> {
84        let src = self.as_pixel_slice()?;
85        let mut out = PixelBuffer::new(self.width, self.height, src.descriptor());
86        {
87            let mut dst = out.as_slice_mut();
88            for y in 0..self.height {
89                let n = src.row(y).len().min(dst.row_mut(y).len());
90                dst.row_mut(y)[..n].copy_from_slice(&src.row(y)[..n]);
91            }
92        }
93        Some(out)
94    }
95}
96
97impl core::fmt::Debug for CollectSink {
98    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
99        f.debug_struct("CollectSink")
100            .field("width", &self.width)
101            .field("height", &self.height)
102            .field("descriptor", &self.descriptor)
103            .field("rows_written", &self.rows_written)
104            .field("finished", &self.finished)
105            .finish_non_exhaustive()
106    }
107}
108
109impl DecodeRowSink for CollectSink {
110    fn begin(
111        &mut self,
112        width: u32,
113        height: u32,
114        descriptor: PixelDescriptor,
115    ) -> Result<(), SinkError> {
116        if self.descriptor.is_some() {
117            return Err("CollectSink: begin called twice".into());
118        }
119        let stride = (width as usize)
120            .checked_mul(descriptor.bytes_per_pixel())
121            .ok_or_else(|| SinkError::from("CollectSink: row too large"))?;
122        let total = stride
123            .checked_mul(height as usize)
124            .ok_or_else(|| SinkError::from("CollectSink: image too large"))?;
125        // Retain the original reservation cause and leave state unchanged on
126        // failure. Prefix slack lets PixelBuffer align U16/F32 storage without
127        // copying; every row is written through that same aligned owner.
128        let capacity = total
129            .checked_add(descriptor.min_alignment() - 1)
130            .ok_or_else(|| SinkError::from("CollectSink: aligned image too large"))?;
131        let mut data = Vec::new();
132        data.try_reserve_exact(capacity)?;
133        data.resize(capacity, 0);
134        let buffer = PixelBuffer::from_vec(data, width, height, descriptor)?;
135        let words = usize::try_from(u64::from(height).div_ceil(64))
136            .map_err(|_| SinkError::from("CollectSink: coverage too large"))?;
137        let mut coverage = Vec::new();
138        coverage.try_reserve_exact(words)?;
139        coverage.resize(words, 0);
140        self.covered_rows = coverage;
141        self.data = Some(buffer);
142        self.width = width;
143        self.height = height;
144        self.descriptor = Some(descriptor);
145        Ok(())
146    }
147
148    fn provide_next_buffer(
149        &mut self,
150        y: u32,
151        height: u32,
152        width: u32,
153        descriptor: PixelDescriptor,
154    ) -> Result<PixelSliceMut<'_>, SinkError> {
155        if self.finished {
156            return Err("CollectSink: rows after finish".into());
157        }
158        let Some(desc) = self.descriptor else {
159            return Err("CollectSink: rows before begin".into());
160        };
161        if width != self.width || descriptor != desc {
162            return Err(format!(
163                "CollectSink: strip {width}x{height} {descriptor:?} does not match image {}x{} {desc:?}",
164                self.width, self.height
165            )
166            .into());
167        }
168        y.checked_add(height)
169            .filter(|end| *end <= self.height)
170            .ok_or_else(|| {
171                SinkError::from(format!(
172                    "CollectSink: rows {y}..{} exceed height {}",
173                    y as u64 + height as u64,
174                    self.height
175                ))
176            })?;
177        for row in y..y + height {
178            let word = &mut self.covered_rows[(row / 64) as usize];
179            let bit = 1u64 << (row % 64);
180            if *word & bit == 0 {
181                self.unique_rows += 1;
182                *word |= bit;
183            }
184        }
185        self.rows_written += height as u64;
186        Ok(self
187            .data
188            .as_mut()
189            .expect("begin installed pixel storage")
190            .rows_mut(y, height))
191    }
192
193    fn finish(&mut self) -> Result<(), SinkError> {
194        if self.descriptor.is_none() {
195            return Err("CollectSink: finish before begin".into());
196        }
197        if self.unique_rows != self.height {
198            return Err("CollectSink: finish before every row was supplied".into());
199        }
200        self.finished = true;
201        Ok(())
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    #[test]
210    fn finish_without_begin_is_not_completion() {
211        let mut sink = CollectSink::new();
212        assert!(sink.finish().is_err());
213        assert!(!sink.finished());
214        assert!(sink.into_buffer().is_none());
215    }
216
217    #[test]
218    fn completed_pixels_cannot_be_reopened_for_writing() {
219        let mut sink = CollectSink::new();
220        sink.begin(1, 1, PixelDescriptor::GRAY8_SRGB).unwrap();
221        sink.provide_next_buffer(0, 1, 1, PixelDescriptor::GRAY8_SRGB)
222            .unwrap();
223        sink.finish().unwrap();
224        assert!(
225            sink.provide_next_buffer(0, 1, 1, PixelDescriptor::GRAY8_SRGB)
226                .is_err()
227        );
228    }
229
230    #[test]
231    fn completion_requires_unique_coverage_and_closes_writes() {
232        let mut sink = CollectSink::new();
233        sink.begin(1, 65, PixelDescriptor::GRAY8_SRGB).unwrap();
234        for _ in 0..65 {
235            sink.provide_next_buffer(0, 1, 1, PixelDescriptor::GRAY8_SRGB)
236                .unwrap();
237        }
238        assert_eq!(sink.rows_written(), 65);
239        assert!(sink.finish().is_err());
240        assert!(!sink.finished());
241        // Out-of-order and overlapping delivery remain valid.
242        sink.provide_next_buffer(32, 33, 1, PixelDescriptor::GRAY8_SRGB)
243            .unwrap();
244        sink.provide_next_buffer(1, 32, 1, PixelDescriptor::GRAY8_SRGB)
245            .unwrap();
246        sink.finish().unwrap();
247        assert!(
248            sink.provide_next_buffer(0, 1, 1, PixelDescriptor::GRAY8_SRGB)
249                .is_err()
250        );
251        assert!(sink.into_buffer().is_some());
252    }
253
254    #[test]
255    fn finished_storage_moves_without_copy_for_integer_and_float_rows() {
256        for descriptor in [
257            PixelDescriptor::RGB8_SRGB,
258            PixelDescriptor::RGB16_SRGB,
259            PixelDescriptor::RGBAF32_LINEAR,
260        ] {
261            let mut sink = CollectSink::new();
262            sink.begin(3, 2, descriptor).unwrap();
263            let original = sink.bytes().as_ptr();
264            for y in 0..2 {
265                let mut rows = sink.provide_next_buffer(y, 1, 3, descriptor).unwrap();
266                rows.row_mut(0).fill(37 + y as u8);
267            }
268            sink.finish().unwrap();
269            let buffer = sink.into_buffer().unwrap();
270            assert_eq!(buffer.as_slice().as_strided_bytes().as_ptr(), original);
271            assert_eq!(buffer.descriptor(), descriptor);
272            for y in 0..2 {
273                assert!(buffer.as_slice().row(y).iter().all(|b| *b == 37 + y as u8));
274            }
275        }
276    }
277
278    #[test]
279    fn unfinished_collectors_do_not_yield_owned_completion() {
280        assert!(CollectSink::new().into_buffer().is_none());
281        let mut sink = CollectSink::new();
282        sink.begin(2, 1, PixelDescriptor::GRAY8_SRGB).unwrap();
283        sink.provide_next_buffer(0, 1, 2, PixelDescriptor::GRAY8_SRGB)
284            .unwrap()
285            .row_mut(0)
286            .copy_from_slice(&[17, 23]);
287        assert!(sink.into_buffer().is_none());
288    }
289
290    #[test]
291    fn collect_sink_refuses_unrepresentable_capacity_without_panic_or_state_change() {
292        let mut sink = CollectSink::new();
293        // On 64-bit this fits usize but exceeds Vec's isize capacity limit;
294        // on 32-bit the multiplication itself must refuse. Neither allocates.
295        assert!(
296            sink.begin(u32::MAX, u32::MAX, PixelDescriptor::GRAY8_SRGB)
297                .is_err()
298        );
299        assert_eq!(sink.descriptor(), None);
300        assert_eq!((sink.width(), sink.height()), (0, 0));
301        assert!(sink.bytes().is_empty());
302        assert!(!sink.finished());
303        sink.begin(3, 2, PixelDescriptor::GRAY8_SRGB).unwrap();
304        for y in 0..2 {
305            sink.provide_next_buffer(y, 1, 3, PixelDescriptor::GRAY8_SRGB)
306                .unwrap()
307                .row_mut(0)
308                .copy_from_slice(&[1 + y as u8, 3, 5]);
309        }
310        sink.finish().unwrap();
311        assert_eq!(sink.bytes(), &[1, 3, 5, 2, 3, 5]);
312        assert_eq!(sink.rows_written(), 2);
313        assert!(sink.finished());
314    }
315
316    #[test]
317    fn collect_sink_refuses_geometry_overflow() {
318        let mut sink = CollectSink::new();
319        assert!(
320            sink.begin(u32::MAX, u32::MAX, PixelDescriptor::RGBAF32_LINEAR)
321                .is_err()
322        );
323        assert!(sink.bytes().is_empty());
324        assert_eq!(sink.descriptor(), None);
325    }
326}
zencodec_io/sinks.rs

zencodec_io/
sinks.rs
