6 files changed, 102 insertions, 115 deletions

diff --git a/src/core/hle/kernel/address_arbiter.cpp b/src/core/hle/kernel/address_arbiter.cpp
index 37eec4c84..b5a0cc3a3 100644
--- a/src/core/hle/kernel/address_arbiter.cpp
+++ b/src/core/hle/kernel/address_arbiter.cpp
@@ -79,8 +79,6 @@ ResultCode AddressArbiter::ArbitrateAddress(ArbitrationType type, VAddr address,
                           ErrorSummary::WrongArgument, ErrorLevel::Usage);
     }
 
-    HLE::Reschedule(__func__);
-
     // The calls that use a timeout seem to always return a Timeout error even if they did not put
     // the thread to sleep
     if (type == ArbitrationType::WaitIfLessThanWithTimeout ||
diff --git a/src/core/hle/kernel/kernel.cpp b/src/core/hle/kernel/kernel.cpp
index 0c8752670..209d35270 100644
--- a/src/core/hle/kernel/kernel.cpp
+++ b/src/core/hle/kernel/kernel.cpp
@@ -3,6 +3,7 @@
 // Refer to the license.txt file included.
 
 #include <algorithm>
+#include <boost/range/algorithm_ext/erase.hpp>
 #include "common/assert.h"
 #include "common/logging/log.h"
 #include "core/hle/config_mem.h"
@@ -31,13 +32,60 @@ void WaitObject::RemoveWaitingThread(Thread* thread) {
         waiting_threads.erase(itr);
 }
 
-void WaitObject::WakeupAllWaitingThreads() {
-    for (auto thread : waiting_threads)
-        thread->ResumeFromWait();
+SharedPtr<Thread> WaitObject::GetHighestPriorityReadyThread() {
+    // Remove the threads that are ready or already running from our waitlist
+    boost::range::remove_erase_if(waiting_threads, [](const SharedPtr<Thread>& thread) {
+        return thread->status == THREADSTATUS_RUNNING || thread->status == THREADSTATUS_READY;
+    });
+
+    // TODO(Subv): This call should be performed inside the loop below to check if an object can be
+    // acquired by a particular thread. This is useful for things like recursive locking of Mutexes.
+    if (ShouldWait())
+        return nullptr;
+
+    Thread* candidate = nullptr;
+    s32 candidate_priority = THREADPRIO_LOWEST + 1;
+
+    for (const auto& thread : waiting_threads) {
+        if (thread->current_priority >= candidate_priority)
+            continue;
 
-    waiting_threads.clear();
+        bool ready_to_run =
+            std::none_of(thread->wait_objects.begin(), thread->wait_objects.end(),
+                         [](const SharedPtr<WaitObject>& object) { return object->ShouldWait(); });
+        if (ready_to_run) {
+            candidate = thread.get();
+            candidate_priority = thread->current_priority;
+        }
+    }
+
+    return candidate;
+}
 
-    HLE::Reschedule(__func__);
+void WaitObject::WakeupAllWaitingThreads() {
+    while (auto thread = GetHighestPriorityReadyThread()) {
+        if (!thread->IsSleepingOnWaitAll()) {
+            Acquire();
+            // Set the output index of the WaitSynchronizationN call to the index of this object.
+            if (thread->wait_set_output) {
+                thread->SetWaitSynchronizationOutput(thread->GetWaitObjectIndex(this));
+                thread->wait_set_output = false;
+            }
+        } else {
+            for (auto& object : thread->wait_objects) {
+                object->Acquire();
+                object->RemoveWaitingThread(thread.get());
+            }
+            // Note: This case doesn't update the output index of WaitSynchronizationN.
+            // Clear the thread's waitlist
+            thread->wait_objects.clear();
+        }
+
+        thread->SetWaitSynchronizationResult(RESULT_SUCCESS);
+        thread->ResumeFromWait();
+        // Note: Removing the thread from the object's waitlist will be
+        // done by GetHighestPriorityReadyThread.
+    }
 }
 
 const std::vector<SharedPtr<Thread>>& WaitObject::GetWaitingThreads() const {
diff --git a/src/core/hle/kernel/kernel.h b/src/core/hle/kernel/kernel.h
index 0b811c5a7..1adcf6c71 100644
--- a/src/core/hle/kernel/kernel.h
+++ b/src/core/hle/kernel/kernel.h
@@ -151,9 +151,15 @@ public:
      */
     void RemoveWaitingThread(Thread* thread);
 
-    /// Wake up all threads waiting on this object
+    /**
+     * Wake up all threads waiting on this object that can be awoken, in priority order,
+     * and set the synchronization result and output of the thread.
+     */
     void WakeupAllWaitingThreads();
 
+    /// Obtains the highest priority thread that is ready to run from this object's waiting list.
+    SharedPtr<Thread> GetHighestPriorityReadyThread();
+
     /// Get a const reference to the waiting threads list for debug use
     const std::vector<SharedPtr<Thread>>& GetWaitingThreads() const;
 
diff --git a/src/core/hle/kernel/thread.cpp b/src/core/hle/kernel/thread.cpp
index 84d6d24c6..4bbc08516 100644
--- a/src/core/hle/kernel/thread.cpp
+++ b/src/core/hle/kernel/thread.cpp
@@ -120,8 +120,6 @@ void Thread::Stop() {
     u32 tls_slot =
         ((tls_address - Memory::TLS_AREA_VADDR) % Memory::PAGE_SIZE) / Memory::TLS_ENTRY_SIZE;
     Kernel::g_current_process->tls_slots[tls_page].reset(tls_slot);
-
-    HLE::Reschedule(__func__);
 }
 
 Thread* ArbitrateHighestPriorityThread(u32 address) {
@@ -181,50 +179,6 @@ static void PriorityBoostStarvedThreads() {
 }
 
 /**
- * Gets the registers for timeout parameter of the next WaitSynchronization call.
- * @param thread a pointer to the thread that is ready to call WaitSynchronization
- * @returns a tuple of two register pointers to low and high part of the timeout parameter
- */
-static std::tuple<u32*, u32*> GetWaitSynchTimeoutParameterRegister(Thread* thread) {
-    bool thumb_mode = (thread->context.cpsr & TBIT) != 0;
-    u16 thumb_inst = Memory::Read16(thread->context.pc & 0xFFFFFFFE);
-    u32 inst = Memory::Read32(thread->context.pc & 0xFFFFFFFC) & 0x0FFFFFFF;
-
-    if ((thumb_mode && thumb_inst == 0xDF24) || (!thumb_mode && inst == 0x0F000024)) {
-        // svc #0x24 (WaitSynchronization1)
-        return std::make_tuple(&thread->context.cpu_registers[2],
-                               &thread->context.cpu_registers[3]);
-    } else if ((thumb_mode && thumb_inst == 0xDF25) || (!thumb_mode && inst == 0x0F000025)) {
-        // svc #0x25 (WaitSynchronizationN)
-        return std::make_tuple(&thread->context.cpu_registers[0],
-                               &thread->context.cpu_registers[4]);
-    }
-
-    UNREACHABLE();
-}
-
-/**
- * Updates the WaitSynchronization timeout parameter according to the difference
- * between ticks of the last WaitSynchronization call and the incoming one.
- * @param timeout_low a pointer to the register for the low part of the timeout parameter
- * @param timeout_high a pointer to the register for the high part of the timeout parameter
- * @param last_tick tick of the last WaitSynchronization call
- */
-static void UpdateTimeoutParameter(u32* timeout_low, u32* timeout_high, u64 last_tick) {
-    s64 timeout = ((s64)*timeout_high << 32) | *timeout_low;
-
-    if (timeout != -1) {
-        timeout -= cyclesToUs(CoreTiming::GetTicks() - last_tick) * 1000; // in nanoseconds
-
-        if (timeout < 0)
-            timeout = 0;
-
-        *timeout_low = timeout & 0xFFFFFFFF;
-        *timeout_high = timeout >> 32;
-    }
-}
-
-/**
  * Switches the CPU's active thread context to that of the specified thread
  * @param new_thread The thread to switch to
  */
@@ -254,32 +208,6 @@ static void SwitchContext(Thread* new_thread) {
 
         current_thread = new_thread;
 
-        // If the thread was waited by a svcWaitSynch call, step back PC by one instruction to rerun
-        // the SVC when the thread wakes up. This is necessary to ensure that the thread can acquire
-        // the requested wait object(s) before continuing.
-        if (new_thread->waitsynch_waited) {
-            // CPSR flag indicates CPU mode
-            bool thumb_mode = (new_thread->context.cpsr & TBIT) != 0;
-
-            // SVC instruction is 2 bytes for THUMB, 4 bytes for ARM
-            new_thread->context.pc -= thumb_mode ? 2 : 4;
-
-            // Get the register for timeout parameter
-            u32 *timeout_low, *timeout_high;
-            std::tie(timeout_low, timeout_high) = GetWaitSynchTimeoutParameterRegister(new_thread);
-
-            // Update the timeout parameter
-            UpdateTimeoutParameter(timeout_low, timeout_high, new_thread->last_running_ticks);
-        }
-
-        // Clean up the thread's wait_objects, they'll be restored if needed during
-        // the svcWaitSynchronization call
-        for (size_t i = 0; i < new_thread->wait_objects.size(); ++i) {
-            SharedPtr<WaitObject> object = new_thread->wait_objects[i];
-            object->RemoveWaitingThread(new_thread);
-        }
-        new_thread->wait_objects.clear();
-
         ready_queue.remove(new_thread->current_priority, new_thread);
         new_thread->status = THREADSTATUS_RUNNING;
 
@@ -319,17 +247,13 @@ static Thread* PopNextReadyThread() {
 void WaitCurrentThread_Sleep() {
     Thread* thread = GetCurrentThread();
     thread->status = THREADSTATUS_WAIT_SLEEP;
-
-    HLE::Reschedule(__func__);
 }
 
 void WaitCurrentThread_WaitSynchronization(std::vector<SharedPtr<WaitObject>> wait_objects,
-                                           bool wait_set_output, bool wait_all) {
+                                           bool wait_set_output) {
     Thread* thread = GetCurrentThread();
     thread->wait_set_output = wait_set_output;
-    thread->wait_all = wait_all;
     thread->wait_objects = std::move(wait_objects);
-    thread->waitsynch_waited = true;
     thread->status = THREADSTATUS_WAIT_SYNCH;
 }
 
@@ -351,15 +275,15 @@ static void ThreadWakeupCallback(u64 thread_handle, int cycles_late) {
         return;
     }
 
-    thread->waitsynch_waited = false;
-
     if (thread->status == THREADSTATUS_WAIT_SYNCH || thread->status == THREADSTATUS_WAIT_ARB) {
+        thread->wait_set_output = false;
+        // Remove the thread from each of its waiting objects' waitlists
+        for (auto& object : thread->wait_objects)
+            object->RemoveWaitingThread(thread.get());
+        thread->wait_objects.clear();
         thread->SetWaitSynchronizationResult(ResultCode(ErrorDescription::Timeout, ErrorModule::OS,
                                                         ErrorSummary::StatusChanged,
                                                         ErrorLevel::Info));
-
-        if (thread->wait_set_output)
-            thread->SetWaitSynchronizationOutput(-1);
     }
 
     thread->ResumeFromWait();
@@ -399,6 +323,7 @@ void Thread::ResumeFromWait() {
 
     ready_queue.push_back(current_priority, this);
     status = THREADSTATUS_READY;
+    HLE::Reschedule(__func__);
 }
 
 /**
@@ -494,13 +419,11 @@ ResultVal<SharedPtr<Thread>> Thread::Create(std::string name, VAddr entry_point,
     thread->last_running_ticks = CoreTiming::GetTicks();
     thread->processor_id = processor_id;
     thread->wait_set_output = false;
-    thread->wait_all = false;
     thread->wait_objects.clear();
     thread->wait_address = 0;
     thread->name = std::move(name);
     thread->callback_handle = wakeup_callback_handle_table.Create(thread).MoveFrom();
     thread->owner_process = g_current_process;
-    thread->waitsynch_waited = false;
 
     // Find the next available TLS index, and mark it as used
     auto& tls_slots = Kernel::g_current_process->tls_slots;
@@ -555,8 +478,6 @@ ResultVal<SharedPtr<Thread>> Thread::Create(std::string name, VAddr entry_point,
     ready_queue.push_back(thread->current_priority, thread.get());
     thread->status = THREADSTATUS_READY;
 
-    HLE::Reschedule(__func__);
-
     return MakeResult<SharedPtr<Thread>>(std::move(thread));
 }
 
@@ -619,14 +540,6 @@ void Reschedule() {
 
     HLE::DoneRescheduling();
 
-    // Don't bother switching to the same thread.
-    // But if the thread was waiting on objects, we still need to switch it
-    // to perform PC modification, change state to RUNNING, etc.
-    // This occurs in the case when an object the thread is waiting on immediately wakes up
-    // the current thread before Reschedule() is called.
-    if (next == cur && (next == nullptr || next->waitsynch_waited == false))
-        return;
-
     if (cur && next) {
         LOG_TRACE(Kernel, "context switch %u -> %u", cur->GetObjectId(), next->GetObjectId());
     } else if (cur) {
diff --git a/src/core/hle/kernel/thread.h b/src/core/hle/kernel/thread.h
index e0ffcea8a..238359fc5 100644
--- a/src/core/hle/kernel/thread.h
+++ b/src/core/hle/kernel/thread.h
@@ -5,7 +5,9 @@
 #pragma once
 
 #include <string>
+#include <unordered_map>
 #include <vector>
+#include <boost/container/flat_map.hpp>
 #include <boost/container/flat_set.hpp>
 #include "common/common_types.h"
 #include "core/core.h"
@@ -125,6 +127,16 @@ public:
     void SetWaitSynchronizationOutput(s32 output);
 
     /**
+     * Retrieves the index that this particular object occupies in the list of objects
+     * that the thread passed to WaitSynchronizationN.
+     * It is used to set the output value of WaitSynchronizationN when the thread is awakened.
+     * @param object Object to query the index of.
+     */
+    s32 GetWaitObjectIndex(const WaitObject* object) const {
+        return wait_objects_index.at(object->GetObjectId());
+    }
+
+    /**
      * Stops a thread, invalidating it from further use
      */
     void Stop();
@@ -137,6 +149,15 @@ public:
         return tls_address;
     }
 
+    /**
+     * Returns whether this thread is waiting for all the objects in
+     * its wait list to become ready, as a result of a WaitSynchronizationN call
+     * with wait_all = true, or a ReplyAndReceive call.
+     */
+    bool IsSleepingOnWaitAll() const {
+        return !wait_objects.empty();
+    }
+
     Core::ThreadContext context;
 
     u32 thread_id;
@@ -154,16 +175,22 @@ public:
 
     VAddr tls_address; ///< Virtual address of the Thread Local Storage of the thread
 
-    bool waitsynch_waited; ///< Set to true if the last svcWaitSynch call caused the thread to wait
-
     /// Mutexes currently held by this thread, which will be released when it exits.
     boost::container::flat_set<SharedPtr<Mutex>> held_mutexes;
 
-    SharedPtr<Process> owner_process;                ///< Process that owns this thread
-    std::vector<SharedPtr<WaitObject>> wait_objects; ///< Objects that the thread is waiting on
-    VAddr wait_address;   ///< If waiting on an AddressArbiter, this is the arbitration address
-    bool wait_all;        ///< True if the thread is waiting on all objects before resuming
-    bool wait_set_output; ///< True if the output parameter should be set on thread wakeup
+    SharedPtr<Process> owner_process; ///< Process that owns this thread
+
+    /// Objects that the thread is waiting on.
+    /// This is only populated when the thread should wait for all the objects to become ready.
+    std::vector<SharedPtr<WaitObject>> wait_objects;
+
+    /// Mapping of Object ids to their position in the last waitlist that this object waited on.
+    boost::container::flat_map<int, s32> wait_objects_index;
+
+    VAddr wait_address; ///< If waiting on an AddressArbiter, this is the arbitration address
+
+    /// True if the WaitSynchronizationN output parameter should be set on thread wakeup.
+    bool wait_set_output;
 
     std::string name;
 
@@ -215,10 +242,9 @@ void WaitCurrentThread_Sleep();
  * @param wait_objects Kernel objects that we are waiting on
  * @param wait_set_output If true, set the output parameter on thread wakeup (for
  * WaitSynchronizationN only)
- * @param wait_all If true, wait on all objects before resuming (for WaitSynchronizationN only)
  */
 void WaitCurrentThread_WaitSynchronization(std::vector<SharedPtr<WaitObject>> wait_objects,
-                                           bool wait_set_output, bool wait_all);
+                                           bool wait_set_output);
 
 /**
  * Waits the current thread from an ArbitrateAddress call
diff --git a/src/core/hle/kernel/timer.cpp b/src/core/hle/kernel/timer.cpp
index eac181f4e..b50cf520d 100644
--- a/src/core/hle/kernel/timer.cpp
+++ b/src/core/hle/kernel/timer.cpp
@@ -60,14 +60,10 @@ void Timer::Set(s64 initial, s64 interval) {
     u64 initial_microseconds = initial / 1000;
     CoreTiming::ScheduleEvent(usToCycles(initial_microseconds), timer_callback_event_type,
                               callback_handle);
-
-    HLE::Reschedule(__func__);
 }
 
 void Timer::Cancel() {
     CoreTiming::UnscheduleEvent(timer_callback_event_type, callback_handle);
-
-    HLE::Reschedule(__func__);
 }
 
 void Timer::Clear() {