Rework tree building logic

https://github.com/ankitects/anki/pull/1638#discussion_r798328734

rslib/src/decks/limits.rs

@@ -11,7 +11,6 @@ use crate::{
 
 #[derive(Clone, Copy, Debug, PartialEq)]
 pub(crate) struct RemainingLimits {
-    pub deck_id: DeckId,
     pub review: u32,
     pub new: u32,
 }
@@ -26,7 +25,6 @@ impl RemainingLimits {
                     rev_today += new_today;
                 }
                 RemainingLimits {
-                    deck_id: deck.id,
                     review: ((config.inner.reviews_per_day as i32) - rev_today).max(0) as u32,
                     new: ((config.inner.new_per_day as i32) - new_today).max(0) as u32,
                 }
@@ -43,7 +41,6 @@ impl RemainingLimits {
 impl Default for RemainingLimits {
     fn default() -> Self {
         RemainingLimits {
-            deck_id: DeckId(1),
             review: 9999,
             new: 9999,
         }

rslib/src/scheduler/queue/builder/context.rs

@@ -1,9 +1,9 @@
 // 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 std::{collections::HashMap, iter::Peekable};
 
-use id_tree::{InsertBehavior, Node, NodeId, Tree, TreeBuilder};
+use id_tree::{InsertBehavior, Node, NodeId, Tree};
 
 use super::{BuryMode, QueueSortOptions};
 use crate::{
@@ -11,13 +11,37 @@ use crate::{
     scheduler::timing::SchedTimingToday,
 };
 
+/// Wrapper of [RemainingLimits] with some additional meta data.
+#[derive(Debug, Clone, Copy)]
+struct NodeLimits {
+    deck_id: DeckId,
+    /// absolute level in the deck hierarchy
+    level: usize,
+    limits: RemainingLimits,
+}
+
+impl NodeLimits {
+    fn new(deck: &Deck, config: &HashMap<DeckConfigId, DeckConfig>, today: u32) -> Self {
+        Self {
+            deck_id: deck.id,
+            level: deck.name.components().count(),
+            limits: RemainingLimits::new(
+                deck,
+                deck.config_id().and_then(|id| config.get(&id)),
+                today,
+                true,
+            ),
+        }
+    }
+}
+
 /// Data container and helper for building queues.
 #[derive(Debug, Clone)]
 pub(super) struct Context {
     pub(super) timing: SchedTimingToday,
     config_map: HashMap<DeckConfigId, DeckConfig>,
     /// The active decks.
-    pub(super) decks: Vec<Deck>,
+    pub(super) root_deck: Deck,
     pub(super) limits: LimitTreeMap,
     pub(super) sort_options: QueueSortOptions,
     deck_map: HashMap<DeckId, Deck>,
@@ -27,10 +51,10 @@ pub(super) struct Context {
 pub(super) struct LimitTreeMap {
     /// A tree representing the remaining limits of the active deck hierarchy.
     //
-    // As long as we never (1) allow a tree without a root and (2) remove nodes,
+    // As long as we never (1) allow a tree without a root, (2) remove nodes,
-    // it's safe to unwrap on Tree::get() and Tree::root_node_id(), even if we
+    // and (3) have more than 1 tree, it's safe to unwrap on Tree::get() and
-    // clone Nodes.
+    // Tree::root_node_id(), even if we clone Nodes.
-    tree: Tree<RemainingLimits>,
+    tree: Tree<NodeLimits>,
     /// A map to access the tree node of a deck. Only decks with a remaining
     /// limit above zero are included.
     map: HashMap<DeckId, NodeId>,
@@ -38,84 +62,101 @@ pub(super) struct LimitTreeMap {
 }
 
 impl LimitTreeMap {
-    /// Returns the newly built [LimitTreeMap] and the represented decks in depth-first order.
+    /// Returns the newly built [LimitTreeMap] and the deck for `deck_id`.
     fn build(
         col: &mut Collection,
         deck_id: DeckId,
         config: &HashMap<DeckConfigId, DeckConfig>,
         today: u32,
-    ) -> Result<(Self, Vec<Deck>)> {
+    ) -> Result<(Self, Deck)> {
-        let mut decks = vec![col.storage.get_deck(deck_id)?.ok_or(AnkiError::NotFound)?];
+        let root_deck = col.storage.get_deck(deck_id)?.ok_or(AnkiError::NotFound)?;
+        let root_limits = NodeLimits::new(&root_deck, config, today);
+        let initial_root_limits = root_limits.limits;
 
-        let root_config = decks[0].config_id().and_then(|id| config.get(&id));
+        let mut tree = Tree::new();
-        let initial_root_limits = RemainingLimits::new(&decks[0], root_config, today, true);
+        let root_id = tree
-        let tree = TreeBuilder::new()
+            .insert(Node::new(root_limits), InsertBehavior::AsRoot)
-            .with_root(Node::new(initial_root_limits))
+            .unwrap();
-            .build();
 
-        let parent_node_id = tree.root_node_id().unwrap().clone();
         let mut map = HashMap::new();
-        map.insert(deck_id, parent_node_id.clone());
+        map.insert(deck_id, root_id.clone());
 
         let mut limits = Self {
             tree,
             map,
             initial_root_limits,
         };
-        decks = limits.add_descendant_nodes(
+        let mut remaining_decks = col.storage.child_decks(&root_deck)?.into_iter().peekable();
-            col,
+        limits.add_child_nodes(root_id, &mut remaining_decks, config, today);
-            &parent_node_id,
-            initial_root_limits,
-            decks,
-            config,
-            today,
-        )?;
 
-        Ok((limits, decks))
+        Ok((limits, root_deck))
     }
 
-    /// Recursively appends all descendants to the provided [NodeMut], adding their
+    /// Recursively appends descendants to the provided parent [Node], and adds
-    /// [NodeId]s to the [HashMap] and appending their [Deck]s to the [Vec<Deck>],
+    /// them to the [HashMap].
-    /// which is returned.
+    /// Given [Deck]s are assumed to arrive in depth-first order.
-    ///
+    /// The tree-from-deck-list logic to is taken from [crate::decks::tree::add_child_nodes].
-    /// The [NodeMut] is assumed to represent the last [Deck] in the [Vec<Deck>].
+    fn add_child_nodes(
-    /// [RemainingLimits] are capped to their parent's limits.
-    /// [Deck]s with empty review limits are _not_ added to the [HashMap].
-    fn add_descendant_nodes(
         &mut self,
-        col: &mut Collection,
+        parent_node_id: NodeId,
-        parent_node_id: &NodeId,
+        remaining_decks: &mut Peekable<impl Iterator<Item = Deck>>,
-        parent_limits: RemainingLimits,
-        mut decks: Vec<Deck>,
         config: &HashMap<DeckConfigId, DeckConfig>,
         today: u32,
-    ) -> Result<Vec<Deck>> {
+    ) {
-        for child_deck in col.storage.immediate_child_decks(&decks[decks.len() - 1])? {
+        let parent = *self.tree.get(&parent_node_id).unwrap().data();
-            let mut child_limits = RemainingLimits::new(
+        while let Some(deck) = remaining_decks.peek() {
-                &child_deck,
+            match deck.name.components().count() {
-                child_deck.config_id().and_then(|id| config.get(&id)),
+                l if l <= parent.level => {
-                today,
+                    // next item is at a higher level
-                true,
+                    break;
-            );
+                }
-            child_limits.cap_to(parent_limits);
+                l if l == parent.level + 1 => {
-
+                    // next item is an immediate descendent of parent
-            let child_node_id = self
+                    self.insert_child_node(deck, parent_node_id.clone(), config, today);
-                .tree
+                    remaining_decks.next();
-                .insert(
+                }
-                    Node::new(child_limits),
+                _ => {
-                    InsertBehavior::UnderNode(&parent_node_id),
+                    // next item is at a lower level
-                )
+                    if let Some(last_child_node_id) = self
-                .unwrap();
+                        .tree
-            if child_limits.review > 0 {
+                        .get(&parent_node_id)
-                self.map.insert(child_deck.id, child_node_id.clone());
+                        .unwrap()
+                        .children()
+                        .last()
+                        .cloned()
+                    {
+                        self.add_child_nodes(last_child_node_id, remaining_decks, config, today)
+                    } else {
+                        // immediate parent is missing, skip the deck until a DB check is run
+                        remaining_decks.next();
+                    }
+                }
             }
-
-            decks.push(child_deck);
-            decks =
-                self.add_descendant_nodes(col, &child_node_id, child_limits, decks, config, today)?;
         }
+    }
 
-        Ok(decks)
+    fn insert_child_node(
+        &mut self,
+        child_deck: &Deck,
+        parent_node_id: NodeId,
+        config: &HashMap<DeckConfigId, DeckConfig>,
+        today: u32,
+    ) {
+        let mut child_limits = NodeLimits::new(child_deck, config, today);
+        child_limits
+            .limits
+            .cap_to(self.tree.get(&parent_node_id).unwrap().data().limits);
+
+        let child_node_id = self
+            .tree
+            .insert(
+                Node::new(child_limits),
+                InsertBehavior::UnderNode(&parent_node_id),
+            )
+            .unwrap();
+        if child_limits.limits.review > 0 {
+            self.map.insert(child_deck.id, child_node_id);
+        }
     }
 }
 
@@ -123,28 +164,21 @@ impl Context {
     pub(super) fn new(col: &mut Collection, deck_id: DeckId) -> Result<Self> {
         let timing = col.timing_for_timestamp(TimestampSecs::now())?;
         let config_map = col.storage.get_deck_config_map()?;
-        let (limits, decks) = LimitTreeMap::build(col, deck_id, &config_map, timing.days_elapsed)?;
+        let (limits, root_deck) =
-        let sort_options = sort_options(&decks[0], &config_map);
+            LimitTreeMap::build(col, deck_id, &config_map, timing.days_elapsed)?;
+        let sort_options = sort_options(&root_deck, &config_map);
         let deck_map = col.storage.get_decks_map()?;
 
         Ok(Self {
             timing,
             config_map,
-            decks,
+            root_deck,
             limits,
             sort_options,
             deck_map,
         })
     }
 
-    pub(super) fn root_deck(&self) -> &Deck {
-        &self.decks[0]
-    }
-
-    pub(super) fn active_deck_ids(&self) -> Vec<DeckId> {
-        self.decks.iter().map(|deck| deck.id).collect()
-    }
-
     pub(super) fn bury_mode(&self, deck_id: DeckId) -> BuryMode {
         self.deck_map
             .get(&deck_id)
@@ -166,39 +200,33 @@ impl LimitTreeMap {
     pub(super) fn is_exhausted(&self, deck_id: DeckId) -> bool {
         self.map.get(&deck_id).is_none()
     }
+
+    pub(super) fn remaining_decks(&self) -> Vec<DeckId> {
+        self.map.keys().copied().collect()
     }
 
     pub(super) fn remaining_node_id(&self, deck_id: DeckId) -> Option<NodeId> {
         self.map.get(&deck_id).map(Clone::clone)
     }
 
-    pub(super) fn decrement_node_and_parent_review(&mut self, node_id: &NodeId) {
+    pub(super) fn decrement_node_and_parent_limits(&mut self, node_id: &NodeId, new: bool) {
         let node = self.tree.get_mut(node_id).unwrap();
-        let parent = node.parent().map(Clone::clone);
+        let parent = node.parent().cloned();
 
-        let mut limit = node.data_mut();
+        let limit = &mut node.data_mut().limits;
-        limit.review -= 1;
+        if if new {
-        if limit.review < 1 {
+            limit.new = limit.new.saturating_sub(1);
+            limit.new
+        } else {
+            limit.review = limit.review.saturating_sub(1);
+            limit.review
+        } < 1
+        {
             self.remove_node_and_descendants_from_map(node_id);
-        }
+        };
 
         if let Some(parent_id) = parent {
-            self.decrement_node_and_parent_review(&parent_id)
+            self.decrement_node_and_parent_limits(&parent_id, new)
-        }
-    }
-
-    pub(super) fn decrement_node_and_parent_new(&mut self, node_id: &NodeId) {
-        let node = self.tree.get_mut(node_id).unwrap();
-        let parent = node.parent().map(Clone::clone);
-
-        let mut limit = node.data_mut();
-        limit.new -= 1;
-        if limit.new < 1 {
-            self.remove_node_and_descendants_from_map(node_id);
-        }
-
-        if let Some(parent_id) = parent {
-            self.decrement_node_and_parent_new(&parent_id)
         }
     }
 
@@ -217,9 +245,9 @@ impl LimitTreeMap {
 
     fn cap_new_to_review_rec(&mut self, node_id: &NodeId, parent_limit: u32) {
         let node = self.tree.get_mut(node_id).unwrap();
-        let limit = node.data_mut();
+        let mut limits = node.data_mut().limits;
-        limit.new = limit.new.min(limit.review).min(parent_limit);
+        limits.new = limits.new.min(limits.review).min(parent_limit);
-        let node_limit = limit.new;
+        let node_limit = limits.new;
 
         for child_id in node.children().clone() {
             self.cap_new_to_review_rec(&child_id, node_limit);
@@ -235,6 +263,7 @@ impl LimitTreeMap {
                     .get(self.tree.root_node_id().unwrap())
                     .unwrap()
                     .data()
+                    .limits
                     .review,
             ),
             ..self.initial_root_limits

rslib/src/scheduler/queue/builder/mod.rs

@@ -208,7 +208,7 @@ fn sort_learning(mut learning: Vec<DueCard>) -> VecDeque<LearningQueueEntry> {
 impl Collection {
     pub(crate) fn build_queues(&mut self, deck_id: DeckId) -> Result<CardQueues> {
         let mut ctx = Context::new(self, deck_id)?;
-        self.storage.update_active_decks(ctx.root_deck())?;
+        self.storage.update_active_decks(&ctx.root_deck)?;
         let mut queues = QueueBuilder::new(ctx.sort_options.clone());
 
         self.add_intraday_learning_cards(&mut queues, &mut ctx)?;
@@ -268,7 +268,7 @@ impl Collection {
                     let bury = ctx.bury_mode(card.original_deck_id.or(card.current_deck_id));
                     if let Some(node_id) = ctx.limits.remaining_node_id(card.current_deck_id) {
                         if queues.add_due_card(card, bury) {
-                            ctx.limits.decrement_node_and_parent_review(&node_id);
+                            ctx.limits.decrement_node_and_parent_limits(&node_id, false);
                         }
                     }
 
@@ -282,7 +282,7 @@ impl Collection {
     fn add_new_cards_by_deck(&self, queues: &mut QueueBuilder, ctx: &mut Context) -> Result<()> {
         // TODO: must own Vec as closure below requires unique access to ctx
         // maybe decks should not be field of Context?
-        for deck_id in ctx.active_deck_ids() {
+        for deck_id in ctx.limits.remaining_decks() {
             if ctx.limits.is_exhausted_root() {
                 break;
             }
@@ -293,7 +293,7 @@ impl Collection {
                     // and only adjusted the parent nodes after this node's limit is reached
                     if let Some(node_id) = ctx.limits.remaining_node_id(deck_id) {
                         if queues.add_new_card(card, bury) {
-                            ctx.limits.decrement_node_and_parent_new(&node_id);
+                            ctx.limits.decrement_node_and_parent_limits(&node_id, true);
                         }
 
                         true
@@ -318,7 +318,7 @@ impl Collection {
                 let bury = ctx.bury_mode(card.original_deck_id.or(card.current_deck_id));
                 if let Some(node_id) = ctx.limits.remaining_node_id(card.current_deck_id) {
                     if queues.add_new_card(card, bury) {
-                        ctx.limits.decrement_node_and_parent_new(&node_id);
+                        ctx.limits.decrement_node_and_parent_limits(&node_id, true);
                     }
 
                     true

rslib/src/storage/deck/mod.rs

@@ -17,7 +17,6 @@ use crate::{
     decks::{immediate_parent_name, DeckCommon, DeckKindContainer, DeckSchema11, DueCounts},
     error::DbErrorKind,
     prelude::*,
-    text::escape_sql_wildcards,
 };
 
 fn row_to_deck(row: &Row) -> Result<Deck> {
@@ -211,19 +210,6 @@ impl SqliteStorage {
             .collect()
     }
 
-    pub(crate) fn immediate_child_decks(&self, parent: &Deck) -> Result<Vec<Deck>> {
-        let prefix_start = format!("{}\x1f", parent.name);
-        let prefix_end = format!("{}\x20", parent.name);
-        let child_descendant = format!("{}%\x1f%", escape_sql_wildcards(&prefix_start));
-        self.db
-            .prepare_cached(concat!(
-                include_str!("get_deck.sql"),
-                " where name >= ? and name < ? and not name like ? escape '\\'"
-            ))?
-            .query_and_then([prefix_start, prefix_end, child_descendant], row_to_deck)?
-            .collect()
-    }
-
     pub(crate) fn deck_id_with_children(&self, parent: &Deck) -> Result<Vec<DeckId>> {
         let prefix_start = format!("{}\x1f", parent.name);
         let prefix_end = format!("{}\x20", parent.name);

rslib/src/text.rs

@@ -355,14 +355,6 @@ pub(crate) fn escape_anki_wildcards_for_search_node(txt: &str) -> String {
     }
 }
 
-/// Unescape everything.
-pub(crate) fn escape_sql_wildcards(txt: &str) -> Cow<str> {
-    lazy_static! {
-        static ref RE: Regex = Regex::new(r"_|%").unwrap();
-    }
-    RE.replace_all(txt, "\\$1")
-}
-
 /// Return a function to match input against `search`,
 /// which may contain wildcards.
 pub(crate) fn glob_matcher(search: &str) -> impl Fn(&str) -> bool + '_ {