From e6225210fc07085b8c5704bd29c9cad4f5319653 Mon Sep 17 00:00:00 2001
From: RumovZ <gp5glkw78@relay.firefox.com>
Date: Fri, 29 Apr 2022 20:56:23 +0200
Subject: [PATCH] Take count_map for IncrementableProgress::get_inner

---
 rslib/src/import_export/mod.rs                       | 12 +++++-------
 rslib/src/import_export/package/apkg/import/media.rs |  6 ++----
 rslib/src/import_export/package/colpkg/import.rs     |  3 ++-
 3 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/rslib/src/import_export/mod.rs b/rslib/src/import_export/mod.rs
index 4ac7b1b48..8f9825e59 100644
--- a/rslib/src/import_export/mod.rs
+++ b/rslib/src/import_export/mod.rs
@@ -77,13 +77,11 @@ impl<P> IncrementableProgress<P> {
     }
 
     /// 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))
+    pub(crate) fn media_db_fn(
+        &mut self,
+        count_map: impl 'static + Fn(usize) -> P,
+    ) -> Result<impl FnMut(usize) -> bool + '_> {
+        Ok(move |count| (self.progress_fn)(count_map(count), true))
     }
 }
 
diff --git a/rslib/src/import_export/package/apkg/import/media.rs b/rslib/src/import_export/package/apkg/import/media.rs
index 451993642..5a1495e4b 100644
--- a/rslib/src/import_export/package/apkg/import/media.rs
+++ b/rslib/src/import_export/package/apkg/import/media.rs
@@ -34,10 +34,8 @@ pub(super) struct MediaUseMap {
 
 impl Context<'_> {
     pub(super) fn prepare_media(&mut self) -> Result<MediaUseMap> {
-        self.progress.set_count_map(ImportProgress::MediaCheck);
-        let existing_sha1s = self
-            .target_col
-            .all_existing_sha1s(self.progress.get_inner()?)?;
+        let db_progress_fn = self.progress.media_db_fn(ImportProgress::MediaCheck)?;
+        let existing_sha1s = self.target_col.all_existing_sha1s(db_progress_fn)?;
         prepare_media(
             &self.meta,
             &mut self.archive,
diff --git a/rslib/src/import_export/package/colpkg/import.rs b/rslib/src/import_export/package/colpkg/import.rs
index 8adb48636..6c69abb31 100644
--- a/rslib/src/import_export/package/colpkg/import.rs
+++ b/rslib/src/import_export/package/colpkg/import.rs
@@ -166,7 +166,8 @@ impl<'a> MediaComparer<'a> {
         Ok(Self(if meta.media_list_is_hashmap() {
             None
         } else {
-            media_manager.register_changes(&mut progress.get_inner()?, log)?;
+            let mut db_progress_fn = progress.media_db_fn(ImportProgress::MediaCheck)?;
+            media_manager.register_changes(&mut db_progress_fn, log)?;
             Some(Box::new(media_manager.checksum_getter()))
         }))
     }