Codecs: decode only the requested EXR part; Cryptomatte uses the cumulative cap (#3721) (#3722)

read_layer decoded every part and picked one; use a first-valid-layer reader
that declines all other parts. Cryptomatte now goes through read_meta and
read_layer. Enable exr/rayon on native targets.

Co-authored-by: Claude Opus 5.5 <[email protected]>

Cargo.lock

@@ -1743,6 +1743,7 @@ dependencies = [
  "miniz_oxide 0.8.9",
  "num-complex",
  "pulp",
+ "rayon-core",
  "smallvec",
  "zune-inflate",
 ]

crates/codecs/Cargo.toml

@@ -36,6 +36,8 @@ photocraft-heif = { path = "../heif", optional = true }
 # Orientation is applied a band of rows per task on native; wasm stays single-threaded.
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
 rayon.workspace = true
+# EXR block (de)compression runs on the rayon pool natively; wasm stays single-threaded.
+exr = { version = "1.73", default-features = false, features = ["rayon"] }
 
 [dev-dependencies]
 proptest.workspace = true

crates/codecs/src/codecs/exr.rs

@@ -35,7 +35,7 @@ fn map_exr(e: exr::error::Error) -> CodecError {
 
 /// Read the headers without the `exr` crate's validation (it rejects deep parts outright),
 /// after checking every part's declared size against the limits, before any pixel allocation.
-fn read_meta(bytes: &[u8], limits: &Limits) -> Result<MetaData, CodecError> {
+pub(super) fn read_meta(bytes: &[u8], limits: &Limits) -> Result<MetaData, CodecError> {
     let meta = MetaData::read_from_buffered(Cursor::new(bytes), false).map_err(err)?;
     // Flat decoding reads every part at once, so the parts' sizes are also capped together.
     let mut total: u64 = 0;
@@ -93,12 +93,31 @@ fn pick_part(meta: &MetaData) -> usize {
     best
 }
 
-/// Decode the flat layers via the `exr` crate (it handles every flat compression) and
-/// return the requested part's layer.
-fn read_layer(bytes: &[u8], index: usize) -> Result<Layer<AnyChannels<FlatSamples>>, CodecError> {
-    let image =
-        read().no_deep_data().largest_resolution_level().all_channels().all_layers().all_attributes().from_buffered(Cursor::new(bytes)).map_err(map_exr)?;
-    image.layer_data.into_iter().nth(index).ok_or_else(|| err(format!("part {index} is missing from the file")))
+/// Selects one part for the `exr` crate's first-valid-layer reader: `create_channels_reader`
+/// is asked about the headers in file order and the first `Ok` wins, so every header before
+/// `index` (and, harmlessly, those after it) is declined. Only the chosen part's chunks are
+/// then decompressed and only its planes allocated; the other parts are skipped unread.
+struct OnlyPart<C> {
+    inner: C,
+    index: usize,
+    seen: std::cell::Cell<usize>,
+}
+
+impl<'s, C: ReadChannels<'s>> ReadChannels<'s> for OnlyPart<C> {
+    type Reader = C::Reader;
+
+    fn create_channels_reader(&'s self, header: &exr::meta::header::Header) -> exr::error::Result<Self::Reader> {
+        let position = self.seen.get();
+        self.seen.set(position.saturating_add(1));
+        if position == self.index { self.inner.create_channels_reader(header) } else { Err(exr::error::Error::Invalid("part not requested".into())) }
+    }
+}
+
+/// Decode one part's flat layer via the `exr` crate (it handles every flat compression).
+/// Other parts are not decompressed, so an AOV file costs one part per request, not the file.
+pub(super) fn read_layer(bytes: &[u8], index: usize) -> Result<Layer<AnyChannels<FlatSamples>>, CodecError> {
+    let only = OnlyPart { inner: read().no_deep_data().largest_resolution_level().all_channels(), index, seen: std::cell::Cell::new(0) };
+    only.first_valid_layer().all_attributes().from_buffered(Cursor::new(bytes)).map(|image| image.layer_data).map_err(map_exr)
 }
 
 fn layer_to_image(layer: &Layer<AnyChannels<FlatSamples>>) -> Result<Image, CodecError> {

crates/codecs/src/codecs/exr_cryptomatte.rs

@@ -11,7 +11,6 @@ use std::io::Cursor;
 
 use exr::meta::MetaData;
 use exr::meta::attribute::AttributeValue;
-use exr::prelude::{ReadChannels, ReadLayers, read};
 
 use crate::Format;
 use crate::error::CodecError;
@@ -200,7 +199,8 @@ pub(crate) fn decode(bytes: &[u8], layer_name: &str, limits: &Limits) -> Result<
         .iter()
         .find(|l| l.name.eq_ignore_ascii_case(layer_name))
         .ok_or_else(|| CodecError::unsupported(F, format!("no Cryptomatte layer named {layer_name}")))?;
-    let meta = MetaData::read_from_buffered(Cursor::new(bytes), false).map_err(|e| CodecError::malformed(F, e))?;
+    // Also caps all parts' declared sizes together, like every other flat decode.
+    let meta = super::exr::read_meta(bytes, limits)?;
     if meta.requirements.has_deep_data {
         return Err(CodecError::unsupported(F, "this file has deep parts; Cryptomatte decoding stays with flat files"));
     }
@@ -214,13 +214,7 @@ pub(crate) fn decode(bytes: &[u8], layer_name: &str, limits: &Limits) -> Result<
     let per_pixel = (layer.channels.len().max(1) as u64).saturating_mul(8).saturating_add(std::mem::size_of::<Vec<(f32, f32)>>() as u64);
     limits.check_bytes(w32, h32, per_pixel)?;
 
-    let image = read().no_deep_data().largest_resolution_level().all_channels().all_layers().all_attributes().from_buffered(Cursor::new(bytes)).map_err(
-        |e| match e {
-            exr::error::Error::NotSupported(m) => CodecError::unsupported(F, format!("not supported: {m}")),
-            e => CodecError::malformed(F, e),
-        },
-    )?;
-    let ex_layer = image.layer_data.get(layer.part).ok_or_else(|| CodecError::malformed(F, "the Cryptomatte part is missing"))?;
+    let ex_layer = super::exr::read_layer(bytes, layer.part)?;
 
     let mut ids: Vec<Vec<f32>> = vec![Vec::new(); layer.channels.len()];
     let mut coverages: Vec<Vec<f32>> = vec![Vec::new(); layer.channels.len()];

crates/codecs/tests/exr_multipart.rs

@@ -197,3 +197,32 @@ fn parts_are_capped_together() {
     let e = exr_info(&aov_file(), &limits).unwrap_err().to_string();
     assert!(e.contains("across 3 parts"), "{e}");
 }
+
+/// Two ZIP1 parts of 4x4 RGBA; every chunk of part 0 is then overwritten with garbage, so
+/// inflating it fails.
+fn file_with_corrupt_first_part() -> Vec<u8> {
+    let px = |k: f32| -> Vec<(&'static str, Vec<f32>)> { ["R", "G", "B", "A"].iter().map(|c| (*c, vec![k; 16])).collect() };
+    let mut bytes = write_image(vec![layer_f32("first", Compression::ZIP1, &px(0.25)), layer_f32("second", Compression::ZIP1, &px(0.75))]);
+    let mut cursor = std::io::Cursor::new(&bytes[..]);
+    exr::meta::MetaData::read_from_buffered(&mut cursor, false).unwrap();
+    let table = cursor.position() as usize;
+    // Part 0 owns the first four of the eight chunks (one scan line per ZIP1 block).
+    for i in 0..4 {
+        let at = table + i * 8;
+        let offset = u64::from_le_bytes(bytes[at..at + 8].try_into().unwrap()) as usize;
+        let size = u32::from_le_bytes(bytes[offset + 8..offset + 12].try_into().unwrap()) as usize;
+        bytes[offset + 12..offset + 12 + size].fill(0xA5);
+    }
+    bytes
+}
+
+#[test]
+fn decoding_one_part_does_not_decode_the_others() {
+    let bytes = file_with_corrupt_first_part();
+    // Part 0 is unreadable ...
+    assert!(decode_exr_part(&bytes, 0, &DecodeOptions::default()).is_err());
+    // ... but asking for part 1 must not touch it.
+    let img = decode_exr_part(&bytes, 1, &DecodeOptions::default()).expect("part 1 decodes on its own");
+    assert_eq!((img.width(), img.height()), (4, 4));
+    assert!(img.to_normalized().iter().all(|v| near(*v, 0.75)), "wrong pixels");
+}