4 files changed, 91 insertions, 17 deletions

diff --git a/src/core/hle/kernel/thread.cpp b/src/core/hle/kernel/thread.cpp
index 11f7d2127..2614a260c 100644
--- a/src/core/hle/kernel/thread.cpp
+++ b/src/core/hle/kernel/thread.cpp
@@ -247,12 +247,15 @@ static void ThreadWakeupCallback(u64 thread_handle, int cycles_late) {
 
     if (thread->status == THREADSTATUS_WAIT_SYNCH_ANY ||
         thread->status == THREADSTATUS_WAIT_SYNCH_ALL || thread->status == THREADSTATUS_WAIT_ARB) {
-        thread->wait_set_output = false;
+
+        // Invoke the wakeup callback before clearing the wait objects
+        if (thread->wakeup_callback)
+            thread->wakeup_callback(ThreadWakeupReason::Timeout, thread, nullptr);
+
         // 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(RESULT_TIMEOUT);
     }
 
     thread->ResumeFromWait();
@@ -278,6 +281,9 @@ void Thread::ResumeFromWait() {
         break;
 
     case THREADSTATUS_READY:
+        // The thread's wakeup callback must have already been cleared when the thread was first
+        // awoken.
+        ASSERT(wakeup_callback == nullptr);
         // If the thread is waiting on multiple wait objects, it might be awoken more than once
         // before actually resuming. We can ignore subsequent wakeups if the thread status has
         // already been set to THREADSTATUS_READY.
@@ -293,6 +299,8 @@ void Thread::ResumeFromWait() {
         return;
     }
 
+    wakeup_callback = nullptr;
+
     ready_queue.push_back(current_priority, this);
     status = THREADSTATUS_READY;
     Core::System::GetInstance().PrepareReschedule();
@@ -395,7 +403,6 @@ ResultVal<SharedPtr<Thread>> Thread::Create(std::string name, VAddr entry_point,
     thread->nominal_priority = thread->current_priority = priority;
     thread->last_running_ticks = CoreTiming::GetTicks();
     thread->processor_id = processor_id;
-    thread->wait_set_output = false;
     thread->wait_objects.clear();
     thread->wait_address = 0;
     thread->name = std::move(name);
diff --git a/src/core/hle/kernel/thread.h b/src/core/hle/kernel/thread.h
index f02e1d43a..4679c2022 100644
--- a/src/core/hle/kernel/thread.h
+++ b/src/core/hle/kernel/thread.h
@@ -41,6 +41,11 @@ enum ThreadStatus {
     THREADSTATUS_DEAD            ///< Run to completion, or forcefully terminated
 };
 
+enum class ThreadWakeupReason {
+    Signal, // The thread was woken up by WakeupAllWaitingThreads due to an object signal.
+    Timeout // The thread was woken up due to a wait timeout.
+};
+
 namespace Kernel {
 
 class Mutex;
@@ -199,14 +204,18 @@ public:
 
     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;
 
     /// Handle used as userdata to reference this object when inserting into the CoreTiming queue.
     Handle callback_handle;
 
+    using WakeupCallback = void(ThreadWakeupReason reason, SharedPtr<Thread> thread,
+                                SharedPtr<WaitObject> object);
+    // Callback that will be invoked when the thread is resumed from a waiting state. If the thread
+    // was waiting via WaitSynchronizationN then the object will be the last object that became
+    // available. In case of a timeout, the object will be nullptr.
+    std::function<WakeupCallback> wakeup_callback;
+
 private:
     Thread();
     ~Thread() override;
diff --git a/src/core/hle/kernel/wait_object.cpp b/src/core/hle/kernel/wait_object.cpp
index 56fdd977f..469554908 100644
--- a/src/core/hle/kernel/wait_object.cpp
+++ b/src/core/hle/kernel/wait_object.cpp
@@ -71,23 +71,20 @@ void WaitObject::WakeupAllWaitingThreads() {
     while (auto thread = GetHighestPriorityReadyThread()) {
         if (!thread->IsSleepingOnWaitAll()) {
             Acquire(thread.get());
-            // 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(thread.get());
             }
-            // Note: This case doesn't update the output index of WaitSynchronizationN.
         }
 
+        // Invoke the wakeup callback before clearing the wait objects
+        if (thread->wakeup_callback)
+            thread->wakeup_callback(ThreadWakeupReason::Signal, thread, this);
+
         for (auto& object : thread->wait_objects)
             object->RemoveWaitingThread(thread.get());
         thread->wait_objects.clear();
 
-        thread->SetWaitSynchronizationResult(RESULT_SUCCESS);
         thread->ResumeFromWait();
     }
 }
diff --git a/src/core/hle/svc.cpp b/src/core/hle/svc.cpp
index 41c82c922..fefd50805 100644
--- a/src/core/hle/svc.cpp
+++ b/src/core/hle/svc.cpp
@@ -271,6 +271,24 @@ static ResultCode WaitSynchronization1(Kernel::Handle handle, s64 nano_seconds)
         // Create an event to wake the thread up after the specified nanosecond delay has passed
         thread->WakeAfterDelay(nano_seconds);
 
+        thread->wakeup_callback = [](ThreadWakeupReason reason,
+                                     Kernel::SharedPtr<Kernel::Thread> thread,
+                                     Kernel::SharedPtr<Kernel::WaitObject> object) {
+
+            ASSERT(thread->status == THREADSTATUS_WAIT_SYNCH_ANY);
+
+            if (reason == ThreadWakeupReason::Timeout) {
+                thread->SetWaitSynchronizationResult(Kernel::RESULT_TIMEOUT);
+                return;
+            }
+
+            ASSERT(reason == ThreadWakeupReason::Signal);
+            thread->SetWaitSynchronizationResult(RESULT_SUCCESS);
+
+            // WaitSynchronization1 doesn't have an output index like WaitSynchronizationN, so we
+            // don't have to do anything else here.
+        };
+
         Core::System::GetInstance().PrepareReschedule();
 
         // Note: The output of this SVC will be set to RESULT_SUCCESS if the thread
@@ -344,6 +362,23 @@ static ResultCode WaitSynchronizationN(s32* out, Kernel::Handle* handles, s32 ha
         // Create an event to wake the thread up after the specified nanosecond delay has passed
         thread->WakeAfterDelay(nano_seconds);
 
+        thread->wakeup_callback = [](ThreadWakeupReason reason,
+                                     Kernel::SharedPtr<Kernel::Thread> thread,
+                                     Kernel::SharedPtr<Kernel::WaitObject> object) {
+
+            ASSERT(thread->status == THREADSTATUS_WAIT_SYNCH_ALL);
+
+            if (reason == ThreadWakeupReason::Timeout) {
+                thread->SetWaitSynchronizationResult(Kernel::RESULT_TIMEOUT);
+                return;
+            }
+
+            ASSERT(reason == ThreadWakeupReason::Signal);
+
+            thread->SetWaitSynchronizationResult(RESULT_SUCCESS);
+            // The wait_all case does not update the output index.
+        };
+
         Core::System::GetInstance().PrepareReschedule();
 
         // This value gets set to -1 by default in this case, it is not modified after this.
@@ -389,12 +424,28 @@ static ResultCode WaitSynchronizationN(s32* out, Kernel::Handle* handles, s32 ha
         // Create an event to wake the thread up after the specified nanosecond delay has passed
         thread->WakeAfterDelay(nano_seconds);
 
+        thread->wakeup_callback = [](ThreadWakeupReason reason,
+                                     Kernel::SharedPtr<Kernel::Thread> thread,
+                                     Kernel::SharedPtr<Kernel::WaitObject> object) {
+
+            ASSERT(thread->status == THREADSTATUS_WAIT_SYNCH_ANY);
+
+            if (reason == ThreadWakeupReason::Timeout) {
+                thread->SetWaitSynchronizationResult(Kernel::RESULT_TIMEOUT);
+                return;
+            }
+
+            ASSERT(reason == ThreadWakeupReason::Signal);
+
+            thread->SetWaitSynchronizationResult(RESULT_SUCCESS);
+            thread->SetWaitSynchronizationOutput(thread->GetWaitObjectIndex(object.get()));
+        };
+
         Core::System::GetInstance().PrepareReschedule();
 
         // Note: The output of this SVC will be set to RESULT_SUCCESS if the thread resumes due to a
         // signal in one of its wait objects.
         // Otherwise we retain the default value of timeout, and -1 in the out parameter
-        thread->wait_set_output = true;
         *out = -1;
         return Kernel::RESULT_TIMEOUT;
     }
@@ -483,8 +534,6 @@ static ResultCode ReplyAndReceive(s32* index, Kernel::Handle* handles, s32 handl
 
     // No objects were ready to be acquired, prepare to suspend the thread.
 
-    // TODO(Subv): Perform IPC translation upon wakeup.
-
     // Put the thread to sleep
     thread->status = THREADSTATUS_WAIT_SYNCH_ANY;
 
@@ -496,12 +545,24 @@ static ResultCode ReplyAndReceive(s32* index, Kernel::Handle* handles, s32 handl
 
     thread->wait_objects = std::move(objects);
 
+    thread->wakeup_callback = [](ThreadWakeupReason reason,
+                                 Kernel::SharedPtr<Kernel::Thread> thread,
+                                 Kernel::SharedPtr<Kernel::WaitObject> object) {
+
+        ASSERT(thread->status == THREADSTATUS_WAIT_SYNCH_ANY);
+        ASSERT(reason == ThreadWakeupReason::Signal);
+
+        thread->SetWaitSynchronizationResult(RESULT_SUCCESS);
+        thread->SetWaitSynchronizationOutput(thread->GetWaitObjectIndex(object.get()));
+
+        // TODO(Subv): Perform IPC translation upon wakeup.
+    };
+
     Core::System::GetInstance().PrepareReschedule();
 
     // Note: The output of this SVC will be set to RESULT_SUCCESS if the thread resumes due to a
     // signal in one of its wait objects, or to 0xC8A01836 if there was a translation error.
     // By default the index is set to -1.
-    thread->wait_set_output = true;
     *index = -1;
     return RESULT_SUCCESS;
 }