diff --git a/proto/backend.proto b/proto/backend.proto index 3fcb56a3a..b66b47db8 100644 --- a/proto/backend.proto +++ b/proto/backend.proto @@ -252,6 +252,7 @@ message SchedTimingTodayOut { message DeckTreeIn { bool include_counts = 1; + int64 top_deck_id = 2; } message DeckTreeNode { diff --git a/pylib/anki/rsbackend.py b/pylib/anki/rsbackend.py index 07f5f0b9b..8bde0a458 100644 --- a/pylib/anki/rsbackend.py +++ b/pylib/anki/rsbackend.py @@ -728,9 +728,9 @@ class RustBackend: def remove_deck(self, did: int) -> None: self._run_command(pb.BackendInput(remove_deck=did)) - def deck_tree(self, include_counts: bool) -> DeckTreeNode: + def deck_tree(self, include_counts: bool, top_deck_id: int = 0) -> DeckTreeNode: return self._run_command( - pb.BackendInput(deck_tree=pb.DeckTreeIn(include_counts=include_counts)) + pb.BackendInput(deck_tree=pb.DeckTreeIn(include_counts=include_counts, top_deck_id=top_deck_id)) ).deck_tree def check_database(self) -> List[str]: diff --git a/pylib/anki/schedv2.py b/pylib/anki/schedv2.py index 1901db49c..fc7c94eaa 100644 --- a/pylib/anki/schedv2.py +++ b/pylib/anki/schedv2.py @@ -128,7 +128,7 @@ class Scheduler: self._removeFromFiltered(card) def _reset_counts(self): - tree = self.deck_due_tree() + tree = self.deck_due_tree(self.col.decks.selected()) node = self.col.decks.find_deck_in_tree(tree, int(self.col.conf["curDeck"])) if not node: print("invalid current deck") @@ -217,9 +217,10 @@ order by due""" ) return self.col.backend.legacy_deck_tree() - def deck_due_tree(self) -> DeckTreeNode: - "Returns a tree of decks with counts." - return self.col.backend.deck_tree(include_counts=True) + def deck_due_tree(self, top_deck_id: int = 0) -> DeckTreeNode: + """Returns a tree of decks with counts. + If top_deck_id provided, counts are limited to that node.""" + return self.col.backend.deck_tree(include_counts=True, top_deck_id=top_deck_id) # Getting the next card ########################################################################## diff --git a/rslib/src/backend/mod.rs b/rslib/src/backend/mod.rs index 1da602bbc..2a98e2316 100644 --- a/rslib/src/backend/mod.rs +++ b/rslib/src/backend/mod.rs @@ -453,7 +453,12 @@ impl Backend { } fn deck_tree(&self, input: pb::DeckTreeIn) -> Result { - self.with_col(|col| col.deck_tree(input.include_counts)) + let lim = if input.top_deck_id > 0 { + Some(DeckID(input.top_deck_id)) + } else { + None + }; + self.with_col(|col| col.deck_tree(input.include_counts, lim)) } fn render_existing_card(&self, input: pb::RenderExistingCardIn) -> Result { diff --git a/rslib/src/decks/counts.rs b/rslib/src/decks/counts.rs index bd2402b33..824ed6806 100644 --- a/rslib/src/decks/counts.rs +++ b/rslib/src/decks/counts.rs @@ -12,10 +12,18 @@ pub(crate) struct DueCounts { } impl Collection { - pub(crate) fn due_counts(&mut self) -> Result> { + pub(crate) fn due_counts( + &mut self, + limit_to: Option<&str>, + ) -> Result> { let days_elapsed = self.timing_today()?.days_elapsed; let learn_cutoff = self.learn_cutoff(); - self.storage - .due_counts(self.sched_ver(), days_elapsed, learn_cutoff) + if let Some(limit) = limit_to { + self.storage + .due_counts_limited(self.sched_ver(), days_elapsed, learn_cutoff, limit) + } else { + self.storage + .due_counts(self.sched_ver(), days_elapsed, learn_cutoff) + } } } diff --git a/rslib/src/decks/tree.rs b/rslib/src/decks/tree.rs index 623c273fa..1ff8c009e 100644 --- a/rslib/src/decks/tree.rs +++ b/rslib/src/decks/tree.rs @@ -195,7 +195,11 @@ impl From for LegacyDueCounts { } impl Collection { - pub fn deck_tree(&mut self, counts: bool) -> Result { + /// Get the deck tree, optionally populating it with due counts. + /// If top_deck_id is provided, only the node starting at the provided deck ID will + /// have the counts populated. Currently the entire tree is returned in this case, but + /// this may change in the future. + pub fn deck_tree(&mut self, counts: bool, top_deck_id: Option) -> Result { let names = self.storage.get_all_deck_names()?; let mut tree = deck_names_to_tree(names); @@ -212,7 +216,14 @@ impl Collection { } if counts { - let counts = self.due_counts()?; + let limit = top_deck_id.and_then(|did| { + if let Some(deck) = decks_map.get(&did) { + Some(deck.name.as_str()) + } else { + None + } + }); + let counts = self.due_counts(limit)?; let today = self.timing_today()?.days_elapsed; let dconf: HashMap<_, _> = self .storage @@ -234,7 +245,7 @@ impl Collection { } pub(crate) fn legacy_deck_tree(&mut self) -> Result { - let tree = self.deck_tree(true)?; + let tree = self.deck_tree(true, None)?; Ok(LegacyDueCounts::from(tree)) } @@ -272,7 +283,7 @@ mod test { col.get_or_create_normal_deck("2::c::A")?; col.get_or_create_normal_deck("3")?; - let tree = col.deck_tree(false)?; + let tree = col.deck_tree(false, None)?; assert_eq!(tree.children.len(), 3); @@ -295,7 +306,7 @@ mod test { col.storage.remove_deck(col.get_deck_id("2")?.unwrap())?; col.storage.remove_deck(col.get_deck_id("2::3")?.unwrap())?; - let tree = col.deck_tree(false)?; + let tree = col.deck_tree(false, None)?; assert_eq!(tree.children.len(), 1); Ok(()) @@ -314,7 +325,7 @@ mod test { note.fields[0] = "{{c1::}} {{c2::}} {{c3::}} {{c4::}}".into(); col.add_note(&mut note, child_deck.id)?; - let tree = col.deck_tree(true)?; + let tree = col.deck_tree(true, None)?; assert_eq!(tree.children[0].new_count, 4); assert_eq!(tree.children[0].children[0].new_count, 4); @@ -325,7 +336,7 @@ mod test { col.add_or_update_deck(&mut parent_deck, false)?; // with the default limit of 20, there should still be 4 due - let tree = col.deck_tree(true)?; + let tree = col.deck_tree(true, None)?; assert_eq!(tree.children[0].new_count, 4); assert_eq!(tree.children[0].children[0].new_count, 4); @@ -334,7 +345,7 @@ mod test { conf.new.per_day = 4; col.add_or_update_deck_config(&mut conf, false)?; - let tree = col.deck_tree(true)?; + let tree = col.deck_tree(true, None)?; assert_eq!(tree.children[0].new_count, 3); assert_eq!(tree.children[0].children[0].new_count, 3); diff --git a/rslib/src/storage/deck/mod.rs b/rslib/src/storage/deck/mod.rs index dd9530876..de10a74e4 100644 --- a/rslib/src/storage/deck/mod.rs +++ b/rslib/src/storage/deck/mod.rs @@ -170,6 +170,39 @@ impl SqliteStorage { .collect() } + pub(crate) fn due_counts_limited( + &self, + sched: SchedulerVersion, + day_cutoff: u32, + learn_cutoff: u32, + top: &str, + ) -> Result> { + let prefix_start = format!("{}\x1f", top); + let prefix_end = format!("{}\x20", top); + self.db + .prepare_cached(concat!( + include_str!("due_counts.sql"), + " and did in (select id from decks where name = ? ", + "or (name >= ? and name < ?)) group by did " + ))? + .query_and_then( + params![ + CardQueue::New as u8, + CardQueue::Review as u8, + day_cutoff, + sched as u8, + CardQueue::Learn as u8, + learn_cutoff, + CardQueue::DayLearn as u8, + top, + prefix_start, + prefix_end, + ], + row_to_due_counts, + )? + .collect() + } + /// Decks referenced by cards but missing. pub(crate) fn missing_decks(&self) -> Result> { self.db