From f50cd43e75a17589cfc83e5af3317b7645833086 Mon Sep 17 00:00:00 2001
From: Damien Elmes <gpg@ankiweb.net>
Date: Sun, 21 Feb 2021 22:37:54 +1000
Subject: [PATCH] fix incorrect nested review counts in v2 scheduler

https://forums.ankiweb.net/t/problem-with-anki-subdecks/7689
---
 pylib/tests/test_schedv2.py |  8 ++---
 rslib/src/decks/tree.rs     | 63 ++++++++++++++++++++++++++++++++-----
 2 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/pylib/tests/test_schedv2.py b/pylib/tests/test_schedv2.py
index 48dd52cc4e0..7b81b9b383c 100644
--- a/pylib/tests/test_schedv2.py
+++ b/pylib/tests/test_schedv2.py
@@ -445,21 +445,21 @@ def test_review_limits():
     tree = col.sched.deck_due_tree().children
     # (('parent', 1514457677462, 5, 0, 0, (('child', 1514457677463, 5, 0, 0, ()),)))
     assert tree[0].review_count == 5  # parent
-    assert tree[0].children[0].review_count == 5  # child
+    assert tree[0].children[0].review_count == 10  # child
 
     # .counts() should match
     col.decks.select(child["id"])
     col.sched.reset()
-    assert col.sched.counts() == (0, 0, 5)
+    assert col.sched.counts() == (0, 0, 10)
 
     # answering a card in the child should decrement parent count
     c = col.sched.getCard()
     col.sched.answerCard(c, 3)
-    assert col.sched.counts() == (0, 0, 4)
+    assert col.sched.counts() == (0, 0, 9)
 
     tree = col.sched.deck_due_tree().children
     assert tree[0].review_count == 4  # parent
-    assert tree[0].children[0].review_count == 4  # child
+    assert tree[0].children[0].review_count == 9  # child
 
 
 def test_button_spacing():
diff --git a/rslib/src/decks/tree.rs b/rslib/src/decks/tree.rs
index 98cc94ca16f..4280dd64735 100644
--- a/rslib/src/decks/tree.rs
+++ b/rslib/src/decks/tree.rs
@@ -5,6 +5,7 @@ use super::{Deck, DeckKind, DueCounts};
 use crate::{
     backend_proto::DeckTreeNode,
     collection::Collection,
+    config::SchedulerVersion,
     deckconf::{DeckConf, DeckConfID},
     decks::DeckID,
     err::Result,
@@ -122,6 +123,44 @@ fn apply_limits(
     node.review_count = (node.review_count + child_rev_total).min(remaining_rev);
 }
 
+/// Apply parent new limits to children, and add child counts to parents.
+/// Unlike v1, reviews are not capped by their parents, and we return the
+/// uncapped review amount to add to the parent. This is a bit of a hack, and
+/// just tides us over until the v2 queue building code can be reworked.
+/// Counts are (new, review).
+fn apply_limits_v2(
+    node: &mut DeckTreeNode,
+    today: u32,
+    decks: &HashMap<DeckID, Deck>,
+    dconf: &HashMap<DeckConfID, DeckConf>,
+    parent_limits: (u32, u32),
+) -> u32 {
+    let original_rev_count = node.review_count;
+
+    let (mut remaining_new, 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);
+
+    // 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 {
+        child_rev_total +=
+            apply_limits_v2(child, today, decks, dconf, (remaining_new, remaining_rev));
+        child_new_total += child.new_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);
+
+    original_rev_count + child_rev_total
+}
+
 fn remaining_counts_for_deck(
     did: DeckID,
     today: u32,
@@ -244,13 +283,23 @@ impl Collection {
             let counts = self.due_counts(days_elapsed, learn_cutoff, limit)?;
             let dconf = self.storage.get_deck_config_map()?;
             add_counts(&mut tree, &counts);
-            apply_limits(
-                &mut tree,
-                days_elapsed,
-                &decks_map,
-                &dconf,
-                (std::u32::MAX, std::u32::MAX),
-            );
+            if self.scheduler_version() == SchedulerVersion::V2 {
+                apply_limits_v2(
+                    &mut tree,
+                    days_elapsed,
+                    &decks_map,
+                    &dconf,
+                    (std::u32::MAX, std::u32::MAX),
+                );
+            } else {
+                apply_limits(
+                    &mut tree,
+                    days_elapsed,
+                    &decks_map,
+                    &dconf,
+                    (std::u32::MAX, std::u32::MAX),
+                );
+            }
         }
 
         Ok(tree)