flashraw_codecs/
leaf.rs

1//! Restart-interval parallel lossless JPEG, as written by Leaf (and Phase One-era
2//! Leaf/Mamiya backs) for MOS files.
3//!
4//! Leaf streams carry a DRI marker whose interval is a whole number of JPEG rows, so
5//! each restart interval is an independently decodable band (predictors reset at a
6//! restart). The scan is split at `RSTn` markers and the bands decode in parallel
7//! straight into the output, which is filled as one linear sample stream mapped onto
8//! raw rows in order.
9
10use crate::ljpeg::{LjpegHeader, decode_rows};
11use flashraw_core::{Error, Result, par};
12
13/// Start offsets (into `data`) of every restart interval of the scan.
14pub fn restart_segments(data: &[u8], hdr: &LjpegHeader) -> Vec<usize> {
15    let mut segs = vec![hdr.scan_offset];
16    let mut p = hdr.scan_offset;
17    while p + 1 < data.len() {
18        // Entropy data never contains FF followed by anything but 00 or a marker.
19        match data[p..].iter().position(|&b| b == 0xFF) {
20            Some(i) => p += i,
21            None => break,
22        }
23        let Some(&m) = data.get(p + 1) else { break };
24        if (0xD0..=0xD7).contains(&m) {
25            segs.push(p + 2);
26            p += 2;
27        } else if m == 0xD9 {
28            break;
29        } else {
30            p += 1;
31        }
32    }
33    segs
34}
35
36/// Decode a lossless JPEG stream as a linear sample stream into `out` (row-major
37/// `out.len()` samples), decoding restart intervals in parallel when they cover whole
38/// rows. Samples beyond `out` are dropped; missing samples are left untouched.
39pub fn decode_linear(data: &[u8], hdr: &LjpegHeader, out: &mut [u16], parallel: bool) -> Result<()> {
40    let rl = hdr.row_len();
41    let width = hdr.width as usize;
42    let restart = hdr.restart_interval as usize;
43    if rl == 0 {
44        return Err(Error::corrupt("lossless JPEG: empty rows"));
45    }
46    let band_rows = if restart > 0 && restart % width == 0 { restart / width } else { 0 };
47    let total_rows = hdr.height as usize;
48    if band_rows == 0 || band_rows >= total_rows {
49        return decode_sequential(data, hdr, out);
50    }
51    let segs = restart_segments(data, hdr);
52    let nbands = total_rows.div_ceil(band_rows);
53    if segs.len() < nbands {
54        return Err(Error::corrupt(format!("lossless JPEG: {} restart intervals, expected {nbands}", segs.len())));
55    }
56    let errs = std::sync::Mutex::new(None::<Error>);
57    let usable = (out.len() / rl) * rl;
58    let (body, tail) = out.split_at_mut(usable);
59    par::for_each_row_chunk(body, rl, band_rows, parallel, |first_row, chunk| {
60        let band = first_row / band_rows;
61        let mut h = hdr.clone();
62        h.height = band_rows.min(total_rows - first_row) as u32;
63        h.restart_interval = 0;
64        h.scan_offset = segs[band];
65        let res = decode_rows(data, &h, |r, s| {
66            if let Some(dst) = chunk.get_mut(r * rl..(r + 1) * rl) {
67                dst.copy_from_slice(s);
68            }
69            Ok(())
70        });
71        if let Err(e) = res {
72            if let Ok(mut g) = errs.lock() {
73                g.get_or_insert(e);
74            }
75        }
76    });
77    if let Some(e) = errs.into_inner().ok().flatten() {
78        return Err(e);
79    }
80    // A partial final row (when `out` is not a multiple of the JPEG row length).
81    if !tail.is_empty() && usable / rl < total_rows {
82        let row = usable / rl;
83        let band = row / band_rows;
84        let mut h = hdr.clone();
85        h.height = (row - band * band_rows + 1) as u32;
86        h.restart_interval = 0;
87        h.scan_offset = segs[band];
88        let target = row - band * band_rows;
89        decode_rows(data, &h, |r, s| {
90            if r == target {
91                let n = tail.len().min(s.len());
92                tail[..n].copy_from_slice(&s[..n]);
93            }
94            Ok(())
95        })?;
96    }
97    Ok(())
98}
99
100fn decode_sequential(data: &[u8], hdr: &LjpegHeader, out: &mut [u16]) -> Result<()> {
101    let mut pos = 0usize;
102    decode_rows(data, hdr, |_, s| {
103        if pos < out.len() {
104            let n = (out.len() - pos).min(s.len());
105            out[pos..pos + n].copy_from_slice(&s[..n]);
106            pos += n;
107        }
108        Ok(())
109    })
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use crate::ljpeg::LjpegOptions;
116
117    /// Insert a DRI segment and RST markers into a test-encoded stream by re-encoding
118    /// bands separately and splicing them (each band restarts its predictors).
119    fn encode_with_restarts(w: usize, h: usize, band: usize, samples: &[u16]) -> Vec<u8> {
120        let mut out = Vec::new();
121        for (k, rows) in samples.chunks(w * band).enumerate() {
122            let bh = rows.len() / w;
123            let enc = crate::ljpeg::encoder::encode(w, bh, 1, 14, 1, rows);
124            let hdr = LjpegHeader::parse(&enc, LjpegOptions::default()).unwrap();
125            if k == 0 {
126                // Header with the full height and a DRI marker.
127                let mut head = enc[..hdr.scan_offset].to_vec();
128                // Patch SOF3 height.
129                let sof = head.windows(2).position(|x| x == [0xFF, 0xC3]).unwrap();
130                head[sof + 5..sof + 7].copy_from_slice(&(h as u16).to_be_bytes());
131                // Insert DRI before SOS.
132                let sos = head.windows(2).rposition(|x| x == [0xFF, 0xDA]).unwrap();
133                let dri = [0xFF, 0xDD, 0x00, 0x04, ((w * band) >> 8) as u8, (w * band) as u8];
134                head.splice(sos..sos, dri);
135                out.extend(head);
136            } else {
137                out.extend([0xFF, 0xD0 + ((k - 1) % 8) as u8]);
138            }
139            // Entropy data of this band (strip trailing EOI).
140            let body = &enc[hdr.scan_offset..enc.len() - 2];
141            out.extend(body);
142        }
143        out.extend([0xFF, 0xD9]);
144        out
145    }
146
147    #[test]
148    fn parallel_bands_match_sequential() {
149        let (w, h, band) = (24usize, 37usize, 4usize);
150        let samples: Vec<u16> = (0..w * h).map(|i| ((i * 7919) % 16384) as u16).collect();
151        let data = encode_with_restarts(w, h, band, &samples);
152        let hdr = LjpegHeader::parse(&data, LjpegOptions::default()).unwrap();
153        assert_eq!(hdr.restart_interval as usize, w * band);
154        for parallel in [false, true] {
155            let mut out = vec![0u16; w * h];
156            decode_linear(&data, &hdr, &mut out, parallel).unwrap();
157            assert_eq!(out, samples);
158        }
159        // Truncated stream: error, no panic.
160        let mut out = vec![0u16; w * h];
161        assert!(decode_linear(&data[..data.len() / 2], &hdr, &mut out, true).is_err());
162    }
163}
flashraw_codecs/leaf.rs

flashraw_codecs/
leaf.rs
