From b5d1e447821eb21158669e0ef1d24d630602f1fe Mon Sep 17 00:00:00 2001
From: Fernando Sahmkow <fsahmkow27@gmail.com>
Date: Fri, 29 Mar 2019 17:13:00 -0400
Subject: Add PrepareReschedule where required.

---
 src/core/hle/kernel/address_arbiter.cpp | 28 ++++++++++++----------------
 src/core/hle/kernel/mutex.cpp           |  2 ++
 src/core/hle/kernel/wait_object.cpp     |  4 ++++
 3 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/src/core/hle/kernel/address_arbiter.cpp b/src/core/hle/kernel/address_arbiter.cpp
index c8842410b..77f7bb451 100644
--- a/src/core/hle/kernel/address_arbiter.cpp
+++ b/src/core/hle/kernel/address_arbiter.cpp
@@ -22,6 +22,8 @@ namespace Kernel {
 namespace {
 // Wake up num_to_wake (or all) threads in a vector.
 void WakeThreads(const std::vector<SharedPtr<Thread>>& waiting_threads, s32 num_to_wake) {
+
+    auto& system = Core::System::GetInstance();
     // Only process up to 'target' threads, unless 'target' is <= 0, in which case process
     // them all.
     std::size_t last = waiting_threads.size();
@@ -35,6 +37,8 @@ void WakeThreads(const std::vector<SharedPtr<Thread>>& waiting_threads, s32 num_
         waiting_threads[i]->SetWaitSynchronizationResult(RESULT_SUCCESS);
         waiting_threads[i]->SetArbiterWaitAddress(0);
         waiting_threads[i]->ResumeFromWait();
+        if (waiting_threads[i]->GetProcessorID() >= 0)
+            system.CpuCore(waiting_threads[i]->GetProcessorID()).PrepareReschedule();
     }
 }
 } // Anonymous namespace
@@ -174,25 +178,17 @@ ResultCode AddressArbiter::WaitForAddressImpl(VAddr address, s64 timeout) {
 }
 
 std::vector<SharedPtr<Thread>> AddressArbiter::GetThreadsWaitingOnAddress(VAddr address) const {
-    const auto RetrieveWaitingThreads = [this](std::size_t core_index,
-                                               std::vector<SharedPtr<Thread>>& waiting_threads,
-                                               VAddr arb_addr) {
-        const auto& scheduler = system.Scheduler(core_index);
-        const auto& thread_list = scheduler.GetThreadList();
-
-        for (const auto& thread : thread_list) {
-            if (thread->GetArbiterWaitAddress() == arb_addr) {
-                waiting_threads.push_back(thread);
-            }
-        }
-    };
 
     // Retrieve all threads that are waiting for this address.
     std::vector<SharedPtr<Thread>> threads;
-    RetrieveWaitingThreads(0, threads, address);
-    RetrieveWaitingThreads(1, threads, address);
-    RetrieveWaitingThreads(2, threads, address);
-    RetrieveWaitingThreads(3, threads, address);
+    const auto& scheduler = system.GlobalScheduler();
+    const auto& thread_list = scheduler.GetThreadList();
+
+    for (const auto& thread : thread_list) {
+        if (thread->GetArbiterWaitAddress() == address) {
+            threads.push_back(thread);
+        }
+    }
 
     // Sort them by priority, such that the highest priority ones come first.
     std::sort(threads.begin(), threads.end(),
diff --git a/src/core/hle/kernel/mutex.cpp b/src/core/hle/kernel/mutex.cpp
index 98e87313b..57f2d8bf3 100644
--- a/src/core/hle/kernel/mutex.cpp
+++ b/src/core/hle/kernel/mutex.cpp
@@ -140,6 +140,8 @@ ResultCode Mutex::Release(VAddr address) {
     thread->SetMutexWaitAddress(0);
     thread->SetWaitHandle(0);
 
+    Core::System::GetInstance().PrepareReschedule();
+
     return RESULT_SUCCESS;
 }
 } // namespace Kernel
diff --git a/src/core/hle/kernel/wait_object.cpp b/src/core/hle/kernel/wait_object.cpp
index 0e96ba872..e035a67e9 100644
--- a/src/core/hle/kernel/wait_object.cpp
+++ b/src/core/hle/kernel/wait_object.cpp
@@ -6,6 +6,8 @@
 #include "common/assert.h"
 #include "common/common_types.h"
 #include "common/logging/log.h"
+#include "core/core.h"
+#include "core/core_cpu.h"
 #include "core/hle/kernel/object.h"
 #include "core/hle/kernel/process.h"
 #include "core/hle/kernel/thread.h"
@@ -95,6 +97,8 @@ void WaitObject::WakeupWaitingThread(SharedPtr<Thread> thread) {
     }
     if (resume) {
         thread->ResumeFromWait();
+        if (thread->GetProcessorID() >= 0)
+            Core::System::GetInstance().CpuCore(thread->GetProcessorID()).PrepareReschedule();
     }
 }
 
-- 
cgit v1.2.3