Feat/support load balance and easy days in rescheduling

rslib/src/scheduler/answering/mod.rs

@@ -608,7 +608,7 @@ impl Card {
 /// Return a consistent seed for a given card at a given number of reps.
 /// If for_reschedule is true, we use card.reps - 1 to match the previous
 /// review.
-fn get_fuzz_seed(card: &Card, for_reschedule: bool) -> Option<u64> {
+pub(crate) fn get_fuzz_seed(card: &Card, for_reschedule: bool) -> Option<u64> {
     let reps = if for_reschedule {
         card.reps.saturating_sub(1)
     } else {

rslib/src/scheduler/fsrs/memory_state.rs

@@ -10,10 +10,12 @@ use fsrs::FSRS;
 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::revlog::RevlogReviewKind;
+use crate::scheduler::answering::get_fuzz_seed;
 use crate::scheduler::fsrs::params::reviews_for_fsrs;
 use crate::scheduler::fsrs::params::Params;
 use crate::scheduler::states::fuzz::with_review_fuzz;
@@ -69,6 +71,10 @@ impl Collection {
             } else {
                 None
             };
+            let mut 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 historical_retention = req.as_ref().map(|w| w.historical_retention);
             let items = fsrs_items_for_memory_states(
@@ -99,26 +105,55 @@ impl Collection {
                                 if let Some(state) = &card.memory_state {
                                     // or in (re)learning
                                     if card.ctype == CardType::Review {
+                                        let deck = self
+                                            .get_deck(card.original_or_current_deck_id())?
+                                            .or_not_found(card.original_or_current_deck_id())?;
+                                        let deckconfig_id = deck.config_id().unwrap();
                                         // reschedule it
                                         let original_interval = card.interval;
                                         let interval = fsrs.next_interval(
                                             Some(state.stability),
                                             card.desired_retention.unwrap(),
                                             0,
                                         );
-                                        card.interval = with_review_fuzz(
-                                            card.get_fuzz_factor(true),
-                                            interval,
-                                            1,
-                                            req.max_interval,
-                                        );
+                                        let new_interval =
+                                            if let Some(rescheduler) = &mut rescheduler {
+                                                rescheduler.find_interval(
+                                                    interval,
+                                                    1,
+                                                    req.max_interval,
+                                                    days_elapsed as u32,
+                                                    deckconfig_id,
+                                                    get_fuzz_seed(&card, true),
+                                                )
+                                            } else {
+                                                None
+                                            };
+                                        if let Some(new_interval) = new_interval {
+                                            card.interval = new_interval;
+                                        } else {
+                                            card.interval = with_review_fuzz(
+                                                card.get_fuzz_factor(true),
+                                                interval,
+                                                1,
+                                                req.max_interval,
+                                            )
+                                        }
                                         let due = if card.original_due != 0 {
                                             &mut card.original_due
                                         } else {
                                             &mut card.due
                                         };
-                                        *due = (timing.days_elapsed as i32) - days_elapsed
+                                        let new_due = (timing.days_elapsed as i32) - days_elapsed
                                             + card.interval as i32;
+                                        if let Some(rescheduler) = &mut rescheduler {
+                                            rescheduler.update_due_cnt_per_day(
+                                                *due,
+                                                new_due,
+                                                deckconfig_id,
+                                            );
+                                        }
+                                        *due = new_due;
                                         // Add a rescheduled revlog entry if the last entry wasn't
                                         // rescheduled
                                         if !last_info.last_revlog_is_rescheduled {

rslib/src/scheduler/fsrs/mod.rs

@@ -4,6 +4,7 @@
 mod error;
 pub mod memory_state;
 pub mod params;
+pub mod rescheduler;
 pub mod retention;
 pub mod simulator;
 pub mod try_collect;

rslib/src/scheduler/fsrs/rescheduler.rs

@@ -0,0 +1,201 @@
+// Copyright: Ankitects Pty Ltd and contributors
+// License: GNU AGPL, version 3 or later; http://www.gnu.org/licenses/agpl.html
+use std::collections::HashMap;
+
+use chrono::Datelike;
+use rand::distributions::Distribution;
+use rand::distributions::WeightedIndex;
+use rand::rngs::StdRng;
+use rand::SeedableRng;
+
+use crate::prelude::*;
+use crate::scheduler::states::fuzz::constrained_fuzz_bounds;
+use crate::scheduler::states::load_balancer::build_easy_days_percentages;
+use crate::scheduler::states::load_balancer::calculate_easy_days_modifiers;
+use crate::scheduler::states::load_balancer::EasyDay;
+
+pub struct Rescheduler {
+    today: i32,
+    next_day_at: TimestampSecs,
+    due_cnt_per_day_by_preset: HashMap<DeckConfigId, HashMap<i32, usize>>,
+    due_today_by_preset: HashMap<DeckConfigId, usize>,
+    reviewed_today_by_preset: HashMap<DeckConfigId, usize>,
+    easy_days_percentages_by_preset: HashMap<DeckConfigId, [EasyDay; 7]>,
+}
+
+impl Rescheduler {
+    pub fn new(col: &mut Collection) -> Result<Self> {
+        let timing = col.timing_today()?;
+        let deck_stats = col.storage.get_deck_due_counts()?;
+        let deck_map = col.storage.get_decks_map()?;
+        let did_to_dcid = deck_map
+            .values()
+            .filter_map(|deck| Some((deck.id, deck.config_id()?)))
+            .collect::<HashMap<_, _>>();
+
+        let mut due_cnt_per_day_by_preset: HashMap<DeckConfigId, HashMap<i32, usize>> =
+            HashMap::new();
+        for (did, due_date, count) in deck_stats {
+            let deck_config_id = did_to_dcid[&did];
+            due_cnt_per_day_by_preset
+                .entry(deck_config_id)
+                .or_default()
+                .entry(due_date)
+                .and_modify(|e| *e += count)
+                .or_insert(count);
+        }
+
+        let today = timing.days_elapsed as i32;
+        let due_today_by_preset = due_cnt_per_day_by_preset
+            .iter()
+            .map(|(deck_config_id, config_dues)| {
+                let due_today = config_dues
+                    .iter()
+                    .filter(|(&due, _)| due <= today)
+                    .map(|(_, &count)| count)
+                    .sum();
+                (*deck_config_id, due_today)
+            })
+            .collect();
+
+        let next_day_at = timing.next_day_at;
+        let reviewed_stats = col.storage.studied_today_by_deck(timing.next_day_at)?;
+        let mut reviewed_today_by_preset: HashMap<DeckConfigId, usize> = HashMap::new();
+        for (did, count) in reviewed_stats {
+            if let Some(&deck_config_id) = &did_to_dcid.get(&did) {
+                *reviewed_today_by_preset.entry(deck_config_id).or_default() += count;
+            }
+        }
+
+        let easy_days_percentages_by_preset =
+            build_easy_days_percentages(col.storage.get_deck_config_map()?)?;
+
+        Ok(Self {
+            today,
+            next_day_at,
+            due_cnt_per_day_by_preset,
+            due_today_by_preset,
+            reviewed_today_by_preset,
+            easy_days_percentages_by_preset,
+        })
+    }
+
+    pub fn update_due_cnt_per_day(
+        &mut self,
+        due_before: i32,
+        due_after: i32,
+        deck_config_id: DeckConfigId,
+    ) {
+        if let Some(counts) = self.due_cnt_per_day_by_preset.get_mut(&deck_config_id) {
+            if let Some(count) = counts.get_mut(&due_before) {
+                *count -= 1;
+            }
+            *counts.entry(due_after).or_default() += 1;
+        }
+
+        if due_before <= self.today && due_after > self.today {
+            if let Some(count) = self.due_today_by_preset.get_mut(&deck_config_id) {
+                *count -= 1;
+            }
+        }
+        if due_before > self.today && due_after <= self.today {
+            *self.due_today_by_preset.entry(deck_config_id).or_default() += 1;
+        }
+    }
+
+    fn due_today(&self, deck_config_id: DeckConfigId) -> usize {
+        *self.due_today_by_preset.get(&deck_config_id).unwrap_or(&0)
+    }
+
+    fn reviewed_today(&self, deck_config_id: DeckConfigId) -> usize {
+        *self
+            .reviewed_today_by_preset
+            .get(&deck_config_id)
+            .unwrap_or(&0)
+    }
+
+    pub fn find_interval(
+        &self,
+        interval: f32,
+        minimum: u32,
+        maximum: u32,
+        days_elapsed: u32,
+        deckconfig_id: DeckConfigId,
+        fuzz_seed: Option<u64>,
+    ) -> Option<u32> {
+        let (before_days, after_days) = constrained_fuzz_bounds(interval, minimum, maximum);
+
+        // Don't reschedule the card when it's overdue
+        if after_days < days_elapsed {
+            return None;
+        }
+        // Don't reschedule the card to the past
+        let before_days = before_days.max(days_elapsed);
+
+        // Generate possible intervals and their review counts
+        let possible_intervals: Vec<u32> = (before_days..=after_days).collect();
+        let review_counts: Vec<usize> = possible_intervals
+            .iter()
+            .map(|&ivl| {
+                if ivl > days_elapsed {
+                    let check_due = self.today + ivl as i32 - days_elapsed as i32;
+                    *self
+                        .due_cnt_per_day_by_preset
+                        .get(&deckconfig_id)
+                        .and_then(|counts| counts.get(&check_due))
+                        .unwrap_or(&0)
+                } else {
+                    // today's workload is the sum of backlogs, cards due today and cards reviewed
+                    // today
+                    self.due_today(deckconfig_id) + self.reviewed_today(deckconfig_id)
+                }
+            })
+            .collect();
+        let weekdays: Vec<usize> = possible_intervals
+            .iter()
+            .map(|&ivl| {
+                self.next_day_at
+                    .adding_secs(days_elapsed as i64 * -86400)
+                    .adding_secs((ivl - 1) as i64 * 86400)
+                    .local_datetime()
+                    .unwrap()
+                    .weekday()
+                    .num_days_from_monday() as usize
+            })
+            .collect();
+
+        let easy_days_load = self.easy_days_percentages_by_preset.get(&deckconfig_id)?;
+
+        let easy_days_modifier =
+            calculate_easy_days_modifiers(easy_days_load, &weekdays, &review_counts);
+
+        // calculate params for each day
+        let intervals_and_params = possible_intervals
+            .iter()
+            .enumerate()
+            .map(|(interval_index, &target_interval)| {
+                let weight = match review_counts[interval_index] {
+                    0 => 1.0, // if theres no cards due on this day, give it the full 1.0 weight
+                    card_count => {
+                        let card_count_weight = (1.0 / card_count as f32).powi(2);
+                        let card_interval_weight = 1.0 / target_interval as f32;
+
+                        card_count_weight
+                            * card_interval_weight
+                            * easy_days_modifier[interval_index]
+                    }
+                };
+
+                (target_interval, weight)
+            })
+            .collect::<Vec<_>>();
+
+        let mut rng = StdRng::seed_from_u64(fuzz_seed?);
+
+        let weighted_intervals =
+            WeightedIndex::new(intervals_and_params.iter().map(|k| k.1)).ok()?;
+
+        let selected_interval_index = weighted_intervals.sample(&mut rng);
+        Some(intervals_and_params[selected_interval_index].0)
+    }
+}