Add support for cancelling encoding (image-rs#83)

Author: RunDevelopment

src/encode/bc.rs

@@ -5,7 +5,7 @@ use crate::{
     encode::write_util::for_each_f32_rgba_rows,
     n4,
     util::{self, clamp_0_1},
-    EncodingError, Report,
+    EncodingError,
 };
 
 use super::{
@@ -26,7 +26,7 @@ fn block_universal<
         image,
         writer,
         options,
-        mut progress,
+        progress,
         ..
     } = args;
     let width = image.width() as usize;
@@ -46,11 +46,13 @@ fn block_universal<
     };
     let block_count = util::div_ceil(width, BLOCK_WIDTH) * util::div_ceil(height, BLOCK_HEIGHT);
     let mut block_index: usize = 0;
-    let mut report_block = || {
-        if block_index % report_frequency == 0 {
-            progress.report(block_index as f32 / block_count as f32);
-        }
+    let mut report_block = || -> Result<(), EncodingError> {
+        progress.checked_report_if(
+            block_index % report_frequency == 0,
+            block_index as f32 / block_count as f32,
+        )?;
         block_index += 1;
+        Ok(())
     };
 
     for_each_f32_rgba_rows(image, BLOCK_HEIGHT, |rows| {
@@ -62,7 +64,7 @@ fn block_universal<
             let encoded = &mut encoded_buffer[block_index];
 
             encode_block(block, width, &options, encoded);
-            report_block();
+            report_block()?;
         }
 
         // handle last partial block
@@ -83,13 +85,13 @@ fn block_universal<
 
             let encoded = &mut encoded_buffer[block_index];
             encode_block(&block_data, BLOCK_WIDTH, &options, encoded);
-            report_block();
+            report_block()?;
         }
 
-        writer.write_all(cast::as_bytes(&encoded_buffer))
-    })?;
+        writer.write_all(cast::as_bytes(&encoded_buffer))?;
 
-    Ok(())
+        Ok(())
+    })
 }
 
 fn get_4x4_rgba(data: &[[f32; 4]], row_pitch: usize) -> [[f32; 4]; 16] {

src/encode/bi_planar.rs

@@ -9,7 +9,7 @@ use super::{
 use crate::{
     cast::{self, ToLe},
     encode::write_util::for_each_f32_rgba_rows,
-    util, yuv10, yuv16, yuv8, EncodingError, Report,
+    util, yuv10, yuv16, yuv8, EncodingError,
 };
 
 #[allow(clippy::type_complexity)]
@@ -24,7 +24,7 @@ fn bi_planar_universal<P1: ToLe + cast::Castable + Default + Copy, P2: ToLe + ca
         image,
         writer,
         options,
-        mut progress,
+        progress,
         ..
     } = args;
     let width = image.width() as usize;
@@ -43,11 +43,12 @@ fn bi_planar_universal<P1: ToLe + cast::Castable + Default + Copy, P2: ToLe + ca
     let line_group_count = util::div_ceil(height, BLOCK_HEIGHT);
     let report_frequency = util::div_ceil(1024 * 1024, width * BLOCK_HEIGHT);
     let mut group_index = 0;
-    for_each_f32_rgba_rows(image, BLOCK_HEIGHT, |rows| {
-        if group_index % report_frequency == 0 {
-            // occasionally report progress
-            progress.report(group_index as f32 / line_group_count as f32);
-        }
+    for_each_f32_rgba_rows(image, BLOCK_HEIGHT, |rows| -> Result<(), EncodingError> {
+        // occasionally report progress
+        progress.checked_report_if(
+            group_index % report_frequency == 0,
+            group_index as f32 / line_group_count as f32,
+        )?;
         group_index += 1;
 
         // handle full blocks
@@ -70,9 +71,12 @@ fn bi_planar_universal<P1: ToLe + cast::Castable + Default + Copy, P2: ToLe + ca
         }
 
         P1::to_le(&mut plane1_buffer);
-        writer.write_all(cast::as_bytes(&plane1_buffer))
+        writer.write_all(cast::as_bytes(&plane1_buffer))?;
+
+        Ok(())
     })?;
 
+    progress.check_cancelled()?;
     P2::to_le(&mut plane2);
     writer.write_all(cast::as_bytes(&plane2))?;
 

src/encode/encoder.rs

@@ -4,7 +4,7 @@ use bitflags::bitflags;
 
 use crate::{
     cast, encode::write_util::for_each_chunk, ColorFormat, ColorFormatSet, EncodingError,
-    ImageView, Precision, Progress, Report,
+    ImageView, Precision, Progress,
 };
 
 use super::{
@@ -14,14 +14,14 @@ use super::{
 pub(crate) struct Args<'a, 'b, 'c, 'd> {
     pub image: ImageView<'a>,
     pub writer: &'b mut dyn Write,
-    pub progress: Option<&'c mut Progress<'d>>,
+    pub progress: &'c mut Progress<'d>,
     pub options: EncodeOptions,
 }
 impl<'a, 'b, 'c, 'd> Args<'a, 'b, 'c, 'd> {
     fn from(
         image: ImageView<'a>,
         writer: &'b mut dyn Write,
-        progress: Option<&'c mut Progress<'d>>,
+        progress: &'c mut Progress<'d>,
         options: EncodeOptions,
     ) -> Result<Self, EncodingError> {
         Ok(Self {
@@ -245,7 +245,7 @@ impl EncoderSet {
         &self,
         writer: &mut dyn Write,
         image: ImageView,
-        progress: Option<&mut Progress>,
+        progress: &mut Progress,
         options: &EncodeOptions,
     ) -> Result<(), EncodingError> {
         let encoder = self.pick_encoder(image.color(), options);
@@ -258,12 +258,12 @@ fn copy_directly(args: Args) -> Result<(), EncodingError> {
     let Args {
         image,
         writer,
-        mut progress,
+        progress,
         ..
     } = args;
     let color = image.color();
 
-    progress.report(0.0);
+    progress.checked_report(0.0)?;
 
     // sometimes we can always just write everything directly without any processing
     if image.is_contiguous() && (cfg!(target_endian = "little") || color.precision == Precision::U8)

src/encode/mod.rs

@@ -110,33 +110,60 @@ pub(crate) const fn get_encoders(format: Format) -> Option<EncoderSet> {
     })
 }
 
+/// Encodes a single surfaces in the given format and writes the encoded pixel
+/// data to the given writer.
+///
+/// If an error is returned (or the operation is cancelled), the writer may be
+/// in an inconsistent state. Some, all, or none of the pixel data may have been
+/// written.
+///
+/// If `progress` is cancelled before this function returns,
+/// [`EncodingError::Cancelled`] is returned (unless another error occurs first).
+/// Passing `None` as `progress` is equivalent to [`Progress::none()`].
+///
+/// ## Panics
+///
+/// Panics if:
+///
+/// 1. `writer` panics during writes.
+/// 2. The underlying reporter function of `progress` panics (if given).
+/// 3. An allocation fails.
 pub fn encode(
     writer: &mut dyn Write,
     image: ImageView,
     format: Format,
     progress: Option<&mut Progress>,
     options: &EncodeOptions,
 ) -> Result<(), EncodingError> {
+    let mut no_reporting = Progress::none();
+    let progress = progress.unwrap_or(&mut no_reporting);
+
+    // ending quickly if cancelled is a good property to have
+    progress.check_cancelled()?;
+
     #[cfg(feature = "rayon")]
     if options.parallel {
         return encode_parallel(writer, image, format, progress, options);
     }
 
     let encoders = get_encoders(format).ok_or(EncodingError::UnsupportedFormat(format))?;
-    encoders.encode(writer, image, progress, options)
+    encoders.encode(writer, image, progress, options)?;
+
+    // error if the operation was cancelled
+    progress.check_cancelled()
 }
 
 #[cfg(feature = "rayon")]
 fn encode_parallel(
     writer: &mut dyn Write,
     image: ImageView,
     format: Format,
-    mut progress: Option<&mut Progress>,
+    mut progress: &mut Progress,
     options: &EncodeOptions,
 ) -> Result<(), EncodingError> {
     use rayon::iter::{IntoParallelIterator, ParallelIterator};
 
-    use crate::{ParallelProgress, PixelInfo, Report, SplitView};
+    use crate::{ParallelProgress, PixelInfo, SplitView};
 
     let mut options = options.clone();
     // don't cause an infinite loop
@@ -146,7 +173,7 @@ fn encode_parallel(
 
     // optimization for single fragment
     if let Some(single) = split.single() {
-        return encode(writer, single, format, progress, &options);
+        return encode(writer, single, format, Some(progress), &options);
     }
 
     // Prepare the parallel progress reporter. The +1 ensures that 100% is
@@ -160,6 +187,7 @@ fn encode_parallel(
         .into_par_iter()
         .map(|fragment_index| -> Result<Vec<u8>, EncodingError> {
             let fragment = split.get(fragment_index).expect("invalid fragment index");
+            parallel_progress.check_cancelled()?;
 
             // allocate exactly the right amount of memory
             let bytes: usize = pixel_info
@@ -169,10 +197,12 @@ fn encode_parallel(
                 .expect("too many bytes");
             let mut buffer: Vec<u8> = Vec::with_capacity(bytes);
 
-            // encode the fragment without any progress reporting.
-            // progress will be reported per fragment
+            // Encode the fragment without any progress reporting or cancellation.
+            // Fragments are typically very small and encoded quickly, so
+            // reporting progress or checking for cancellation is not necessary.
             encode(&mut buffer, fragment, format, None, &options)?;
 
+            parallel_progress.check_cancelled()?;
             parallel_progress.submit(fragment.height() as u64);
 
             debug_assert_eq!(buffer.len(), bytes);
@@ -182,13 +212,12 @@ fn encode_parallel(
 
     let encoded_fragments = result?;
     for fragment in encoded_fragments {
+        progress.check_cancelled()?;
         writer.write_all(&fragment)?;
     }
 
-    // Report 100%
-    progress.report(1.0);
-
-    Ok(())
+    // report completion
+    progress.checked_report(1.0)
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]

src/encode/sub_sampled.rs

@@ -1,4 +1,4 @@
-use crate::{as_rgba_f32, cast, ch, n1, n8, util, yuv16, yuv8, EncodingError, Report};
+use crate::{as_rgba_f32, cast, ch, n1, n8, util, yuv16, yuv8, EncodingError};
 
 use super::encoder::{Args, Encoder, EncoderSet, Flags};
 
@@ -38,7 +38,7 @@ where
     let Args {
         image,
         writer,
-        mut progress,
+        progress,
         ..
     } = args;
     let color = image.color();
@@ -62,10 +62,11 @@ where
         debug_assert!(y_line.len() == width * bytes_per_pixel);
 
         for chunk in y_line.chunks(chunk_size) {
-            if chunk_index % 4096 == 0 {
-                // occasionally report progress
-                progress.report(chunk_index as f32 / chunk_count as f32);
-            }
+            // occasionally report progress
+            progress.checked_report_if(
+                chunk_index % 4096 == 0,
+                chunk_index as f32 / chunk_count as f32,
+            )?;
             chunk_index += 1;
 
             let pixels = chunk.len() / bytes_per_pixel;