Less buffer copies

CHANGELOG.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Removed
+- `jxl-oxide`: Remove `Render::image` (#334). Use `Render::stream` instead.
+
 ### Fixed
 - `jxl-modular`: Fix incorrect color with complex inverse palette (#312).
 

crates/jxl-color/src/convert.rs

@@ -524,6 +524,22 @@ impl ColorTransform {
         self.ops.is_empty()
     }
 
+    #[inline]
+    pub fn input_channels(&self) -> usize {
+        self.begin_channels
+    }
+
+    #[inline]
+    pub fn output_channels(&self) -> usize {
+        let mut channels = self.begin_channels;
+        for op in &self.ops {
+            if let Some(x) = op.outputs() {
+                channels = x;
+            }
+        }
+        channels
+    }
+
     /// Performs the prepared color transformation on the samples.
     ///
     /// Returns the number of final channels after transformation.

crates/jxl-oxide-cli/src/decode.rs

@@ -99,7 +99,7 @@ pub fn handle_decode(args: DecodeArgs) -> Result<()> {
     let mps = total_pixels as f64 / 1e6;
 
     if args.output_format == OutputFormat::Npy {
-        image.set_render_spot_colour(false);
+        image.set_render_spot_color(false);
     }
 
     let keyframes = if let Some(num_reps @ 2..) = args.num_reps {
@@ -379,7 +379,7 @@ fn write_npy<W: Write>(
 ) -> std::io::Result<()> {
     let channels = {
         let first_frame = keyframes.first().unwrap();
-        first_frame.color_channels().len() + first_frame.extra_channels().len()
+        first_frame.color_channels().len() + first_frame.extra_channels().0.len()
     };
 
     let mut output = std::io::BufWriter::new(output);

crates/jxl-oxide-wasm/src/lib.rs

@@ -254,10 +254,16 @@ impl RenderResult {
     #[wasm_bindgen(js_name = encodeToPng)]
     pub fn into_png(self) -> Result<Vec<u8>, String> {
         let image = self.image;
+        let mut stream = image.stream();
+        let mut fb = jxl_oxide::FrameBuffer::new(
+            stream.width() as usize,
+            stream.height() as usize,
+            stream.channels() as usize,
+        );
+        stream.write_to_buffer(fb.buf_mut());
 
-        let fb = image.image();
         let mut out = Vec::new();
-        let mut encoder = png::Encoder::new(&mut out, fb.width() as u32, fb.height() as u32);
+        let mut encoder = png::Encoder::new(&mut out, stream.width(), stream.height());
         let color = match self.pixfmt {
             PixelFormat::Gray => png::ColorType::Grayscale,
             PixelFormat::Graya => png::ColorType::GrayscaleAlpha,

crates/jxl-oxide/src/fb.rs

@@ -166,3 +166,244 @@ impl FrameBuffer {
         }
     }
 }
+
+/// Image stream that writes to borrowed buffer.
+pub struct ImageStream<'r> {
+    orientation: u32,
+    width: u32,
+    height: u32,
+    grids: Vec<&'r ImageBuffer>,
+    start_offset_xy: Vec<(i32, i32)>,
+    bit_depth: Vec<BitDepth>,
+    spot_colors: Vec<ImageStreamSpotColor<'r>>,
+    y: u32,
+    x: u32,
+    c: u32,
+}
+
+impl<'r> ImageStream<'r> {
+    pub(crate) fn from_render(render: &'r crate::Render) -> Self {
+        use jxl_image::ExtraChannelType;
+
+        let orientation = render.orientation;
+        assert!((1..=8).contains(&orientation));
+        let Region {
+            left,
+            top,
+            mut width,
+            mut height,
+        } = render.target_frame_region;
+        if orientation >= 5 {
+            std::mem::swap(&mut width, &mut height);
+        }
+
+        let fb = render.image.buffer();
+        let color_channels = render.image.color_channels();
+        let regions_and_shifts = render.image.regions_and_shifts();
+
+        let mut grids: Vec<_> = render.color_channels().iter().collect();
+        let mut bit_depth = vec![render.color_bit_depth; grids.len()];
+
+        let mut start_offset_xy = Vec::new();
+        for (region, _) in &regions_and_shifts[..color_channels] {
+            start_offset_xy.push((left - region.left, top - region.top));
+        }
+
+        // Find black
+        for (ec_idx, (ec, (region, _))) in render
+            .extra_channels
+            .iter()
+            .zip(&regions_and_shifts[color_channels..])
+            .enumerate()
+        {
+            if ec.is_black() {
+                grids.push(&fb[color_channels + ec_idx]);
+                bit_depth.push(ec.bit_depth);
+                start_offset_xy.push((left - region.left, top - region.top));
+                break;
+            }
+        }
+        // Find alpha
+        for (ec_idx, (ec, (region, _))) in render
+            .extra_channels
+            .iter()
+            .zip(&regions_and_shifts[color_channels..])
+            .enumerate()
+        {
+            if ec.is_alpha() {
+                grids.push(&fb[color_channels + ec_idx]);
+                bit_depth.push(ec.bit_depth);
+                start_offset_xy.push((left - region.left, top - region.top));
+                break;
+            }
+        }
+
+        let mut spot_colors = Vec::new();
+        if render.render_spot_color && color_channels == 3 {
+            for (ec_idx, (ec, (region, _))) in render
+                .extra_channels
+                .iter()
+                .zip(&regions_and_shifts[color_channels..])
+                .enumerate()
+            {
+                if let ExtraChannelType::SpotColour {
+                    red,
+                    green,
+                    blue,
+                    solidity,
+                } = ec.ty
+                {
+                    let grid = &fb[color_channels + ec_idx];
+                    let xy = (left - region.left, top - region.top);
+                    spot_colors.push(ImageStreamSpotColor {
+                        grid,
+                        start_offset_xy: xy,
+                        bit_depth: ec.bit_depth,
+                        rgb: (red, green, blue),
+                        solidity,
+                    });
+                }
+            }
+        }
+
+        ImageStream {
+            orientation,
+            width,
+            height,
+            grids,
+            bit_depth,
+            start_offset_xy,
+            spot_colors,
+            y: 0,
+            x: 0,
+            c: 0,
+        }
+    }
+}
+
+impl ImageStream<'_> {
+    /// Returns width of the image.
+    #[inline]
+    pub fn width(&self) -> u32 {
+        self.width
+    }
+
+    /// Returns height of the image.
+    #[inline]
+    pub fn height(&self) -> u32 {
+        self.height
+    }
+
+    /// Returns the number of channels of the image.
+    #[inline]
+    pub fn channels(&self) -> u32 {
+        self.grids.len() as u32
+    }
+
+    /// Writes next samples to the buffer, returning how many samples are written.
+    pub fn write_to_buffer(&mut self, buf: &mut [f32]) -> usize {
+        let channels = self.grids.len() as u32;
+        let mut buf_it = buf.iter_mut();
+        let mut count = 0usize;
+        'outer: while self.y < self.height {
+            while self.x < self.width {
+                while self.c < channels {
+                    let Some(v) = buf_it.next() else {
+                        break 'outer;
+                    };
+                    let (start_x, start_y) = self.start_offset_xy[self.c as usize];
+                    let (orig_x, orig_y) = self.to_original_coord(self.x, self.y);
+                    let (Some(x), Some(y)) = (
+                        orig_x.checked_add_signed(start_x),
+                        orig_y.checked_add_signed(start_y),
+                    ) else {
+                        *v = 0.0;
+                        count += 1;
+                        self.c += 1;
+                        continue;
+                    };
+                    let x = x as usize;
+                    let y = y as usize;
+                    let grid = &self.grids[self.c as usize];
+                    let bit_depth = self.bit_depth[self.c as usize];
+                    *v = match grid {
+                        ImageBuffer::F32(g) => g.get(x, y).copied().unwrap_or(0.0),
+                        ImageBuffer::I32(g) => {
+                            bit_depth.parse_integer_sample(g.get(x, y).copied().unwrap_or(0))
+                        }
+                        ImageBuffer::I16(g) => {
+                            bit_depth.parse_integer_sample(g.get(x, y).copied().unwrap_or(0) as i32)
+                        }
+                    };
+
+                    if self.c < 3 {
+                        for spot in &self.spot_colors {
+                            let ImageStreamSpotColor {
+                                grid,
+                                start_offset_xy: (start_x, start_y),
+                                bit_depth,
+                                rgb: (r, g, b),
+                                solidity,
+                            } = *spot;
+                            let color = [r, g, b][self.c as usize];
+                            let xy = (
+                                orig_x.checked_add_signed(start_x),
+                                orig_y.checked_add_signed(start_y),
+                            );
+                            let mix = if let (Some(x), Some(y)) = xy {
+                                let x = x as usize;
+                                let y = y as usize;
+                                let val = match grid {
+                                    ImageBuffer::F32(g) => g.get(x, y).copied().unwrap_or(0.0),
+                                    ImageBuffer::I32(g) => bit_depth
+                                        .parse_integer_sample(g.get(x, y).copied().unwrap_or(0)),
+                                    ImageBuffer::I16(g) => bit_depth.parse_integer_sample(
+                                        g.get(x, y).copied().unwrap_or(0) as i32,
+                                    ),
+                                };
+                                val * solidity
+                            } else {
+                                0.0
+                            };
+
+                            *v = color * mix + *v * (1.0 - mix);
+                        }
+                    }
+
+                    count += 1;
+                    self.c += 1;
+                }
+                self.c = 0;
+                self.x += 1;
+            }
+            self.x = 0;
+            self.y += 1;
+        }
+        count
+    }
+
+    #[inline]
+    fn to_original_coord(&self, x: u32, y: u32) -> (u32, u32) {
+        let width = self.width;
+        let height = self.height;
+        match self.orientation {
+            1 => (x, y),
+            2 => (width - x - 1, y),
+            3 => (width - x - 1, height - y - 1),
+            4 => (x, height - y - 1),
+            5 => (y, x),
+            6 => (y, width - x - 1),
+            7 => (height - y - 1, width - x - 1),
+            8 => (height - y - 1, x),
+            _ => unreachable!(),
+        }
+    }
+}
+
+struct ImageStreamSpotColor<'r> {
+    grid: &'r ImageBuffer,
+    start_offset_xy: (i32, i32),
+    bit_depth: BitDepth,
+    rgb: (f32, f32, f32),
+    solidity: f32,
+}