feat(jpeg): add incremental checkpoint mode

Author: telecos

benchmarks/benches/decode_jpeg.rs

@@ -175,6 +175,13 @@ fn decode_jpeg_opts(buf: &[u8], options: DecoderOptions) -> Vec<u8> {
     d.decode().unwrap()
 }
 
+fn decode_jpeg_incremental_mode(buf: &[u8]) -> Vec<u8> {
+    let mut d = JpegDecoder::new(ZCursor::new(buf));
+    d.set_incremental_mode(true);
+
+    d.decode().unwrap()
+}
+
 fn decode_no_samp_opts(c: &mut Criterion) {
     let a = sample_path().join("test-images/jpeg/benchmarks/speed_bench.jpg");
 
@@ -336,6 +343,66 @@ fn decode_restart_resume(c: &mut Criterion) {
     });
 }
 
+fn decode_streaming_mode(c: &mut Criterion) {
+    let data = read(sample_path().join("test-images/jpeg/benchmarks/speed_bench_hv_subsampling.jpg"))
+        .unwrap();
+    let mut group = c.benchmark_group("jpeg: Incremental mode checkpoints");
+    group.throughput(Throughput::Bytes(data.len() as u64));
+
+    group.bench_function("one-shot default", |b| {
+        b.iter(|| black_box(decode_jpeg(data.as_slice())));
+    });
+
+    group.bench_function("one-shot incremental-mode", |b| {
+        b.iter(|| black_box(decode_jpeg_incremental_mode(data.as_slice())));
+    });
+
+    let partial = data.len() * 80 / 100;
+
+    group.bench_function("first retry default", |b| {
+        b.iter(|| {
+            let limit = Rc::new(Cell::new(partial));
+            let cursor = GrowableCursor::new(data.as_slice(), Rc::clone(&limit));
+            let mut decoder = JpegDecoder::new(cursor);
+            decoder.decode_headers().expect("headers should fit in 80%");
+            let mut out = vec![0u8; decoder.output_buffer_size().unwrap()];
+
+            match decoder.decode_into(&mut out) {
+                Ok(()) => panic!("scan should not complete at 80% visibility"),
+                Err(e) => assert!(e.is_recoverable_eof(), "got: {e:?}"),
+            }
+
+            limit.set(data.len());
+            decoder
+                .decode_into(&mut out)
+                .expect("scan should complete with full data");
+            black_box(out);
+        });
+    });
+
+    group.bench_function("first retry incremental-mode", |b| {
+        b.iter(|| {
+            let limit = Rc::new(Cell::new(partial));
+            let cursor = GrowableCursor::new(data.as_slice(), Rc::clone(&limit));
+            let mut decoder = JpegDecoder::new(cursor);
+            decoder.set_incremental_mode(true);
+            decoder.decode_headers().expect("headers should fit in 80%");
+            let mut out = vec![0u8; decoder.output_buffer_size().unwrap()];
+
+            match decoder.decode_into(&mut out) {
+                Ok(()) => panic!("scan should not complete at 80% visibility"),
+                Err(e) => assert!(e.is_recoverable_eof(), "got: {e:?}"),
+            }
+
+            limit.set(data.len());
+            decoder
+                .decode_into(&mut out)
+                .expect("scan should complete with full data");
+            black_box(out);
+        });
+    });
+}
+
 criterion_group!(name=benches;
       config={
       let c = Criterion::default();
@@ -345,6 +412,6 @@ criterion_group!(name=benches;
     decode_hv_samp,criterion_benchmark_grayscale,
     decode_hv_samp_prog,decode_h_samp_prog,decode_no_samp_prog,decode_v_samp_prog,
     decode_no_samp_opts,
-    decode_restart_full,decode_restart_resume);
+    decode_restart_full,decode_restart_resume,decode_streaming_mode);
 
 criterion_main!(benches);

crates/zune-jpeg/README.md

@@ -40,6 +40,13 @@ kept across retries. If scan decoding returns recoverable EOF,
 `decoded_output_bytes()` and `decoded_scanlines()` report the stable prefix of
 the output buffer that can be displayed or copied before retrying.
 
+By default, row checkpoints are recorded only after a previous scan decode
+attempt, so one-shot decoding keeps the lowest-overhead path. Call
+`set_incremental_mode(true)` before the first `decode_into()` attempt when the
+caller expects input to arrive incrementally; this records checkpoints during
+the first baseline Huffman scan attempt and can reduce replay work on the next
+retry.
+
 ```Rust
 use zune_core::bytestream::ZCursor;
 use zune_jpeg::JpegDecoder;
@@ -57,6 +64,7 @@ loop {
 }
 
 let mut pixels = vec![0; decoder.output_buffer_size().unwrap()];
+decoder.set_incremental_mode(true);
 
 loop {
   match decoder.decode_into(&mut pixels) {

crates/zune-jpeg/src/decoder.rs

@@ -302,11 +302,25 @@ pub struct JpegDecoder<T> {
     pub(crate) progressive_mcus_buffer: [Vec<i16>; MAX_COMPONENTS],
     /// Whether per-row checkpointing is enabled for the current decode.
     ///
-    /// Only `true` on a retry `decode_into` call (when `scan_state` was
-    /// already `Some` on entry). This keeps the one-shot decode path free
-    /// of per-row overhead while still enabling fine-grained resume on
-    /// incremental retries.
+    /// By default this becomes `true` after a previous scan attempt has run,
+    /// keeping one-shot decode free of per-row overhead. `incremental_mode`
+    /// enables the same checkpoints on the first scan attempt for streaming
+    /// callers.
     pub(crate) mcu_checkpoints_enabled: bool,
+    /// Whether row checkpoints should also be recorded on the first scan
+    /// decode attempt.
+    ///
+    /// Disabled by default to keep one-shot decode free of checkpoint work;
+    /// streaming callers can opt in before `decode_into` to avoid replaying
+    /// from scan start after the first recoverable scan EOF.
+    incremental_mode: bool,
+    /// Whether this decoder has already attempted scan decoding.
+    ///
+    /// `scan_state` becomes `Some` as soon as headers reach SOS, including
+    /// after an explicit `decode_headers` call. This flag tracks the narrower
+    /// condition needed for default checkpoint gating: a previous
+    /// `decode_into` scan attempt actually ran.
+    scan_decode_attempted: bool,
     /// Scratch buffer that header marker parsers fill with the marker body
     /// before mutating decoder state.
     ///
@@ -541,6 +555,8 @@ where
             scan_state: None,
             pixels_decoded: 0,
             mcu_checkpoints_enabled: false,
+            incremental_mode: false,
+            scan_decode_attempted: false,
             progressive_mcus_buffer: core::array::from_fn(|_| Vec::new()),
             marker_body_scratch: Vec::new()
         }
@@ -632,6 +648,33 @@ where
         Some(self.pixels_decoded.min(self.output_buffer_size()?))
     }
 
+    /// Return whether incremental mode is enabled.
+    ///
+    /// Incremental mode records per-row checkpoints during the first scan
+    /// decode attempt, allowing a later retry after recoverable EOF to resume
+    /// from the latest stable row instead of replaying from scan start.
+    ///
+    /// It is disabled by default so one-shot decoding keeps the lowest
+    /// overhead path.
+    #[must_use]
+    pub const fn incremental_mode(&self) -> bool {
+        self.incremental_mode
+    }
+
+    /// Enable or disable incremental mode.
+    ///
+    /// Call this before the first `decode_into` scan attempt when the caller
+    /// expects input to arrive incrementally. In this mode baseline Huffman
+    /// single-SOS scans save row checkpoints on the first attempt, trading a
+    /// small amount of checkpoint work for less replay on the next retry.
+    ///
+    /// The default is `false`, which preserves the zero-overhead one-shot
+    /// path and only enables row checkpoints after a previous scan decode
+    /// attempt has run.
+    pub fn set_incremental_mode(&mut self, enabled: bool) {
+        self.incremental_mode = enabled;
+    }
+
     /// Return the number of output scanlines known to be stable after the
     /// most recent `decode_into` attempt.
     ///
@@ -1238,6 +1281,12 @@ where
     /// from hard failures: `Err(e)` where `e.is_recoverable_eof()` means feed
     /// more input and retry, while any other `Err` is non-recoverable.
     ///
+    /// By default, row checkpoints are enabled after a previous scan decode
+    /// attempt, so the first one-shot decode avoids checkpoint overhead. Call
+    /// [`set_incremental_mode`](Self::set_incremental_mode) before the first
+    /// scan attempt to record row checkpoints immediately when input is
+    /// expected to arrive incrementally.
+    ///
     /// On success the decoder keeps scan-start replay state, so a later
     /// `decode_into` call is well-defined and produces bit-identical pixels.
     /// Replay re-runs entropy decoding from the first SOS.
@@ -1283,15 +1332,6 @@ where
             pixels_written:  usize,
             dc_predictions:  [(i32, i32); MAX_COMPONENTS]
         }
-        // Enable per-row checkpointing only on retry calls (when scan_state
-        // already existed before this decode_into invocation). One-shot
-        // decoding skips checkpoint overhead entirely.
-        let retrying_scan = self.scan_state.is_some();
-        self.mcu_checkpoints_enabled = retrying_scan;
-        if !retrying_scan {
-            self.pixels_decoded = 0;
-        }
-
         let scan_plan = self.scan_state.as_deref().map(|state| ScanPlan {
             scan_start_position:   state.scan_start_position,
             outer_append_snapshot: state.append_snapshot,
@@ -1388,6 +1428,17 @@ where
         let out_len = core::cmp::min(out.len(), expected_size);
         let out = &mut out[0..out_len];
 
+        // By default, enable per-row checkpointing only after a previous
+        // scan decode attempt has run. Incremental mode opts into the same
+        // checkpoints on the first scan attempt so a streaming caller avoids
+        // one scan-start replay.
+        let previous_scan_attempt = self.scan_decode_attempted;
+        self.mcu_checkpoints_enabled = previous_scan_attempt || self.incremental_mode;
+        if !previous_scan_attempt {
+            self.pixels_decoded = 0;
+        }
+        self.scan_decode_attempted = true;
+
         let result: Result<(), DecodeErrors>;
         if self.is_arithmetic {
             #[cfg(feature = "arith")]

crates/zune-jpeg/src/mcu.rs

@@ -257,8 +257,9 @@ impl<T: ZByteReaderTrait> JpegDecoder<T> {
                 // Per-row checkpoint: save bitstream state at the start of
                 // each MCU row so we can resume from here (rather than
                 // replaying from scan start) if ExhaustedData fires mid-row.
-                // Only active on retry calls to keep one-shot decode free of
-                // any overhead (this code is outside the hot MCU inner loop).
+                // Active on retry calls, or on the first attempt when the
+                // caller explicitly opted into incremental mode. Default
+                // one-shot decode keeps this disabled.
                 //
                 // Only for single-SOS baseline (all_components_in_first_scan):
                 // multi-SOS scans can't safely resume mid-scan because later

crates/zune-jpeg/tests/incremental.rs

@@ -1256,10 +1256,67 @@ fn entropy_start(data: &[u8]) -> usize {
     sos_pos + 2 + sos_len
 }
 
+fn first_retry_seek_after_scan_eof(data: &[u8], incremental_mode: bool) -> usize {
+    let expected = decode_oneshot(data);
+    let entropy_start = entropy_start(data);
+    let cutoff = entropy_start + (data.len() - entropy_start) * 60 / 100;
+
+    let limit = Rc::new(Cell::new(cutoff));
+    let seek_log = Rc::new(RefCell::new(Vec::new()));
+    let cursor = GrowableCursor::with_seek_log(data, Rc::clone(&limit), Rc::clone(&seek_log));
+    let mut decoder = JpegDecoder::new(cursor);
+    decoder.set_incremental_mode(incremental_mode);
+
+    decoder
+        .decode_headers()
+        .expect("headers should be fully visible at cutoff");
+    assert_eq!(decoder.incremental_mode(), incremental_mode);
+    let mut out = vec![0u8; decoder.output_buffer_size().unwrap()];
+
+    let err = decoder
+        .decode_into(&mut out)
+        .expect_err("truncated scan should give recoverable EOF");
+    assert!(
+        err.is_recoverable_eof(),
+        "expected recoverable EOF, got {err:?}"
+    );
+
+    seek_log.borrow_mut().clear();
+    limit.set(data.len());
+    decoder
+        .decode_into(&mut out)
+        .expect("full data should allow decode to complete");
+    assert_pixels_match(&out, &expected, "first_retry_seek", data.len());
+
+    let seeks = seek_log.borrow();
+    *seeks
+        .first()
+        .expect("retry must seek to scan start or a checkpoint")
+}
+
+#[test]
+fn incremental_mode_records_checkpoint_on_first_scan_attempt() {
+    let data = include_bytes!("../../../test-images/jpeg/sampling_factors.jpg");
+    let entropy_start = entropy_start(data);
+
+    let default_seek = first_retry_seek_after_scan_eof(data, false);
+    assert_eq!(
+        default_seek, entropy_start,
+        "default mode should replay from scan start on the first retry"
+    );
+
+    let incremental_seek = first_retry_seek_after_scan_eof(data, true);
+    assert!(
+        incremental_seek > entropy_start,
+        "incremental mode should resume from an entropy-data row checkpoint; \
+         entropy_start={entropy_start}, seek={incremental_seek}"
+    );
+}
+
 /// Per-row checkpoint: truncating a non-RST image mid-scan and retrying
 /// must eventually resume from a row checkpoint rather than replaying
-/// from scan start. Per-row checkpointing is only enabled on retry calls
-/// (when scan_state already exists), so the sequence is:
+/// from scan start. In default mode, per-row checkpointing is only enabled
+/// after a previous scan decode attempt has run, so the sequence is:
 ///   1. First call: partial data → ExhaustedData (no checkpoints saved)
 ///   2. Second call (retry, checkpoints now enabled): still partial → ExhaustedData
 ///      (per-row checkpoints ARE saved this time)
@@ -1309,9 +1366,9 @@ fn per_row_checkpoint_avoids_full_scan_replay() {
         "expected a stable partial prefix, got {first_scanlines} scanlines"
     );
 
-    // Second attempt — still truncated. Now per-row checkpoints are enabled
-    // (scan_state exists from the first call). This replays from scan start
-    // and saves per-row checkpoints as it decodes.
+    // Second attempt — still truncated. Now a previous scan decode attempt
+    // has run, so this replays from scan start and saves per-row checkpoints
+    // as it decodes.
     let err = decoder
         .decode_into(&mut out)
         .expect_err("still truncated, should give recoverable EOF again");