Make TIFF chunk-length inference and header-plan merging fast

Fixes #2128

Co-Authored-By: Claude Haiku 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

crates/raw/src/sparse.rs

@@ -75,6 +75,21 @@ impl Ranges {
         self.0 = out;
     }
 
+    /// Add every `[start, end)` of `spans` in one sort and merge (empty ones add nothing).
+    pub fn extend(&mut self, spans: impl IntoIterator<Item = (u64, u64)>) {
+        let mut all: Vec<(u64, u64)> = std::mem::take(&mut self.0);
+        all.extend(spans.into_iter().filter(|&(a, b)| a < b));
+        all.sort_unstable();
+        let mut out: Vec<(u64, u64)> = Vec::with_capacity(all.len());
+        for (a, b) in all {
+            match out.last_mut() {
+                Some(last) if a <= last.1 => last.1 = last.1.max(b),
+                _ => out.push((a, b)),
+            }
+        }
+        self.0 = out;
+    }
+
     /// Whether all of `[start, end)` is covered (an empty range always is).
     pub fn contains(&self, start: u64, end: u64) -> bool {
         start >= end || self.0.iter().any(|&(a, b)| a <= start && end <= b)
@@ -232,6 +247,7 @@ fn tiff_plan(bytes: &[u8], loaded: &Ranges, preview: bool) -> HeaderPlan {
     }
     let Ok(tiff) = Tiff::parse(bytes) else { return HeaderPlan::Whole("TIFF structure") };
     let mut holes = Ranges::new();
+    let mut spans = Vec::new();
     for ifd in tiff.all_ifds() {
         let chunks: Vec<(u64, u64)> = match ifd.image() {
             // a preview's strips are read when the preview is wanted: only the raw image is skipped
@@ -249,11 +265,12 @@ fn tiff_plan(bytes: &[u8], loaded: &Ranges, preview: bool) -> HeaderPlan {
         if total < BULK {
             continue;
         }
-        for (i, &(offset, size)) in chunks.iter().enumerate() {
+        spans.extend(chunks.iter().enumerate().map(|(i, &(offset, size))| {
             let start = if i == 0 { offset.saturating_add(HEAD) } else { offset };
-            holes.insert(start, offset.saturating_add(size).min(len));
-        }
+            (start, offset.saturating_add(size).min(len))
+        }));
     }
+    holes.extend(spans);
     // A Panasonic RW2 keeps its raw data from `RawDataOffset` to the end, outside any TIFF image.
     if holes.is_empty()
         && bytes.starts_with(b"IIU\0")
@@ -567,4 +584,40 @@ mod tests {
         assert!(crate::probe_info(&file).is_ok());
         assert_eq!(run(&file, 16, 0).0, HeaderPlan::Whole("an ARW6 header probe reads its tiles"));
     }
+
+    /// A TIFF with a million one-byte strips listed in descending order and all byte counts 0: the plan is
+    /// built in one sort and merge, not one pass over all strips (or all holes) per strip.
+    #[test]
+    fn many_strips_without_byte_counts_plan_fast() {
+        let n = 1_000_000u32;
+        let (offsets_at, counts_at) = (8 + 2 + 5 * 12 + 4, 8 + 2 + 5 * 12 + 4 + 4 * n);
+        let data = counts_at + 4 * n;
+        let mut tiff = b"II*\0\x08\0\0\0".to_vec();
+        tiff.extend_from_slice(&5u16.to_le_bytes());
+        let entry = |tiff: &mut Vec<u8>, tag: u16, typ: u16, count: u32, value: u32| {
+            tiff.extend_from_slice(&tag.to_le_bytes());
+            tiff.extend_from_slice(&typ.to_le_bytes());
+            tiff.extend_from_slice(&count.to_le_bytes());
+            tiff.extend_from_slice(&value.to_le_bytes());
+        };
+        entry(&mut tiff, 256, 4, 1, 1);
+        entry(&mut tiff, 257, 4, 1, n);
+        entry(&mut tiff, 273, 4, n, offsets_at);
+        entry(&mut tiff, 278, 4, 1, 1);
+        entry(&mut tiff, 279, 4, n, counts_at);
+        tiff.extend_from_slice(&0u32.to_le_bytes());
+        for i in (0..n).rev() {
+            tiff.extend_from_slice(&(data + i).to_le_bytes());
+        }
+        tiff.resize(data as usize + n as usize, 0);
+        let len = tiff.len() as u64;
+        let (tx, rx) = std::sync::mpsc::channel();
+        std::thread::spawn(move || {
+            let mut loaded = Ranges::new();
+            loaded.insert(0, len);
+            let _ = tx.send(header_plan(&tiff, &loaded));
+        });
+        let plan = rx.recv_timeout(std::time::Duration::from_secs(20)).expect("header_plan took longer than 20 s");
+        assert!(!matches!(plan, HeaderPlan::Need(_)), "{plan:?}");
+    }
 }

crates/tiff/src/image.rs

@@ -127,6 +127,13 @@ impl ImageInfo {
         let (cw, ch) = self.chunk_size();
         let n = self.offsets.len().min(per_plane.saturating_mul(planes));
         let mut out = Vec::with_capacity(n);
+        // the distinct offsets in order, for chunks without a byte count (sorted once, not scanned per chunk)
+        let mut sorted: Vec<u64> = Vec::new();
+        if (0..n).any(|i| self.byte_counts.get(i).is_none_or(|&l| l == 0)) {
+            sorted = self.offsets.clone();
+            sorted.sort_unstable();
+            sorted.dedup();
+        }
         for i in 0..n {
             let plane = i / per_plane.max(1);
             let j = i % per_plane.max(1);
@@ -142,7 +149,7 @@ impl ImageInfo {
                 Some(&l) if l > 0 => l,
                 _ => {
                     // infer: up to the next larger offset, else end of data
-                    let next = self.offsets.iter().copied().filter(|&o| o > offset).min().unwrap_or(data_len);
+                    let next = sorted.get(sorted.partition_point(|&o| o <= offset)).copied().unwrap_or(data_len);
                     next.saturating_sub(offset)
                 }
             };
@@ -162,3 +169,48 @@ pub fn chunk_bytes<'a>(data: &'a [u8], c: &Chunk) -> Option<&'a [u8]> {
     let end = (c.offset.saturating_add(c.len)).min(data.len() as u64) as usize;
     Some(&data[start..end])
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn strips(offsets: Vec<u64>, byte_counts: Vec<u64>) -> ImageInfo {
+        ImageInfo {
+            width: 1,
+            height: offsets.len() as u32,
+            bits_per_sample: vec![8],
+            samples_per_pixel: 1,
+            compression: 1,
+            photometric: 1,
+            planar: 1,
+            predictor: 1,
+            sample_format: 1,
+            new_subfile_type: 0,
+            layout: Layout::Strips { rows_per_strip: 1 },
+            offsets,
+            byte_counts,
+        }
+    }
+
+    #[test]
+    fn missing_byte_counts_run_to_the_next_larger_offset() {
+        let info = strips(vec![50, 10, 30, 10, 70], vec![0, 0, 5]);
+        let lens: Vec<u64> = info.chunks(100).iter().map(|c| c.len).collect();
+        assert_eq!(lens, vec![20, 20, 5, 20, 30]);
+    }
+
+    /// A million strips without byte counts (offsets in descending order) are measured in one sort, not one
+    /// scan of all offsets per strip.
+    #[test]
+    fn many_strips_without_byte_counts_are_fast() {
+        let n = 1_000_000u64;
+        let info = strips((0..n).rev().map(|i| 1000 + i * 2).collect(), Vec::new());
+        let (tx, rx) = std::sync::mpsc::channel();
+        std::thread::spawn(move || {
+            let chunks = info.chunks(1000 + n * 2);
+            let _ = tx.send(chunks.iter().map(|c| c.len).collect::<Vec<u64>>());
+        });
+        let lens = rx.recv_timeout(std::time::Duration::from_secs(20)).expect("chunks() took longer than 20 s");
+        assert!(lens.iter().all(|&l| l == 2));
+    }
+}