Cleanup

rslib/src/scheduler/fsrs/memory_state.rs

@@ -13,8 +13,6 @@ use itertools::Itertools;
 
 use super::params::ignore_revlogs_before_ms_from_config;
 use super::rescheduler::Rescheduler;
-use crate::card::CardType;
-use crate::prelude::*;
 use crate::revlog::RevlogEntry;
 use crate::scheduler::answering::get_fuzz_seed;
 use crate::scheduler::fsrs::params::reviews_for_fsrs;
@@ -23,6 +21,8 @@ use crate::scheduler::states::fuzz::with_review_fuzz;
 use crate::search::Negated;
 use crate::search::SearchNode;
 use crate::search::StateKind;
+use crate::{card::CardType, progress::ThrottlingProgressHandler};
+use crate::{prelude::*, scheduler::timing::SchedTimingToday};
 
 #[derive(Debug, Clone, Copy, Default)]
 pub struct ComputeMemoryProgress {
@@ -80,9 +80,6 @@ impl Collection {
         &mut self,
         entries: Vec<UpdateMemoryStateEntry>,
     ) -> Result<()> {
-        const ITEM_CHUNK_SIZE: usize = 100_000;
-        const FSRS_CHUNK_SIZE: usize = 1000;
-
         let timing = self.timing_today()?;
         let usn = self.usn()?;
         for UpdateMemoryStateEntry {
@@ -94,37 +91,83 @@ impl Collection {
             let search =
                 SearchBuilder::all([search.into(), SearchNode::State(StateKind::New).negated()]);
             let revlog = self.revlog_for_srs(search)?;
-            let reschedule = req.as_ref().map(|e| e.reschedule).unwrap_or_default();
+            let reschedule = req.as_ref().map(|e| e.reschedule).unwrap_or(false);
             let last_revlog_info = reschedule.then(|| get_last_revlog_info(&revlog));
-            let mut rescheduler = self
+            let rescheduler = self
                 .get_config_bool(BoolKey::LoadBalancerEnabled)
                 .then(|| Rescheduler::new(self))
                 .transpose()?;
             let fsrs = FSRS::new(req.as_ref().map(|w| &w.params[..]).or(Some([].as_slice())))?;
-            let decay = req.as_ref().map(|w| get_decay_from_params(&w.params));
             let historical_retention = req.as_ref().map(|w| w.historical_retention);
-            let mut items = fsrs_items_for_memory_states(
+            let items = fsrs_items_for_memory_states(
                 &fsrs,
                 revlog,
                 timing.next_day_at,
                 historical_retention.unwrap_or(0.9),
                 ignore_before,
             )?;
+
             let mut progress = self.new_progress_handler::<ComputeMemoryProgress>();
             progress.update(false, |s| s.total_cards = items.len() as u32)?;
-
             let Some(req) = &req else {
                 // clear FSRS data if FSRS is disabled
-                for (idx, (card_id, _)) in items.into_iter().enumerate() {
+                self.clear_fsrs_data_for_cards(
+                    items.into_iter().map(|(card_id, _)| card_id),
+                    usn,
+                    &mut progress,
+                )?;
+                return Ok(());
+            };
+            let last_revlog_info = last_revlog_info.unwrap();
+
+            self.update_memory_state_for_items(
+                items,
+                last_revlog_info,
+                req,
+                &fsrs,
+                rescheduler,
+                &mut progress,
+                timing,
+                usn,
+            )?;
+        }
+        Ok(())
+    }
+
+    fn clear_fsrs_data_for_cards(
+        &mut self,
+        cards: impl Iterator<Item = CardId>,
+        usn: Usn,
+        progress: &mut ThrottlingProgressHandler<ComputeMemoryProgress>,
+    ) -> Result<()> {
+        for (idx, card_id) in cards.enumerate() {
             progress.update(true, |state| state.current_cards = idx as u32 + 1)?;
             let mut card = self.storage.get_card(card_id)?.or_not_found(card_id)?;
             let original = card.clone();
             card.clear_fsrs_data();
             self.update_card_inner(&mut card, original, usn)?;
         }
-                return Ok(());
+        Ok(())
-            };
+    }
-            let preset_desired_retention = req.preset_desired_retention;
+
+    #[expect(clippy::too_many_arguments)]
+    fn update_memory_state_for_items(
+        &mut self,
+        mut items: Vec<(CardId, Option<FsrsItemForMemoryState>)>,
+        last_revlog_info: HashMap<CardId, LastRevlogInfo>,
+
+        req: &UpdateMemoryStateRequest,
+        fsrs: &FSRS,
+        mut rescheduler: Option<Rescheduler>,
+        progress: &mut ThrottlingProgressHandler<ComputeMemoryProgress>,
+        timing: SchedTimingToday,
+
+        usn: Usn,
+    ) -> Result<()> {
+        const ITEM_CHUNK_SIZE: usize = 100_000;
+        const FSRS_CHUNK_SIZE: usize = 1000;
+
+        let decay = get_decay_from_params(&req.params);
 
         let mut to_update = Vec::new();
         let mut fsrs_items = Vec::new();
@@ -145,9 +188,9 @@ impl Collection {
                 let desired_retention = *req
                     .deck_desired_retention
                     .get(&deck_id)
-                        .unwrap_or(&preset_desired_retention);
+                    .unwrap_or(&req.preset_desired_retention);
                 card.desired_retention = Some(desired_retention);
-                    card.decay = decay;
+                card.decay = Some(decay);
                 if let Some(item) = item {
                     to_update.push((card, original));
                     fsrs_items.push(item.item);
@@ -161,8 +204,7 @@ impl Collection {
 
             // fsrs.memory_state_batch is O(nm) where n is the number of cards and m is the max review count between all items.
             // Therefore we want to pass batches to fsrs.memory_state_batch where the review count is relatively even.
-                let mut p =
+            let mut p = permutation::sort_unstable_by_key(&fsrs_items, |item| item.reviews.len());
-                    permutation::sort_unstable_by_key(&fsrs_items, |item| item.reviews.len());
             p.apply_slice_in_place(&mut to_update);
             p.apply_slice_in_place(&mut fsrs_items);
             p.apply_slice_in_place(&mut starting_states);
@@ -174,25 +216,19 @@ impl Collection {
             {
                 let memory_states = fsrs.memory_state_batch(fsrs_items, starting_states)?;
 
-                    for ((mut card, original), memory_state) in
+                for ((mut card, original), memory_state) in to_update.into_iter().zip(memory_states)
-                        to_update.into_iter().zip(memory_states)
                 {
                     card.memory_state = Some(memory_state.into());
 
                     'reschedule_card: {
-                            // if rescheduling
+                        // we have a last review time for the card
-                            let Some(reviews) = &last_revlog_info else {
+                        let Some(last_info) = last_revlog_info.get(&card.id) else {
-                                break 'reschedule_card;
-                            };
-
-                            // and we have a last review time for the card
-                            let Some(last_info) = reviews.get(&card.id) else {
                             break 'reschedule_card;
                         };
                         let Some(last_review) = &last_info.last_reviewed_at else {
                             break 'reschedule_card;
                         };
-                            // or in (re)learning
+                        // the card isn't in (re)learning
                         if card.ctype != CardType::Review {
                             break 'reschedule_card;
                         };
@@ -249,7 +285,6 @@ impl Collection {
                 }
             }
         }
-        }
         Ok(())
     }