From b3556403aac6af43614e429054a4b6672583b700 Mon Sep 17 00:00:00 2001
From: Desconocidosmh <Desconocidosmh@users.noreply.github.com>
Date: Mon, 19 Aug 2019 16:18:25 +0200
Subject: [PATCH] Make GetNextObject() a virtual method

---
 .../Replays/ManiaAutoGenerator.cs             | 24 ++++++++---------
 .../Replays/TaikoAutoGenerator.cs             | 26 +++++++++++++++----
 osu.Game/Rulesets/Replays/AutoGenerator.cs    |  9 +++++++
 3 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/osu.Game.Rulesets.Mania/Replays/ManiaAutoGenerator.cs b/osu.Game.Rulesets.Mania/Replays/ManiaAutoGenerator.cs
index 43150958d0..b9dd32208e 100644
--- a/osu.Game.Rulesets.Mania/Replays/ManiaAutoGenerator.cs
+++ b/osu.Game.Rulesets.Mania/Replays/ManiaAutoGenerator.cs
@@ -5,7 +5,7 @@ using System.Collections.Generic;
 using System.Linq;
 using osu.Game.Replays;
 using osu.Game.Rulesets.Mania.Beatmaps;
-using osu.Game.Rulesets.Mania.Objects;
+using osu.Game.Rulesets.Objects;
 using osu.Game.Rulesets.Objects.Types;
 using osu.Game.Rulesets.Replays;
 
@@ -81,7 +81,7 @@ namespace osu.Game.Rulesets.Mania.Replays
             for (int i = 0; i < Beatmap.HitObjects.Count; i++)
             {
                 var currentObject = Beatmap.HitObjects[i];
-                var nextObjectInTheSameColumn = getNextObjectInTheSameColumn(i);
+                var nextObjectInTheSameColumn = GetNextObject(i);
 
                 double endTime = (currentObject as IHasEndTime)?.EndTime ?? currentObject.StartTime;
 
@@ -94,19 +94,19 @@ namespace osu.Game.Rulesets.Mania.Replays
 
                 yield return new ReleasePoint { Time = endTime + calculatedDelay, Column = currentObject.Column };
             }
+        }
 
-            ManiaHitObject getNextObjectInTheSameColumn(int currentIndex)
+        protected override HitObject GetNextObject(int currentIndex)
+        {
+            int desiredColumn = Beatmap.HitObjects[currentIndex++].Column;
+
+            for (; currentIndex < Beatmap.HitObjects.Count; currentIndex++)
             {
-                int desiredColumn = Beatmap.HitObjects[currentIndex++].Column;
-
-                for (; currentIndex < Beatmap.HitObjects.Count; currentIndex++)
-                {
-                    if (Beatmap.HitObjects[currentIndex].Column == desiredColumn)
-                        return Beatmap.HitObjects[currentIndex];
-                }
-
-                return null;
+                if (Beatmap.HitObjects[currentIndex].Column == desiredColumn)
+                    return Beatmap.HitObjects[currentIndex];
             }
+
+            return null;
         }
 
         private interface IActionPoint
diff --git a/osu.Game.Rulesets.Taiko/Replays/TaikoAutoGenerator.cs b/osu.Game.Rulesets.Taiko/Replays/TaikoAutoGenerator.cs
index 1d35393de0..67f93a0259 100644
--- a/osu.Game.Rulesets.Taiko/Replays/TaikoAutoGenerator.cs
+++ b/osu.Game.Rulesets.Taiko/Replays/TaikoAutoGenerator.cs
@@ -10,6 +10,7 @@ using osu.Game.Rulesets.Objects.Types;
 using osu.Game.Rulesets.Taiko.Objects;
 using osu.Game.Rulesets.Replays;
 using osu.Game.Rulesets.Taiko.Beatmaps;
+using osu.Game.Rulesets.Objects;
 
 namespace osu.Game.Rulesets.Taiko.Replays
 {
@@ -113,16 +114,17 @@ namespace osu.Game.Rulesets.Taiko.Replays
                 else
                     throw new InvalidOperationException("Unknown hit object type.");
 
-                TaikoHitObject nextHitObject = i < Beatmap.HitObjects.Count - 1 ? Beatmap.HitObjects[i + 1] : null;
+                var nextHitObject = GetNextObject(i);
 
                 bool canDelayKeyUp = nextHitObject == null || nextHitObject.StartTime > endTime + KEY_UP_DELAY;
 
-                if (canDelayKeyUp)
-                    Frames.Add(new TaikoReplayFrame(endTime + KEY_UP_DELAY));
+                double calculatedDelay = canDelayKeyUp ? KEY_UP_DELAY : nextHitObject.StartTime - endTime;
 
-                if (nextHitObject != null)
+                Frames.Add(new TaikoReplayFrame(endTime + calculatedDelay));
+
+                if (i < Beatmap.HitObjects.Count - 1)
                 {
-                    double waitTime = nextHitObject.StartTime - 1000;
+                    double waitTime = Beatmap.HitObjects[i + 1].StartTime - 1000;
                     if (waitTime > endTime)
                         Frames.Add(new TaikoReplayFrame(waitTime));
                 }
@@ -132,5 +134,19 @@ namespace osu.Game.Rulesets.Taiko.Replays
 
             return Replay;
         }
+
+        protected override HitObject GetNextObject(int currentIndex)
+        {
+            Type desiredType = Beatmap.HitObjects[currentIndex++].GetType();
+
+            for (; currentIndex < Beatmap.HitObjects.Count; currentIndex++)
+            {
+                var currentObj = Beatmap.HitObjects[currentIndex];
+                if (currentObj.GetType().Equals(desiredType) || currentObj is DrumRoll)
+                    return Beatmap.HitObjects[currentIndex];
+            }
+
+            return null;
+        }
     }
 }
diff --git a/osu.Game/Rulesets/Replays/AutoGenerator.cs b/osu.Game/Rulesets/Replays/AutoGenerator.cs
index 1d4cdbf04c..3319f30a6f 100644
--- a/osu.Game/Rulesets/Replays/AutoGenerator.cs
+++ b/osu.Game/Rulesets/Replays/AutoGenerator.cs
@@ -3,6 +3,7 @@
 
 using osu.Game.Beatmaps;
 using osu.Game.Replays;
+using osu.Game.Rulesets.Objects;
 
 namespace osu.Game.Rulesets.Replays
 {
@@ -34,5 +35,13 @@ namespace osu.Game.Rulesets.Replays
         protected const double KEY_UP_DELAY = 50;
 
         #endregion
+
+        protected virtual HitObject GetNextObject(int currentIndex)
+        {
+            if (currentIndex >= Beatmap.HitObjects.Count - 1)
+                return null;
+
+            return Beatmap.HitObjects[currentIndex + 1];
+        }
     }
 }