add option to limit deck tree counts to a particular node

proto/backend.proto

@@ -252,6 +252,7 @@ message SchedTimingTodayOut {
 
 message DeckTreeIn {
     bool include_counts = 1;
+    int64 top_deck_id = 2;
 }
 
 message DeckTreeNode {

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]:

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:
+    def deck_due_tree(self, top_deck_id: int = 0) -> DeckTreeNode:
-        "Returns a tree of decks with counts."
+        """Returns a tree of decks with counts.
-        return self.col.backend.deck_tree(include_counts=True)
+        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
     ##########################################################################

rslib/src/backend/mod.rs

@@ -453,7 +453,12 @@ impl Backend {
     }
 
     fn deck_tree(&self, input: pb::DeckTreeIn) -> Result<pb::DeckTreeNode> {
-        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<pb::RenderCardOut> {

rslib/src/decks/counts.rs

@@ -12,10 +12,18 @@ pub(crate) struct DueCounts {
 }
 
 impl Collection {
-    pub(crate) fn due_counts(&mut self) -> Result<HashMap<DeckID, DueCounts>> {
+    pub(crate) fn due_counts(
+        &mut self,
+        limit_to: Option<&str>,
+    ) -> Result<HashMap<DeckID, DueCounts>> {
         let days_elapsed = self.timing_today()?.days_elapsed;
         let learn_cutoff = self.learn_cutoff();
-        self.storage
+        if let Some(limit) = limit_to {
-            .due_counts(self.sched_ver(), days_elapsed, learn_cutoff)
+            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)
+        }
     }
 }

rslib/src/decks/tree.rs

@@ -195,7 +195,11 @@ impl From<DeckTreeNode> for LegacyDueCounts {
 }
 
 impl Collection {
-    pub fn deck_tree(&mut self, counts: bool) -> Result<DeckTreeNode> {
+    /// 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<DeckID>) -> Result<DeckTreeNode> {
         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<LegacyDueCounts> {
-        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);
 

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<HashMap<DeckID, DueCounts>> {
+        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<Vec<DeckID>> {
         self.db