Fix osu!mania replays actuating incorrect keys when multiple stages are involved

osu.Game.Rulesets.Mania.Tests/ManiaLegacyReplayTest.cs

@@ -0,0 +1,51 @@
+// Copyright (c) ppy Pty Ltd <contact@ppy.sh>. Licensed under the MIT Licence.
+// See the LICENCE file in the repository root for full licence text.
+
+using NUnit.Framework;
+using osu.Game.Rulesets.Mania.Beatmaps;
+using osu.Game.Rulesets.Mania.Replays;
+
+namespace osu.Game.Rulesets.Mania.Tests
+{
+    [TestFixture]
+    public class ManiaLegacyReplayTest
+    {
+        [TestCase(ManiaAction.Key1)]
+        [TestCase(ManiaAction.Key1, ManiaAction.Key2)]
+        [TestCase(ManiaAction.Special1)]
+        [TestCase(ManiaAction.Key8)]
+        public void TestEncodeDecodeSingleStage(params ManiaAction[] actions)
+        {
+            var beatmap = new ManiaBeatmap(new StageDefinition { Columns = 9 });
+
+            var frame = new ManiaReplayFrame(0, actions);
+            var legacyFrame = frame.ToLegacy(beatmap);
+
+            var decodedFrame = new ManiaReplayFrame();
+            decodedFrame.FromLegacy(legacyFrame, beatmap);
+
+            Assert.That(decodedFrame.Actions, Is.EquivalentTo(frame.Actions));
+        }
+
+        [TestCase(ManiaAction.Key1)]
+        [TestCase(ManiaAction.Key1, ManiaAction.Key2)]
+        [TestCase(ManiaAction.Special1)]
+        [TestCase(ManiaAction.Special2)]
+        [TestCase(ManiaAction.Special1, ManiaAction.Special2)]
+        [TestCase(ManiaAction.Special1, ManiaAction.Key5)]
+        [TestCase(ManiaAction.Key8)]
+        public void TestEncodeDecodeDualStage(params ManiaAction[] actions)
+        {
+            var beatmap = new ManiaBeatmap(new StageDefinition { Columns = 5 });
+            beatmap.Stages.Add(new StageDefinition { Columns = 5 });
+
+            var frame = new ManiaReplayFrame(0, actions);
+            var legacyFrame = frame.ToLegacy(beatmap);
+
+            var decodedFrame = new ManiaReplayFrame();
+            decodedFrame.FromLegacy(legacyFrame, beatmap);
+
+            Assert.That(decodedFrame.Actions, Is.EquivalentTo(frame.Actions));
+        }
+    }
+}

osu.Game.Rulesets.Mania/Replays/ManiaReplayFrame.cs

@@ -1,8 +1,8 @@
 // Copyright (c) ppy Pty Ltd <contact@ppy.sh>. Licensed under the MIT Licence.
 // See the LICENCE file in the repository root for full licence text.
 
+using System;
 using System.Collections.Generic;
-using System.Linq;
 using osu.Game.Beatmaps;
 using osu.Game.Replays.Legacy;
 using osu.Game.Rulesets.Mania.Beatmaps;
@@ -37,7 +37,7 @@ public void FromLegacy(LegacyReplayFrame legacyFrame, IBeatmap beatmap, ReplayFr
 
             while (activeColumns > 0)
             {
-                var isSpecial = maniaBeatmap.Stages.First().IsSpecialColumn(counter);
+                bool isSpecial = isColumnAtIndexSpecial(maniaBeatmap, counter);
 
                 if ((activeColumns & 1) > 0)
                     Actions.Add(isSpecial ? specialAction : normalAction);
@@ -58,33 +58,90 @@ public LegacyReplayFrame ToLegacy(IBeatmap beatmap)
 
             int keys = 0;
 
-            var specialColumns = new List<int>();
-
-            for (int i = 0; i < maniaBeatmap.TotalColumns; i++)
-            {
-                if (maniaBeatmap.Stages.First().IsSpecialColumn(i))
-                    specialColumns.Add(i);
-            }
-
             foreach (var action in Actions)
             {
                 switch (action)
                 {
                     case ManiaAction.Special1:
-                        keys |= 1 << specialColumns[0];
+                        keys |= 1 << getSpecialColumnIndex(maniaBeatmap, 0);
                         break;
 
                     case ManiaAction.Special2:
-                        keys |= 1 << specialColumns[1];
+                        keys |= 1 << getSpecialColumnIndex(maniaBeatmap, 1);
                         break;
 
                     default:
-                        keys |= 1 << (action - ManiaAction.Key1);
+                        // the index in lazer, which doesn't include special keys.
+                        int nonSpecialKeyIndex = action - ManiaAction.Key1;
+
+                        int overallIndex = 0;
+
+                        // iterate to find the index including special keys.
+                        while (true)
+                        {
+                            if (!isColumnAtIndexSpecial(maniaBeatmap, overallIndex))
+                            {
+                                // found a non-special column we could use.
+                                if (nonSpecialKeyIndex == 0)
+                                    break;
+
+                                // found a non-special column but not ours.
+                                nonSpecialKeyIndex--;
+                            }
+
+                            overallIndex++;
+                        }
+
+                        keys |= 1 << overallIndex;
                         break;
                 }
             }
 
             return new LegacyReplayFrame(Time, keys, null, ReplayButtonState.None);
         }
+
+        /// <summary>
+        /// Find the overall index (across all stages) for a specified special key.
+        /// </summary>
+        /// <param name="maniaBeatmap">The beatmap.</param>
+        /// <param name="specialOffset">The special key offset (0 is S1).</param>
+        /// <returns>The overall index for the special column.</returns>
+        private int getSpecialColumnIndex(ManiaBeatmap maniaBeatmap, int specialOffset)
+        {
+            for (int i = 0; i < maniaBeatmap.TotalColumns; i++)
+            {
+                if (isColumnAtIndexSpecial(maniaBeatmap, i))
+                {
+                    if (specialOffset == 0)
+                        return i;
+
+                    specialOffset--;
+                }
+            }
+
+            throw new InvalidOperationException("Special key index too high");
+        }
+
+        /// <summary>
+        /// Check whether the column at an overall index (across all stages) is a special column.
+        /// </summary>
+        /// <param name="beatmap">The beatmap.</param>
+        /// <param name="index">The overall index to check.</param>
+        /// <returns></returns>
+        private bool isColumnAtIndexSpecial(ManiaBeatmap beatmap, int index)
+        {
+            foreach (var stage in beatmap.Stages)
+            {
+                for (int stageIndex = 0; stageIndex < stage.Columns; stageIndex++)
+                {
+                    if (index == 0)
+                        return stage.IsSpecialColumn(stageIndex);
+
+                    index--;
+                }
+            }
+
+            return false;
+        }
     }
 }