Generalise IncrementableProgress

And use it in entire import_export code instead.

rslib/src/backend/import_export.rs

@@ -56,7 +56,7 @@ impl ImportExportService for Backend {
         &self,
         input: pb::ImportAnkiPackageRequest,
     ) -> Result<pb::ImportAnkiPackageResponse> {
-        self.with_col(|col| col.import_apkg(&input.package_path, &mut self.import_progress_fn()))
+        self.with_col(|col| col.import_apkg(&input.package_path, self.import_progress_fn()))
             .map(Into::into)
     }
 
@@ -90,30 +90,14 @@ impl SearchNode {
 }
 
 impl Backend {
-    fn import_progress_fn(&self) -> impl FnMut(ImportProgress) -> Result<()> {
+    fn import_progress_fn(&self) -> impl FnMut(ImportProgress, bool) -> bool {
         let mut handler = self.new_progress_handler();
-        move |progress| {
-            let throttle = matches!(
-                progress,
-                ImportProgress::Media(_) | ImportProgress::Notes(_)
-            );
-            if handler.update(Progress::Import(progress), throttle) {
-                Ok(())
-            } else {
-                Err(AnkiError::Interrupted)
-            }
-        }
+        move |progress, throttle| handler.update(Progress::Import(progress), throttle)
     }
 
-    fn export_progress_fn(&self) -> impl FnMut(usize) -> Result<()> {
+    fn export_progress_fn(&self) -> impl FnMut(usize, bool) -> bool {
         let mut handler = self.new_progress_handler();
-        move |media_files| {
-            if handler.update(Progress::Export(media_files), true) {
-                Ok(())
-            } else {
-                Err(AnkiError::Interrupted)
-            }
-        }
+        move |progress, throttle| handler.update(Progress::Export(progress), throttle)
     }
 }
 

rslib/src/import_export/mod.rs

@@ -15,23 +15,84 @@ pub enum ImportProgress {
     Notes(usize),
 }
 
-pub(crate) struct IncrementalProgress<F: FnMut(usize) -> Result<()>> {
-    progress_fn: F,
-    counter: usize,
+/// Wrapper around a progress function, usually passed by the [crate::backend::Backend],
+/// to make repeated calls more ergonomic.
+pub(crate) struct IncrementableProgress<P> {
+    progress_fn: Box<dyn FnMut(P, bool) -> bool>,
+    count_map: Option<Box<dyn FnMut(usize) -> P>>,
+    count: usize,
 }
 
-impl<F: FnMut(usize) -> Result<()>> IncrementalProgress<F> {
-    pub(crate) fn new(progress_fn: F) -> Self {
+const UPDATE_INTERVAL: usize = 17;
+
+impl<P> IncrementableProgress<P> {
+    /// `progress_fn: (progress, throttle) -> should_continue`
+    pub(crate) fn new(progress_fn: impl 'static + FnMut(P, bool) -> bool) -> Self {
         Self {
-            progress_fn,
-            counter: 0,
+            progress_fn: Box::new(progress_fn),
+            count_map: None,
+            count: 0,
         }
     }
 
+    /// Resets the count, and defines how it should be mapped to a progress value
+    /// in the future.
+    pub(crate) fn set_count_map(&mut self, count_map: impl 'static + FnMut(usize) -> P) {
+        self.count_map = Some(Box::new(count_map));
+        self.count = 0;
+    }
+
+    /// Increment the progress counter, periodically triggering an update.
+    /// Returns [AnkiError::Interrupted] if the operation should be cancelled.
+    /// Must have called `set_count_map()` before calling this.
     pub(crate) fn increment(&mut self) -> Result<()> {
-        self.counter += 1;
-        if self.counter % 17 == 0 {
-            (self.progress_fn)(self.counter)
+        self.count += 1;
+        if self.count % UPDATE_INTERVAL != 0 {
+            return Ok(());
+        }
+        let progress = self.mapped_progress()?;
+        self.update(progress, true)
+    }
+
+    /// Manually trigger an update.
+    /// Returns [AnkiError::Interrupted] if the operation should be cancelled.
+    pub(crate) fn call(&mut self, progress: P) -> Result<()> {
+        self.update(progress, false)
+    }
+
+    fn mapped_progress(&mut self) -> Result<P> {
+        if let Some(count_map) = self.count_map.as_mut() {
+            Ok(count_map(self.count))
+        } else {
+            Err(AnkiError::invalid_input("count_map not set"))
+        }
+    }
+
+    fn update(&mut self, progress: P, throttle: bool) -> Result<()> {
+        if (self.progress_fn)(progress, throttle) {
+            Ok(())
+        } else {
+            Err(AnkiError::Interrupted)
+        }
+    }
+
+    /// Stopgap for returning a progress fn compliant with the media code.
+    pub(crate) fn get_inner(&mut self) -> Result<impl FnMut(usize) -> bool + '_> {
+        let count_map = self
+            .count_map
+            .as_mut()
+            .ok_or_else(|| AnkiError::invalid_input("count_map not set"))?;
+        let progress_fn = &mut self.progress_fn;
+        Ok(move |count| progress_fn(count_map(count), true))
+    }
+}
+
+impl IncrementableProgress<usize> {
+    /// Allows incrementing without a map, if the progress is of type [usize].
+    pub(crate) fn increment_flat(&mut self) -> Result<()> {
+        self.count += 1;
+        if self.count % 17 == 0 {
+            self.update(self.count, true)
         } else {
             Ok(())
         }

rslib/src/import_export/package/apkg/export.rs

@@ -16,6 +16,7 @@ use crate::{
             colpkg::export::{export_collection, MediaIter},
             Meta,
         },
+        IncrementableProgress,
     },
     io::{atomic_rename, tempfile_in_parent_of},
     prelude::*,
@@ -32,8 +33,9 @@ impl Collection {
         with_media: bool,
         legacy: bool,
         media_fn: Option<Box<dyn FnOnce(HashSet<PathBuf>) -> MediaIter>>,
-        progress_fn: impl FnMut(usize) -> Result<()>,
+        progress_fn: impl 'static + FnMut(usize, bool) -> bool,
     ) -> Result<usize> {
+        let mut progress = IncrementableProgress::new(progress_fn);
         let temp_apkg = tempfile_in_parent_of(out_path.as_ref())?;
         let mut temp_col = NamedTempFile::new()?;
         let temp_col_path = temp_col
@@ -67,7 +69,7 @@ impl Collection {
             col_size,
             media,
             &self.tr,
-            progress_fn,
+            &mut progress,
         )?;
         atomic_rename(temp_apkg, out_path.as_ref(), true)?;
         Ok(data.notes.len())

rslib/src/import_export/package/apkg/import/media.rs

@@ -12,7 +12,7 @@ use crate::{
             media::{extract_media_entries, SafeMediaEntry},
             Meta,
         },
-        ImportProgress, IncrementalProgress,
+        ImportProgress, IncrementableProgress,
     },
     media::{
         files::{add_hash_suffix_to_file_stem, sha1_of_reader},
@@ -34,21 +34,22 @@ pub(super) struct MediaUseMap {
 
 impl Context<'_> {
     pub(super) fn prepare_media(&mut self) -> Result<MediaUseMap> {
-        let progress_fn = |u| (&mut self.progress_fn)(ImportProgress::MediaCheck(u)).is_ok();
-        let existing_sha1s = self.target_col.all_existing_sha1s(progress_fn)?;
+        self.progress.set_count_map(ImportProgress::MediaCheck);
+        let existing_sha1s = self
+            .target_col
+            .all_existing_sha1s(self.progress.get_inner()?)?;
         prepare_media(
             &self.meta,
             &mut self.archive,
             &existing_sha1s,
-            &mut self.progress_fn,
+            &mut self.progress,
         )
     }
 
     pub(super) fn copy_media(&mut self, media_map: &mut MediaUseMap) -> Result<()> {
-        let mut progress =
-            IncrementalProgress::new(|u| (&mut self.progress_fn)(ImportProgress::Media(u)));
+        self.progress.set_count_map(ImportProgress::Media);
         for entry in media_map.used_entries() {
-            progress.increment()?;
+            self.progress.increment()?;
             entry.copy_from_archive(&mut self.archive, &self.target_col.media_folder)?;
         }
         Ok(())
@@ -69,10 +70,10 @@ fn prepare_media(
     meta: &Meta,
     archive: &mut ZipArchive<File>,
     existing_sha1s: &HashMap<String, Sha1Hash>,
-    progress_fn: &mut impl FnMut(ImportProgress) -> Result<()>,
+    progress: &mut IncrementableProgress<ImportProgress>,
 ) -> Result<MediaUseMap> {
     let mut media_map = MediaUseMap::default();
-    let mut progress = IncrementalProgress::new(|u| progress_fn(ImportProgress::MediaCheck(u)));
+    progress.set_count_map(ImportProgress::MediaCheck);
 
     for mut entry in extract_media_entries(meta, archive)? {
         progress.increment()?;

rslib/src/import_export/package/apkg/import/mod.rs

@@ -19,30 +19,27 @@ use crate::{
     import_export::{
         gather::ExchangeData,
         package::{Meta, NoteLog},
-        ImportProgress,
+        ImportProgress, IncrementableProgress,
     },
     prelude::*,
     search::SearchNode,
 };
 
-type ProgressFn = dyn FnMut(ImportProgress) -> Result<()>;
-
 struct Context<'a> {
     target_col: &'a mut Collection,
     archive: ZipArchive<File>,
     meta: Meta,
     data: ExchangeData,
     usn: Usn,
-    progress_fn: &'a mut ProgressFn,
+    progress: IncrementableProgress<ImportProgress>,
 }
 
 impl Collection {
     pub fn import_apkg(
         &mut self,
         path: impl AsRef<Path>,
-        progress_fn: &mut ProgressFn,
+        progress_fn: impl 'static + FnMut(ImportProgress, bool) -> bool,
     ) -> Result<OpOutput<NoteLog>> {
-        progress_fn(ImportProgress::File)?;
         let file = File::open(path)?;
         let archive = ZipArchive::new(file)?;
 
@@ -57,8 +54,9 @@ impl<'a> Context<'a> {
     fn new(
         mut archive: ZipArchive<File>,
         target_col: &'a mut Collection,
-        progress_fn: &'a mut ProgressFn,
+        progress_fn: impl 'static + FnMut(ImportProgress, bool) -> bool,
     ) -> Result<Self> {
+        let progress = IncrementableProgress::new(progress_fn);
         let meta = Meta::from_archive(&mut archive)?;
         let data = ExchangeData::gather_from_archive(
             &mut archive,
@@ -73,13 +71,13 @@ impl<'a> Context<'a> {
             meta,
             data,
             usn,
-            progress_fn,
+            progress,
         })
     }
 
     fn import(&mut self) -> Result<NoteLog> {
         let mut media_map = self.prepare_media()?;
-        (self.progress_fn)(ImportProgress::File)?;
+        self.progress.call(ImportProgress::File)?;
         let note_imports = self.import_notes_and_notetypes(&mut media_map)?;
         let imported_decks = self.import_decks_and_configs()?;
         self.import_cards_and_revlog(&note_imports.id_map, &imported_decks)?;

rslib/src/import_export/package/apkg/import/notes.rs

@@ -10,11 +10,11 @@ use std::{
 
 use sha1::Sha1;
 
-use super::{media::MediaUseMap, Context, ProgressFn};
+use super::{media::MediaUseMap, Context};
 use crate::{
     import_export::{
         package::{media::safe_normalized_file_name, LogNote, NoteLog},
-        ImportProgress, IncrementalProgress,
+        ImportProgress, IncrementableProgress,
     },
     prelude::*,
     text::{
@@ -108,7 +108,7 @@ impl Context<'_> {
     ) -> Result<NoteImports> {
         let mut ctx = NoteContext::new(self.usn, self.target_col, media_map)?;
         ctx.import_notetypes(mem::take(&mut self.data.notetypes))?;
-        ctx.import_notes(mem::take(&mut self.data.notes), self.progress_fn)?;
+        ctx.import_notes(mem::take(&mut self.data.notes), &mut self.progress)?;
         Ok(ctx.imports)
     }
 }
@@ -184,8 +184,12 @@ impl<'n> NoteContext<'n> {
         Ok(())
     }
 
-    fn import_notes(&mut self, notes: Vec<Note>, progress_fn: &mut ProgressFn) -> Result<()> {
-        let mut progress = IncrementalProgress::new(|u| progress_fn(ImportProgress::Notes(u)));
+    fn import_notes(
+        &mut self,
+        notes: Vec<Note>,
+        progress: &mut IncrementableProgress<ImportProgress>,
+    ) -> Result<()> {
+        progress.set_count_map(ImportProgress::Notes);
 
         for mut note in notes {
             progress.increment()?;
@@ -325,14 +329,14 @@ mod test {
             let mut media_map = MediaUseMap::default();
             let mut ctx = NoteContext::new(Usn(1), &mut $col, &mut media_map).unwrap();
             ctx.remapped_notetypes.insert($old_notetype, $new_notetype);
-            ctx.import_notes(vec![$note], &mut |_progress| Ok(()))
-                .unwrap();
+            let mut progress = IncrementableProgress::new(|_, _| true);
+            ctx.import_notes(vec![$note], &mut progress).unwrap();
             ctx.imports.log
         }};
         ($col:expr, $note:expr, $media_map:expr) => {{
             let mut ctx = NoteContext::new(Usn(1), &mut $col, &mut $media_map).unwrap();
-            ctx.import_notes(vec![$note], &mut |_progress| Ok(()))
-                .unwrap();
+            let mut progress = IncrementableProgress::new(|_, _| true);
+            ctx.import_notes(vec![$note], &mut progress).unwrap();
             ctx.imports.log
         }};
         ($col:expr, $note:expr) => {{

rslib/src/import_export/package/apkg/tests.rs

@@ -37,10 +37,10 @@ fn roundtrip() {
             true,
             true,
             None,
-            |_| Ok(()),
+            |_, _| true,
         )
         .unwrap();
-    target_col.import_apkg(&apkg_path, &mut |_| Ok(())).unwrap();
+    target_col.import_apkg(&apkg_path, |_, _| true).unwrap();
 
     target_col.assert_decks();
     target_col.assert_notetype(&notetype);

rslib/src/import_export/package/colpkg/export.rs

@@ -22,6 +22,7 @@ use zstd::{
 use super::super::{MediaEntries, MediaEntry, Meta, Version};
 use crate::{
     collection::CollectionBuilder,
+    import_export::IncrementableProgress,
     io::{atomic_rename, read_dir_files, tempfile_in_parent_of},
     media::files::filename_if_normalized,
     prelude::*,
@@ -39,8 +40,9 @@ impl Collection {
         out_path: impl AsRef<Path>,
         include_media: bool,
         legacy: bool,
-        progress_fn: impl FnMut(usize) -> Result<()>,
+        progress_fn: impl 'static + FnMut(usize, bool) -> bool,
     ) -> Result<()> {
+        let mut progress = IncrementableProgress::new(progress_fn);
         let colpkg_name = out_path.as_ref();
         let temp_colpkg = tempfile_in_parent_of(colpkg_name)?;
         let src_path = self.col_path.clone();
@@ -62,7 +64,7 @@ impl Collection {
             src_media_folder,
             legacy,
             &tr,
-            progress_fn,
+            &mut progress,
         )?;
         atomic_rename(temp_colpkg, colpkg_name, true)
     }
@@ -103,7 +105,7 @@ fn export_collection_file(
     media_dir: Option<PathBuf>,
     legacy: bool,
     tr: &I18n,
-    progress_fn: impl FnMut(usize) -> Result<()>,
+    progress: &mut IncrementableProgress<usize>,
 ) -> Result<()> {
     let meta = if legacy {
         Meta::new_legacy()
@@ -118,15 +120,7 @@ fn export_collection_file(
         MediaIter::empty()
     };
 
-    export_collection(
-        meta,
-        out_path,
-        &mut col_file,
-        col_size,
-        media,
-        tr,
-        progress_fn,
-    )
+    export_collection(meta, out_path, &mut col_file, col_size, media, tr, progress)
 }
 
 /// Write copied collection data without any media.
@@ -143,7 +137,7 @@ pub(crate) fn export_colpkg_from_data(
         col_size,
         MediaIter::empty(),
         tr,
-        |_| Ok(()),
+        &mut IncrementableProgress::new(|_, _| true),
     )
 }
 
@@ -154,7 +148,7 @@ pub(crate) fn export_collection(
     col_size: usize,
     media: MediaIter,
     tr: &I18n,
-    progress_fn: impl FnMut(usize) -> Result<()>,
+    progress: &mut IncrementableProgress<usize>,
 ) -> Result<()> {
     let out_file = File::create(&out_path)?;
     let mut zip = ZipWriter::new(out_file);
@@ -165,7 +159,7 @@ pub(crate) fn export_collection(
     zip.write_all(&meta_bytes)?;
     write_collection(&meta, &mut zip, col, col_size)?;
     write_dummy_collection(&mut zip, tr)?;
-    write_media(&meta, &mut zip, media, progress_fn)?;
+    write_media(&meta, &mut zip, media, progress)?;
     zip.finish()?;
 
     Ok(())
@@ -240,10 +234,10 @@ fn write_media(
     meta: &Meta,
     zip: &mut ZipWriter<File>,
     media: MediaIter,
-    progress_fn: impl FnMut(usize) -> Result<()>,
+    progress: &mut IncrementableProgress<usize>,
 ) -> Result<()> {
     let mut media_entries = vec![];
-    write_media_files(meta, zip, media, &mut media_entries, progress_fn)?;
+    write_media_files(meta, zip, media, &mut media_entries, progress)?;
     write_media_map(meta, media_entries, zip)?;
     Ok(())
 }
@@ -284,12 +278,12 @@ fn write_media_files(
     zip: &mut ZipWriter<File>,
     media: MediaIter,
     media_entries: &mut Vec<MediaEntry>,
-    mut progress_fn: impl FnMut(usize) -> Result<()>,
+    progress: &mut IncrementableProgress<usize>,
 ) -> Result<()> {
     let mut copier = MediaCopier::new(meta);
     for (index, res) in media.0.enumerate() {
+        progress.increment_flat()?;
         let path = res?;
-        progress_fn(index)?;
 
         zip.start_file(index.to_string(), file_options_stored())?;
 

rslib/src/import_export/package/colpkg/import.rs

@@ -18,7 +18,7 @@ use crate::{
             media::{extract_media_entries, SafeMediaEntry},
             Meta,
         },
-        ImportProgress, IncrementalProgress,
+        ImportProgress, IncrementableProgress,
     },
     io::{atomic_rename, tempfile_in_parent_of},
     media::MediaManager,
@@ -30,10 +30,11 @@ pub fn import_colpkg(
     target_col_path: &str,
     target_media_folder: &Path,
     media_db: &Path,
-    mut progress_fn: impl FnMut(ImportProgress) -> Result<()>,
+    progress_fn: impl 'static + FnMut(ImportProgress, bool) -> bool,
     log: &Logger,
 ) -> Result<()> {
-    progress_fn(ImportProgress::File)?;
+    let mut progress = IncrementableProgress::new(progress_fn);
+    progress.call(ImportProgress::File)?;
     let col_path = PathBuf::from(target_col_path);
     let mut tempfile = tempfile_in_parent_of(&col_path)?;
 
@@ -42,13 +43,13 @@ pub fn import_colpkg(
     let meta = Meta::from_archive(&mut archive)?;
 
     copy_collection(&mut archive, &mut tempfile, &meta)?;
-    progress_fn(ImportProgress::File)?;
+    progress.call(ImportProgress::File)?;
     check_collection_and_mod_schema(tempfile.path())?;
-    progress_fn(ImportProgress::File)?;
+    progress.call(ImportProgress::File)?;
 
     restore_media(
         &meta,
-        &mut progress_fn,
+        &mut progress,
         &mut archive,
         target_media_folder,
         media_db,
@@ -76,7 +77,7 @@ fn check_collection_and_mod_schema(col_path: &Path) -> Result<()> {
 
 fn restore_media(
     meta: &Meta,
-    mut progress_fn: impl FnMut(ImportProgress) -> Result<()>,
+    progress: &mut IncrementableProgress<ImportProgress>,
     archive: &mut ZipArchive<File>,
     media_folder: &Path,
     media_db: &Path,
@@ -89,9 +90,9 @@ fn restore_media(
 
     std::fs::create_dir_all(media_folder)?;
     let media_manager = MediaManager::new(media_folder, media_db)?;
-    let mut media_comparer = MediaComparer::new(meta, &mut progress_fn, &media_manager, log)?;
+    let mut media_comparer = MediaComparer::new(meta, progress, &media_manager, log)?;
 
-    let mut progress = IncrementalProgress::new(|u| progress_fn(ImportProgress::Media(u)));
+    progress.set_count_map(ImportProgress::Media);
     for mut entry in media_entries {
         progress.increment()?;
         maybe_restore_media_file(meta, media_folder, archive, &mut entry, &mut media_comparer)?;
@@ -157,15 +158,14 @@ struct MediaComparer<'a>(Option<Box<dyn FnMut(&str) -> Result<Option<Sha1Hash>>
 impl<'a> MediaComparer<'a> {
     fn new(
         meta: &Meta,
-        mut progress_fn: impl FnMut(ImportProgress) -> Result<()>,
+        progress: &mut IncrementableProgress<ImportProgress>,
         media_manager: &'a MediaManager,
         log: &Logger,
     ) -> Result<Self> {
         Ok(Self(if meta.media_list_is_hashmap() {
             None
         } else {
-            let mut db_progress_fn = |u| progress_fn(ImportProgress::MediaCheck(u)).is_ok();
-            media_manager.register_changes(&mut db_progress_fn, log)?;
+            media_manager.register_changes(&mut progress.get_inner()?, log)?;
             Some(Box::new(media_manager.checksum_getter()))
         }))
     }

rslib/src/import_export/package/colpkg/tests.rs

@@ -41,7 +41,7 @@ fn roundtrip() -> Result<()> {
         // export to a file
         let col = collection_with_media(dir, name)?;
         let colpkg_name = dir.join(format!("{name}.colpkg"));
-        col.export_colpkg(&colpkg_name, true, legacy, |_| Ok(()))?;
+        col.export_colpkg(&colpkg_name, true, legacy, |_, _| true)?;
 
         // import into a new collection
         let anki2_name = dir
@@ -57,7 +57,7 @@ fn roundtrip() -> Result<()> {
             &anki2_name,
             &import_media_dir,
             &import_media_db,
-            |_| Ok(()),
+            |_, _| true,
             &terminal(),
         )?;
 
@@ -89,7 +89,7 @@ fn normalization_check_on_export() -> Result<()> {
     // manually write a file in the wrong encoding.
     std::fs::write(col.media_folder.join("ぱぱ.jpg"), "nfd encoding")?;
     assert_eq!(
-        col.export_colpkg(&colpkg_name, true, false, |_| Ok(()))
+        col.export_colpkg(&colpkg_name, true, false, |_, _| true,)
             .unwrap_err(),
         AnkiError::MediaCheckRequired
     );