From a341052692306fed2ad638f675124421a378724a Mon Sep 17 00:00:00 2001
From: Barinade <poco0317@users.noreply.github.com>
Date: Fri, 5 Nov 2021 19:02:39 -0500
Subject: [PATCH] fix musicwheel setting grades for filtered difficulties on
 wheel

---
 src/Etterna/Actor/Menus/MusicWheelItem.cpp | 33 ++++++++++++++++------
 1 file changed, 24 insertions(+), 9 deletions(-)

diff --git a/src/Etterna/Actor/Menus/MusicWheelItem.cpp b/src/Etterna/Actor/Menus/MusicWheelItem.cpp
index 21108eb84a..ea39b7d32d 100644
--- a/src/Etterna/Actor/Menus/MusicWheelItem.cpp
+++ b/src/Etterna/Actor/Menus/MusicWheelItem.cpp
@@ -314,23 +314,38 @@ MusicWheelItem::RefreshGrades()
 	Difficulty dcBest = Difficulty_Invalid;
 	if (pWID->m_pSong != nullptr) {
 		bool hasCurrentStyleSteps = false;
+		auto& allSteps = pWID->m_pSong->GetChartsMatchingFilter();
+		std::unordered_map<Difficulty, std::vector<Steps*>> difficultyToSteps{
+			NUM_Difficulty
+		};
+		for (auto& s : allSteps) {
+			auto d = s->GetDifficulty();
+			if (difficultyToSteps.count(d) == 0u) {
+				std::vector<Steps*> v;
+				difficultyToSteps[d] = v;
+			}
+			difficultyToSteps[d].push_back(s);
+		}
 		FOREACH_ENUM_N(Difficulty, 6, i)
 		{
-			Steps* pSteps =
-			  SongUtil::GetStepsByDifficulty(pWID->m_pSong, st, i);
-			if (pSteps != nullptr) {
-				hasCurrentStyleSteps = true;
-				Grade dcg = SCOREMAN->GetBestGradeFor(pSteps->GetChartKey());
-				if (gradeBest >= dcg) {
-					dcBest = i;
-					gradeBest = dcg;
+			if (difficultyToSteps.count(i) == 0u)
+				continue;
+			auto& stepslist = difficultyToSteps[i];
+			for (auto& s : stepslist) {
+				if (s->m_StepsType == st) {
+					hasCurrentStyleSteps = true;
+					Grade dcg =
+					  SCOREMAN->GetBestGradeFor(s->GetChartKey());
+					if (gradeBest >= dcg) {
+						dcBest = i;
+						gradeBest = dcg;
+					}
 				}
 			}
 		}
 		// If no grade was found for the current style/stepstype
 		if (!hasCurrentStyleSteps) {
 			// Get the best grade among all steps
-			auto& allSteps = pWID->m_pSong->GetAllSteps();
 			for (auto& stepsPtr : allSteps) {
 				if (stepsPtr->m_StepsType ==
 					st) // Skip already checked steps of type st