Remove some allocations from HDR image

src/codecs/hdr/decoder.rs

@@ -156,14 +156,7 @@ impl<R: BufRead> HdrAdapter<R> {
     fn read_image_data(&mut self, buf: &mut [u8]) -> ImageResult<()> {
         assert_eq!(u64::try_from(buf.len()), Ok(self.total_bytes()));
         match self.inner.take() {
-            Some(decoder) => {
-                let img: Vec<Rgb<u8>> = decoder.read_image_ldr()?;
-                for (i, Rgb(data)) in img.into_iter().enumerate() {
-                    buf[(i * 3)..][..3].copy_from_slice(&data);
-                }
-
-                Ok(())
-            }
+            Some(decoder) => decoder.read_image_ldr_buf(buf),
             None => Err(ImageError::Parameter(ParameterError::from_kind(
                 ParameterErrorKind::NoMoreData,
             ))),
@@ -247,6 +240,17 @@ pub struct HdrDecoder<R> {
     meta: HdrMetadata,
 }
 
+impl<R> HdrDecoder<R> {
+    /// The width of this HDR decoder
+    pub fn width(&self) -> u32 {
+        self.width
+    }
+    /// The height of this HDR decoder
+    pub fn height(&self) -> u32 {
+        self.height
+    }
+}
+
 /// Refer to [wikipedia](https://en.wikipedia.org/wiki/RGBE_image_format)
 #[repr(C)]
 #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
@@ -443,30 +447,33 @@ impl<R: BufRead> HdrDecoder<R> {
     }
 
     /// Consumes decoder and returns a vector of transformed pixels
-    pub fn read_image_transform<T: Send, F: Send + Sync + Fn(Rgbe8Pixel) -> T>(
+    pub fn read_image_transform<T, F: Send + Sync + Fn(Rgbe8Pixel) -> [T; N], const N: usize>(
         mut self,
         f: F,
         output_slice: &mut [T],
-    ) -> ImageResult<()> {
+    ) -> ImageResult<()>
+    where
+        T: Send + Copy,
+    {
         assert_eq!(
             output_slice.len(),
-            self.width as usize * self.height as usize
+            self.width as usize * self.height as usize * N
         );
 
         // Don't read anything if image is empty
         if self.width == 0 || self.height == 0 {
             return Ok(());
         }
 
-        let chunks_iter = output_slice.chunks_mut(self.width as usize);
+        let chunks_iter = output_slice.chunks_mut(self.width as usize * N);
 
         let mut buf = vec![Default::default(); self.width as usize];
         for chunk in chunks_iter {
             // read_scanline overwrites the entire buffer or returns an Err,
             // so not resetting the buffer here is ok.
             read_scanline(&mut self.r, &mut buf[..])?;
-            for (dst, &pix) in chunk.iter_mut().zip(buf.iter()) {
-                *dst = f(pix);
+            for (dst, &pix) in chunk.chunks_mut(N).zip(buf.iter()) {
+                dst.copy_from_slice(&f(pix));
             }
         }
         Ok(())
@@ -475,17 +482,31 @@ impl<R: BufRead> HdrDecoder<R> {
     /// Consumes decoder and returns a vector of `Rgb<u8>` pixels.
     /// scale = 1, gamma = 2.2
     pub fn read_image_ldr(self) -> ImageResult<Vec<Rgb<u8>>> {
-        let mut ret = vec![Rgb([0, 0, 0]); self.width as usize * self.height as usize];
-        self.read_image_transform(|pix| pix.to_ldr(), &mut ret[..])?;
-        Ok(ret)
+        let sz = (self.width * self.height) as usize;
+        let mut buf = vec![Rgb([0; 3]); sz];
+        self.read_image_ldr_buf(bytemuck::try_cast_slice_mut(buf.as_mut_slice()).unwrap())?;
+        Ok(buf)
+    }
+    fn read_image_ldr_buf(self, buf: &mut [u8]) -> ImageResult<()> {
+        let sz = (self.width * self.height * 3) as usize;
+        assert!(buf.len() >= sz);
+        self.read_image_transform(|pix| pix.to_ldr().0, &mut buf[..sz])
     }
-
     /// Consumes decoder and returns a vector of `Rgb<f32>` pixels.
     ///
     pub fn read_image_hdr(self) -> ImageResult<Vec<Rgb<f32>>> {
-        let mut ret = vec![Rgb([0.0, 0.0, 0.0]); self.width as usize * self.height as usize];
-        self.read_image_transform(|pix| pix.to_hdr(), &mut ret[..])?;
-        Ok(ret)
+        let sz = (self.width * self.height) as usize;
+        let mut buf = vec![Rgb([0.0f32; 3]); sz];
+        self.read_image_hdr_buf(bytemuck::try_cast_slice_mut(buf.as_mut_slice()).unwrap())?;
+        Ok(buf)
+    }
+
+    /// Consumes decoder and returns a vector of `Rgb<f32>` pixels.
+    ///
+    fn read_image_hdr_buf(self, buf: &mut [f32]) -> ImageResult<()> {
+        let sz = (self.width * self.height * 3) as usize;
+        assert!(buf.len() >= sz);
+        self.read_image_transform(|pix| pix.to_hdr().0, &mut buf[..sz])
     }
 }
 
@@ -971,15 +992,14 @@ fn split_at_first_test() {
 //   or return None to indicate end of file
 fn read_line_u8<R: BufRead>(r: &mut R) -> ::std::io::Result<Option<Vec<u8>>> {
     let mut ret = Vec::with_capacity(16);
-    match r.read_until(b'\n', &mut ret) {
-        Ok(0) => Ok(None),
-        Ok(_) => {
-            if let Some(&b'\n') = ret[..].last() {
+    match r.read_until(b'\n', &mut ret)? {
+        0 => Ok(None),
+        _ => {
+            if let Some(&b'\n') = ret.last() {
                 let _ = ret.pop();
             }
             Ok(Some(ret))
         }
-        Err(err) => Err(err),
     }
 }
 

src/codecs/openexr.rs

@@ -467,12 +467,11 @@ mod test {
                 .clone()
                 .join("overexposed gradient - data window equals display window.exr");
 
-            let hdr: Vec<Rgb<f32>> = crate::codecs::hdr::HdrDecoder::new(std::io::BufReader::new(
+            let hdr_decoder = crate::codecs::hdr::HdrDecoder::new(std::io::BufReader::new(
                 std::fs::File::open(&reference_path).unwrap(),
             ))
-            .unwrap()
-            .read_image_hdr()
             .unwrap();
+            let hdr = hdr_decoder.read_image_hdr().unwrap();
 
             let exr_pixels: Rgb32FImage = read_as_rgb_image_from_file(exr_path).unwrap();
             assert_eq!(

src/color.rs

@@ -350,6 +350,13 @@ impl<T> From<[T; $channels]> for $ident<T> {
     }
 }
 
+unsafe impl<T> bytemuck::Zeroable for $ident<T> where T: bytemuck::Zeroable {
+  fn zeroed() -> Self {
+    Self(<[T; $channels] as bytemuck::Zeroable>::zeroed())
+  }
+}
+unsafe impl<T> bytemuck::Pod for $ident<T> where T: bytemuck::Pod {}
+
 )* // END Structure definitions
 
     }