use the backend for the deck due tree

- approx 3x faster on a large test deck
- counts are no longer capped to 1000 in the tree

proto/backend.proto

@@ -89,6 +89,7 @@ message BackendInput {
         AddOrUpdateDeckLegacyIn add_or_update_deck_legacy = 74;
         bool new_deck_legacy = 75;
         int64 remove_deck = 76;
+        Empty deck_tree_legacy = 77;
     }
 }
 
@@ -157,6 +158,7 @@ message BackendOutput {
         int64 add_or_update_deck_legacy = 74;
         bytes new_deck_legacy = 75;
         Empty remove_deck = 76;
+        bytes deck_tree_legacy = 77;
 
         BackendError error = 2047;
     }

pylib/anki/rsbackend.py

@@ -738,6 +738,9 @@ class RustBackend:
     def check_database(self) -> None:
         self._run_command(pb.BackendInput(check_database=pb.Empty()))
 
+    def legacy_deck_tree(self) -> Sequence:
+        bytes = self._run_command(pb.BackendInput(deck_tree_legacy=pb.Empty())).deck_tree_legacy
+        return orjson.loads(bytes)[5]
 
 def translate_string_in(
     key: TR, **kwargs: Union[str, int, float]

pylib/anki/schedv2.py

@@ -264,11 +264,7 @@ order by due"""
         return data
 
     def deckDueTree(self) -> Any:
-        self.col.decks._enable_dconf_cache()
-        try:
-            return self._groupChildren(self.deckDueList())
-        finally:
-            self.col.decks._disable_dconf_cache()
+        return self.col.backend.legacy_deck_tree()
 
     def _groupChildren(self, grps: List[List[Any]]) -> Any:
         # first, split the group names into components

qt/aqt/deckbrowser.py

@@ -215,8 +215,6 @@ where id > ?""",
         def nonzeroColour(cnt, klass):
             if not cnt:
                 klass = "zero-count"
-            if cnt >= 1000:
-                cnt = "1000+"
             return f'<span class="{klass}">{cnt}</span>'
 
         buf += "<td align=right>%s</td><td align=right>%s</td>" % (

rslib/src/backend/mod.rs

@@ -360,6 +360,7 @@ impl Backend {
                 self.check_database()?;
                 OValue::CheckDatabase(pb::Empty {})
             }
+            Value::DeckTreeLegacy(_) => OValue::DeckTreeLegacy(self.deck_tree_legacy()?),
         })
     }
 
@@ -441,13 +442,7 @@ impl Backend {
     }
 
     fn deck_tree(&self, input: pb::DeckTreeIn) -> Result<pb::DeckTreeNode> {
-        self.with_col(|col| {
-            if input.include_counts {
-                todo!()
-            } else {
-                col.deck_tree()
-            }
-        })
+        self.with_col(|col| col.deck_tree(input.include_counts))
     }
 
     fn render_template(&self, input: pb::RenderCardIn) -> Result<pb::RenderCardOut> {
@@ -1054,6 +1049,13 @@ impl Backend {
     fn check_database(&self) -> Result<()> {
         self.with_col(|col| col.transact(None, |col| col.check_database()))
     }
+
+    fn deck_tree_legacy(&self) -> Result<Vec<u8>> {
+        self.with_col(|col| {
+            let tree = col.legacy_deck_tree()?;
+            serde_json::to_vec(&tree).map_err(Into::into)
+        })
+    }
 }
 
 fn translate_arg_to_fluent_val(arg: &pb::TranslateArgValue) -> FluentValue {

rslib/src/collection.rs

@@ -169,6 +169,10 @@ impl Collection {
         Ok(timing)
     }
 
+    pub(crate) fn learn_cutoff(&self) -> u32 {
+        TimestampSecs::now().0 as u32 + self.learn_ahead_secs()
+    }
+
     pub(crate) fn usn(&self) -> Result<Usn> {
         // if we cache this in the future, must make sure to invalidate cache when usn bumped in sync.finish()
         self.storage.usn(self.server)

rslib/src/config.rs

@@ -42,6 +42,7 @@ pub(crate) enum ConfigKey {
     CurrentNoteTypeID,
     NextNewCardPosition,
     SchedulerVersion,
+    LearnAheadSecs,
 }
 #[derive(PartialEq, Serialize_repr, Deserialize_repr, Clone, Copy)]
 #[repr(u8)]
@@ -62,6 +63,7 @@ impl From<ConfigKey> for &'static str {
             ConfigKey::CurrentNoteTypeID => "curModel",
             ConfigKey::NextNewCardPosition => "nextPos",
             ConfigKey::SchedulerVersion => "schedVer",
+            ConfigKey::LearnAheadSecs => "collapseTime",
         }
     }
 }
@@ -156,6 +158,11 @@ impl Collection {
         self.get_config_optional(ConfigKey::SchedulerVersion)
             .unwrap_or(SchedulerVersion::V1)
     }
+
+    pub(crate) fn learn_ahead_secs(&self) -> u32 {
+        self.get_config_optional(ConfigKey::LearnAheadSecs)
+            .unwrap_or(1200)
+    }
 }
 
 #[derive(Deserialize, PartialEq, Debug, Clone, Copy)]

rslib/src/decks/counts.rs

@@ -0,0 +1,21 @@
+// Copyright: Ankitects Pty Ltd and contributors
+// License: GNU AGPL, version 3 or later; http://www.gnu.org/licenses/agpl.html
+
+use crate::{collection::Collection, decks::DeckID, err::Result};
+use std::collections::HashMap;
+
+#[derive(Debug)]
+pub(crate) struct DueCounts {
+    pub new: u32,
+    pub review: u32,
+    pub learning: u32,
+}
+
+impl Collection {
+    pub(crate) fn due_counts(&mut self) -> Result<HashMap<DeckID, DueCounts>> {
+        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)
+    }
+}

rslib/src/decks/mod.rs

@@ -15,8 +15,10 @@ use crate::{
     timestamp::TimestampSecs,
     types::Usn,
 };
+mod counts;
 mod schema11;
 mod tree;
+pub(crate) use counts::DueCounts;
 pub use schema11::DeckSchema11;
 use std::{borrow::Cow, sync::Arc};
 
@@ -86,6 +88,16 @@ impl Deck {
         self.mtime_secs = TimestampSecs::now();
         self.usn = usn;
     }
+
+    /// Return the studied counts if studied today.
+    /// May be negative if user has extended limits.
+    pub(crate) fn new_rev_counts(&self, today: u32) -> (i32, i32) {
+        if self.common.last_day_studied == today {
+            (self.common.new_studied, self.common.review_studied)
+        } else {
+            (0, 0)
+        }
+    }
 }
 
 // fixme: need to bump usn on upgrade if we rename

rslib/src/decks/tree.rs

@@ -1,8 +1,15 @@
 // Copyright: Ankitects Pty Ltd and contributors
 // License: GNU AGPL, version 3 or later; http://www.gnu.org/licenses/agpl.html
 
-use super::Deck;
-use crate::{backend_proto::DeckTreeNode, collection::Collection, decks::DeckID, err::Result};
+use super::{Deck, DeckKind, DueCounts};
+use crate::{
+    backend_proto::DeckTreeNode,
+    collection::Collection,
+    deckconf::{DeckConf, DeckConfID},
+    decks::DeckID,
+    err::Result,
+};
+use serde_tuple::Serialize_tuple;
 use std::{
     collections::{HashMap, HashSet},
     iter::Peekable,
@@ -68,8 +75,107 @@ fn add_collapsed(node: &mut DeckTreeNode, decks: &HashMap<DeckID, Deck>, browser
     }
 }
 
+fn add_counts(node: &mut DeckTreeNode, counts: &HashMap<DeckID, DueCounts>) {
+    if let Some(counts) = counts.get(&DeckID(node.deck_id)) {
+        node.new_count = counts.new;
+        node.review_count = counts.review;
+        node.learn_count = counts.learning;
+    }
+    for child in &mut node.children {
+        add_counts(child, counts);
+    }
+}
+
+/// Apply parent limits to children, and add child counts to parents.
+/// Counts are (new, review).
+fn apply_limits(
+    node: &mut DeckTreeNode,
+    today: u32,
+    decks: &HashMap<DeckID, Deck>,
+    dconf: &HashMap<DeckConfID, DeckConf>,
+    parent_limits: (u32, u32),
+) {
+    let (mut remaining_new, mut remaining_rev) =
+        remaining_counts_for_deck(DeckID(node.deck_id), today, decks, dconf);
+
+    // cap remaining to parent limits
+    remaining_new = remaining_new.min(parent_limits.0);
+    remaining_rev = remaining_rev.min(parent_limits.1);
+
+    // apply our limit to children and tally their counts
+    let mut child_new_total = 0;
+    let mut child_rev_total = 0;
+    for child in &mut node.children {
+        apply_limits(child, today, decks, dconf, (remaining_new, remaining_rev));
+        child_new_total += child.new_count;
+        child_rev_total += child.review_count;
+        // no limit on learning cards
+        node.learn_count += child.learn_count;
+    }
+
+    // add child counts to our count, capped to remaining limit
+    node.new_count = (node.new_count + child_new_total).min(remaining_new);
+    node.review_count = (node.review_count + child_rev_total).min(remaining_rev);
+}
+
+fn remaining_counts_for_deck(
+    did: DeckID,
+    today: u32,
+    decks: &HashMap<DeckID, Deck>,
+    dconf: &HashMap<DeckConfID, DeckConf>,
+) -> (u32, u32) {
+    if let Some(deck) = decks.get(&did) {
+        match &deck.kind {
+            DeckKind::Normal(norm) => {
+                let (new_today, rev_today) = deck.new_rev_counts(today);
+                if let Some(conf) = dconf
+                    .get(&DeckConfID(norm.config_id))
+                    .or_else(|| dconf.get(&DeckConfID(1)))
+                {
+                    let new = (conf.new.per_day as i32).saturating_sub(new_today).max(0);
+                    let rev = (conf.rev.per_day as i32).saturating_sub(rev_today).max(0);
+                    (new as u32, rev as u32)
+                } else {
+                    // missing dconf and fallback
+                    (0, 0)
+                }
+            }
+            DeckKind::Filtered(_) => {
+                // filtered decks have no limit
+                (std::u32::MAX, std::u32::MAX)
+            }
+        }
+    } else {
+        // top level deck with id 0
+        (std::u32::MAX, std::u32::MAX)
+    }
+}
+
+#[derive(Serialize_tuple)]
+pub(crate) struct LegacyDueCounts {
+    name: String,
+    deck_id: i64,
+    review: u32,
+    learn: u32,
+    new: u32,
+    children: Vec<LegacyDueCounts>,
+}
+
+impl From<DeckTreeNode> for LegacyDueCounts {
+    fn from(n: DeckTreeNode) -> Self {
+        LegacyDueCounts {
+            name: n.name,
+            deck_id: n.deck_id,
+            review: n.review_count,
+            learn: n.learn_count,
+            new: n.new_count,
+            children: n.children.into_iter().map(From::from).collect(),
+        }
+    }
+}
+
 impl Collection {
-    pub fn deck_tree(&self) -> Result<DeckTreeNode> {
+    pub fn deck_tree(&mut self, counts: bool) -> Result<DeckTreeNode> {
         let names = self.storage.get_all_deck_names()?;
         let mut tree = deck_names_to_tree(names);
 
@@ -80,11 +186,35 @@ impl Collection {
             .map(|d| (d.id, d))
             .collect();
 
-        add_collapsed(&mut tree, &decks_map, true);
+        add_collapsed(&mut tree, &decks_map, !counts);
+
+        if counts {
+            let counts = self.due_counts()?;
+            let today = self.timing_today()?.days_elapsed;
+            let dconf: HashMap<_, _> = self
+                .storage
+                .all_deck_config()?
+                .into_iter()
+                .map(|d| (d.id, d))
+                .collect();
+            add_counts(&mut tree, &counts);
+            apply_limits(
+                &mut tree,
+                today,
+                &decks_map,
+                &dconf,
+                (std::u32::MAX, std::u32::MAX),
+            );
+        }
 
         Ok(tree)
     }
 
+    pub(crate) fn legacy_deck_tree(&mut self) -> Result<LegacyDueCounts> {
+        let tree = self.deck_tree(true)?;
+        Ok(LegacyDueCounts::from(tree))
+    }
+
     pub(crate) fn add_missing_decks(&mut self, names: &[(DeckID, String)]) -> Result<()> {
         let mut parents = HashSet::new();
         for (_id, name) in names {
@@ -117,7 +247,7 @@ mod test {
         col.get_or_create_normal_deck("2::c::A")?;
         col.get_or_create_normal_deck("3")?;
 
-        let tree = col.deck_tree()?;
+        let tree = col.deck_tree(false)?;
 
         // 4 including default
         assert_eq!(tree.children.len(), 4);
@@ -141,7 +271,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()?;
+        let tree = col.deck_tree(false)?;
         assert_eq!(tree.children.len(), 2);
 
         Ok(())

rslib/src/search/sqlwriter.rs

@@ -226,7 +226,7 @@ impl SqlWriter<'_> {
             "filtered" => write!(self.sql, "c.odid != 0").unwrap(),
             deck => {
                 // rewrite "current" to the current deck name
-                let deck = if deck == "current" {
+                let native_deck = if deck == "current" {
                     let current_did = self.col.get_current_deck_id();
                     self.col
                         .storage
@@ -234,12 +234,11 @@ impl SqlWriter<'_> {
                         .map(|d| d.name)
                         .unwrap_or_else(|| "Default".into())
                 } else {
-                    deck.into()
+                    human_deck_name_to_native(deck)
                 };
 
                 // convert to a regex that includes child decks
-                let human_deck = human_deck_name_to_native(&deck);
-                let re = text_to_re(&human_deck);
+                let re = text_to_re(&native_deck);
                 self.args.push(format!("(?i)^{}($|\x1f)", re));
                 let arg_idx = self.args.len();
                 self.sql.push_str(&format!(concat!(

rslib/src/storage/deck/due_counts.sql

@@ -0,0 +1,39 @@
+select
+  did,
+  -- new
+  sum(queue = ?1),
+  -- reviews
+  sum(
+    queue = ?2
+    and due <= ?3
+  ),
+  -- learning
+  sum(
+    (
+      case
+        -- v2 scheduler
+        ?4
+        when 2 then (
+          queue = ?5
+          and due < ?6
+        )
+        or (
+          queue = ?7
+          and due <= ?3
+        )
+        else (
+          -- v1 scheduler
+          case
+            when queue = ?5
+            and due < ?6 then left / 1000
+            when queue = ?7
+            and due <= ?3 then 1
+            else 0
+          end
+        )
+      end
+    )
+  )
+from cards
+where
+  queue >= 0

rslib/src/storage/deck/mod.rs

@@ -4,7 +4,9 @@
 use super::SqliteStorage;
 use crate::{
     card::CardID,
-    decks::{Deck, DeckCommon, DeckID, DeckKindProto, DeckSchema11},
+    card::CardQueue,
+    config::SchedulerVersion,
+    decks::{Deck, DeckCommon, DeckID, DeckKindProto, DeckSchema11, DueCounts},
     err::{AnkiError, DBErrorKind, Result},
     i18n::{I18n, TR},
     timestamp::TimestampMillis,
@@ -31,6 +33,17 @@ fn row_to_deck(row: &Row) -> Result<Deck> {
     })
 }
 
+fn row_to_due_counts(row: &Row) -> Result<(DeckID, DueCounts)> {
+    Ok((
+        row.get(0)?,
+        DueCounts {
+            new: row.get(1)?,
+            review: row.get(2)?,
+            learning: row.get(3)?,
+        },
+    ))
+}
+
 impl SqliteStorage {
     pub(crate) fn get_all_decks_as_schema11(&self) -> Result<HashMap<DeckID, DeckSchema11>> {
         self.get_all_decks()
@@ -134,6 +147,29 @@ impl SqliteStorage {
             .collect()
     }
 
+    pub(crate) fn due_counts(
+        &self,
+        sched: SchedulerVersion,
+        day_cutoff: u32,
+        learn_cutoff: u32,
+    ) -> Result<HashMap<DeckID, DueCounts>> {
+        self.db
+            .prepare_cached(concat!(include_str!("due_counts.sql"), " 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,
+                ],
+                row_to_due_counts,
+            )?
+            .collect()
+    }
+
     // Upgrading/downgrading/legacy
 
     pub(super) fn add_default_deck(&self, i18n: &I18n) -> Result<()> {