fix: handle taiko passed objects properly & improve testing

src/taiko/difficulty_object.rs

@@ -111,8 +111,7 @@ fn closest_rhythm(
         .unwrap()
 }
 
-// TODO: Remove Default impl and replace with `with_capacity` method for efficiency
+#[derive(Clone, Debug)]
-#[derive(Clone, Debug, Default)]
 pub(crate) struct ObjectLists {
     pub(crate) all: Vec<Rc<RefCell<TaikoDifficultyObject>>>,
     pub(crate) centres: Vec<usize>,
@@ -121,6 +120,15 @@ pub(crate) struct ObjectLists {
 }
 
 impl ObjectLists {
+    pub(crate) fn with_capacity(capacity: usize) -> Self {
+        Self {
+            all: Vec::with_capacity(capacity),
+            centres: Vec::new(),
+            rims: Vec::new(),
+            notes: Vec::new(),
+        }
+    }
+
     pub(crate) fn prev_mono(
         &self,
         curr: usize,

src/taiko/gradual_difficulty.rs

@@ -51,9 +51,18 @@ pub struct TaikoGradualDifficulty {
     lists: ObjectLists,
     peaks: Peaks,
     total_hits: usize,
+    first_combos: FirstTwoCombos,
     is_convert: bool,
 }
 
+#[derive(Copy, Clone, Debug)]
+enum FirstTwoCombos {
+    None,
+    OnlyFirst,
+    OnlySecond,
+    Both,
+}
+
 impl TaikoGradualDifficulty {
     /// Create a new difficulty attributes iterator for osu!taiko maps.
     pub fn new(map: &Beatmap, mods: u32) -> Self {
@@ -68,41 +77,44 @@ impl TaikoGradualDifficulty {
             .clock_rate(clock_rate)
             .hit_windows();
 
-        let mut attrs = TaikoDifficultyAttributes {
+        let attrs = TaikoDifficultyAttributes {
-            stamina: 0.0,
-            rhythm: 0.0,
-            colour: 0.0,
-            peak: 0.0,
             hit_window,
-            stars: 0.0,
+            ..Default::default()
-            max_combo: 0,
         };
 
         if map.hit_objects.len() < 2 {
             return Self {
                 idx: 0,
                 diff_objects: Vec::new().into_iter(),
-                lists: ObjectLists::default(),
+                lists: ObjectLists::with_capacity(0),
                 peaks,
                 attrs,
                 total_hits: 0,
+                first_combos: FirstTwoCombos::None,
                 is_convert,
             };
         }
 
-        attrs.max_combo += map.hit_objects[0].is_circle() as usize;
+        let first_combos = match (
-        attrs.max_combo += map.hit_objects[1].is_circle() as usize;
+            map.hit_objects[0].is_circle(),
-        let mut total_hits = attrs.max_combo;
+            map.hit_objects[1].is_circle(),
-        let mut diff_objects = ObjectLists::default();
+        ) {
+            (false, false) => FirstTwoCombos::None,
+            (true, false) => FirstTwoCombos::OnlyFirst,
+            (false, true) => FirstTwoCombos::OnlySecond,
+            (true, true) => FirstTwoCombos::Both,
+        };
+
+        let mut total_hits = 0;
+        let mut diff_objects = ObjectLists::with_capacity(map.hit_objects.len().saturating_sub(2));
 
         map.taiko_objects()
+            .inspect(|(h, _)| total_hits += h.is_hit as usize)
             .skip(2)
             .zip(map.hit_objects.iter().skip(1))
             .zip(map.hit_objects.iter())
             .enumerate()
             .for_each(|(idx, (((base, base_start_time), last), last_last))| {
-                total_hits += base.is_hit as usize;
-
                 let diff_obj = TaikoDifficultyObject::new(
                     base,
                     base_start_time,
@@ -135,6 +147,7 @@ impl TaikoGradualDifficulty {
             peaks,
             attrs,
             total_hits,
+            first_combos,
             is_convert,
         }
     }
@@ -162,6 +175,14 @@ impl Iterator for TaikoGradualDifficulty {
             }
         } else if self.lists.all.is_empty() {
             return None;
+        } else {
+            match self.first_combos {
+                FirstTwoCombos::OnlyFirst => self.attrs.max_combo = 1,
+                FirstTwoCombos::OnlySecond if self.idx == 1 => self.attrs.max_combo = 1,
+                FirstTwoCombos::Both if self.idx == 0 => self.attrs.max_combo = 1,
+                FirstTwoCombos::Both if self.idx == 1 => self.attrs.max_combo = 2,
+                _ => {}
+            }
         }
 
         self.idx += 1;
@@ -211,11 +232,43 @@ impl Iterator for TaikoGradualDifficulty {
     fn nth(&mut self, n: usize) -> Option<Self::Item> {
         let mut take = n.min(self.len().saturating_sub(1));
 
-        // The first two notes have no difficulty object
+        // The first two notes have no difficulty object but might add to combo
-        if self.idx < 2 && take > 0 {
+        match (take, self.idx) {
-            let skipped = take.min(2);
+            (_, 2..) | (0, _) => {}
-            take -= skipped;
+            (1, 0) => {
-            self.idx += skipped;
+                take -= 1;
+                self.idx += 1;
+
+                match self.first_combos {
+                    FirstTwoCombos::None => {}
+                    FirstTwoCombos::OnlyFirst => self.attrs.max_combo = 1,
+                    FirstTwoCombos::OnlySecond => {}
+                    FirstTwoCombos::Both => self.attrs.max_combo = 1,
+                }
+            }
+            (_, 0) => {
+                take -= 2;
+                self.idx += 2;
+
+                match self.first_combos {
+                    FirstTwoCombos::None => {}
+                    FirstTwoCombos::OnlyFirst => self.attrs.max_combo = 1,
+                    FirstTwoCombos::OnlySecond => self.attrs.max_combo = 1,
+                    FirstTwoCombos::Both => self.attrs.max_combo = 2,
+                }
+            }
+            (_, 1) => {
+                take -= 1;
+                self.idx += 1;
+
+                match self.first_combos {
+                    FirstTwoCombos::None => {}
+                    FirstTwoCombos::OnlyFirst => self.attrs.max_combo = 1,
+                    FirstTwoCombos::OnlySecond => self.attrs.max_combo = 1,
+                    FirstTwoCombos::Both => self.attrs.max_combo = 2,
+                }
+            }
+            _ => unreachable!(),
         }
 
         for _ in 0..take {

src/taiko/mod.rs

@@ -227,26 +227,19 @@ fn calculate_skills(params: TaikoStars<'_>) -> (Peaks, usize) {
         is_convert: _,
     } = params;
 
-    let mut take = passed_objects.unwrap_or(map.hit_objects.len());
+    let take = passed_objects.unwrap_or(map.hit_objects.len());
     let clock_rate = clock_rate.unwrap_or_else(|| mods.clock_rate());
 
     let mut peaks = Peaks::new();
     let mut max_combo = 0;
+    let mut n_diff_objects = 0;
 
-    let mut diff_objects = ObjectLists::default();
+    let mut diff_objects = ObjectLists::with_capacity(map.hit_objects.len().saturating_sub(2));
 
     map.taiko_objects()
-        .take_while(|(h, _)| {
+        .inspect(|(h, _)| {
-            if h.is_hit {
+            n_diff_objects += (max_combo < take) as usize;
-                if take == 0 {
+            max_combo += (max_combo < take && h.is_hit) as usize;
-                    return false;
-                }
-
-                max_combo += 1;
-                take -= 1;
-            }
-
-            true
         })
         .skip(2)
         .zip(map.hit_objects.iter().skip(1))
@@ -276,9 +269,14 @@ fn calculate_skills(params: TaikoStars<'_>) -> (Peaks, usize) {
             diff_objects.all.push(Rc::new(RefCell::new(diff_obj)));
         });
 
+    map.hit_objects
+        .iter()
+        .take(2)
+        .for_each(|h| n_diff_objects = n_diff_objects.saturating_sub(h.is_circle() as usize));
+
     ColourDifficultyPreprocessor::process_and_assign(&mut diff_objects);
 
-    for hit_object in diff_objects.all.iter() {
+    for hit_object in diff_objects.all.iter().take(n_diff_objects) {
         peaks.process(&hit_object.borrow(), &diff_objects);
     }
 

tests/gradual_taiko.rs

@@ -4,8 +4,11 @@
 ))]
 
 use rosu_pp::{
-    taiko::{TaikoGradualDifficulty, TaikoGradualPerformance, TaikoScoreState},
+    taiko::{
-    Beatmap, TaikoPP, TaikoStars,
+        TaikoGradualDifficulty, TaikoGradualPerformance, TaikoOwnedGradualPerformance,
+        TaikoScoreState,
+    },
+    Beatmap, TaikoPP,
 };
 
 mod common;
@@ -18,62 +21,6 @@ fn empty_map() {
     assert!(attrs.next().is_none());
 }
 
-#[test]
-fn iter_end_eq_regular() {
-    let map = test_map!(Taiko);
-    let regular = TaikoStars::new(&map).calculate();
-
-    let iter_end = TaikoGradualDifficulty::new(&map, 0)
-        .last()
-        .expect("empty iter");
-
-    assert_eq!(regular, iter_end);
-}
-
-#[test]
-fn correct_empty() {
-    let map = test_map!(Taiko);
-    let mut gradual = TaikoGradualPerformance::new(&map, 0);
-    let state = TaikoScoreState::default();
-
-    let first_attrs = gradual.nth(state.clone(), usize::MAX);
-
-    assert!(first_attrs.is_some());
-    assert!(gradual.next(state).is_none());
-}
-
-#[test]
-fn next_and_next_n() {
-    let map = test_map!(Taiko);
-    let state = TaikoScoreState::default();
-
-    let mut gradual1 = TaikoGradualPerformance::new(&map, 0);
-    let mut gradual2 = TaikoGradualPerformance::new(&map, 0);
-
-    for _ in 0..50 {
-        let _ = gradual1.next(state.clone());
-        let _ = gradual2.next(state.clone());
-    }
-
-    let n = 200;
-
-    for _ in 1..n {
-        let _ = gradual1.next(state.clone());
-    }
-
-    let state = TaikoScoreState {
-        max_combo: 246,
-        n300: 200,
-        n100: 40,
-        n_misses: 6,
-    };
-
-    let next = gradual1.next(state.clone());
-    let next_n = gradual2.nth(state, n - 1);
-
-    assert_eq!(next_n, next);
-}
-
 #[test]
 fn gradual_end_eq_regular() {
     let map = test_map!(Taiko);
@@ -87,27 +34,57 @@ fn gradual_end_eq_regular() {
         n_misses: 0,
     };
 
-    let gradual_end = gradual.nth(state, usize::MAX).unwrap();
+    let gradual_end = gradual.last(state).unwrap();
 
     assert_eq!(regular, gradual_end);
 }
 
 #[test]
-fn gradual_eq_regular_passed() {
+fn gradual_complete_next() {
     let map = test_map!(Taiko);
-    let n = 250;
+    let mods = 88; // HDHRDT
 
-    let regular = TaikoPP::new(&map).passed_objects(n).calculate();
+    let mut gradual = TaikoGradualPerformance::new(map, mods);
-    let mut gradual = TaikoGradualPerformance::new(&map, 0);
+    let mut gradual_2nd = TaikoGradualPerformance::new(map, mods);
+    let mut gradual_3rd = TaikoGradualPerformance::new(map, mods);
+    let mut gradual_owned = TaikoOwnedGradualPerformance::new(map.to_owned(), mods);
 
-    let state = TaikoScoreState {
+    let mut state = TaikoScoreState::default();
-        max_combo: 250,
-        n300: 250,
-        n100: 0,
-        n_misses: 0,
-    };
 
-    let gradual = gradual.nth(state, n - 1).unwrap();
+    for i in 1.. {
+        state.n_misses += 1;
 
-    assert_eq!(regular, gradual);
+        let Some(next_gradual) = gradual.next(state.clone()) else {
+            assert_eq!(i, map.n_circles as usize + 1);
+            assert!(gradual_2nd.last(state.clone()).is_some() || map.hit_objects.len() % 2 == 0);
+            assert!(gradual_3rd.last(state.clone()).is_some() || map.hit_objects.len() % 3 == 0);
+            assert!(gradual_owned.next(state.clone()).is_none());
+            break;
+        };
+
+        if i % 2 == 0 {
+            let next_gradual_2nd = gradual_2nd.nth(state.clone(), 1).unwrap();
+            assert_eq!(next_gradual, next_gradual_2nd, "i={i}");
+        }
+
+        if i % 3 == 0 {
+            let next_gradual_3rd = gradual_3rd.nth(state.clone(), 2).unwrap();
+            assert_eq!(next_gradual, next_gradual_3rd, "i={i}");
+        }
+
+        let next_gradual_owned = gradual_owned.next(state.clone()).unwrap();
+
+        let mut regular_calc = TaikoPP::new(&map)
+            .mods(mods)
+            .passed_objects(i)
+            .state(state.clone());
+
+        let regular_state = regular_calc.generate_state();
+        assert_eq!(state, regular_state);
+
+        let regular = regular_calc.calculate();
+
+        assert_eq!(next_gradual, next_gradual_owned, "i={i}");
+        assert_eq!(next_gradual, regular, "i={i}");
+    }
 }