7 files changed, 250 insertions, 75 deletions

diff --git a/src/core/hle/kernel/k_scoped_scheduler_lock_and_sleep.h b/src/core/hle/kernel/k_scoped_scheduler_lock_and_sleep.h
index 61dc2858f..2995c492d 100644
--- a/src/core/hle/kernel/k_scoped_scheduler_lock_and_sleep.h
+++ b/src/core/hle/kernel/k_scoped_scheduler_lock_and_sleep.h
@@ -8,6 +8,7 @@
 #pragma once
 
 #include "common/common_types.h"
+#include "core/hle/kernel/global_scheduler_context.h"
 #include "core/hle/kernel/k_thread.h"
 #include "core/hle/kernel/kernel.h"
 #include "core/hle/kernel/time_manager.h"
diff --git a/src/core/hle/kernel/k_synchronization_object.cpp b/src/core/hle/kernel/k_synchronization_object.cpp
index f168b4f21..ba8fc4010 100644
--- a/src/core/hle/kernel/k_synchronization_object.cpp
+++ b/src/core/hle/kernel/k_synchronization_object.cpp
@@ -8,11 +8,70 @@
 #include "core/hle/kernel/k_scoped_scheduler_lock_and_sleep.h"
 #include "core/hle/kernel/k_synchronization_object.h"
 #include "core/hle/kernel/k_thread.h"
+#include "core/hle/kernel/k_thread_queue.h"
 #include "core/hle/kernel/kernel.h"
 #include "core/hle/kernel/svc_results.h"
 
 namespace Kernel {
 
+namespace {
+
+class ThreadQueueImplForKSynchronizationObjectWait final : public KThreadQueueWithoutEndWait {
+private:
+    using ThreadListNode = KSynchronizationObject::ThreadListNode;
+
+private:
+    KSynchronizationObject** m_objects;
+    ThreadListNode* m_nodes;
+    s32 m_count;
+
+public:
+    ThreadQueueImplForKSynchronizationObjectWait(KernelCore& kernel_, KSynchronizationObject** o,
+                                                 ThreadListNode* n, s32 c)
+        : KThreadQueueWithoutEndWait(kernel_), m_objects(o), m_nodes(n), m_count(c) { // ...
+    }
+
+    virtual void NotifyAvailable(KThread* waiting_thread, KSynchronizationObject* signaled_object,
+                                 ResultCode wait_result) override {
+        // Determine the sync index, and unlink all nodes.
+        s32 sync_index = -1;
+        for (auto i = 0; i < m_count; ++i) {
+            // Check if this is the signaled object.
+            if (m_objects[i] == signaled_object && sync_index == -1) {
+                sync_index = i;
+            }
+
+            // Unlink the current node from the current object.
+            m_objects[i]->UnlinkNode(std::addressof(m_nodes[i]));
+        }
+
+        // Set the waiting thread's sync index.
+        waiting_thread->SetSyncedIndex(sync_index);
+
+        // Set the waiting thread as not cancellable.
+        waiting_thread->ClearCancellable();
+
+        // Invoke the base end wait handler.
+        KThreadQueue::EndWait(waiting_thread, wait_result);
+    }
+
+    virtual void CancelWait(KThread* waiting_thread, ResultCode wait_result,
+                            bool cancel_timer_task) override {
+        // Remove all nodes from our list.
+        for (auto i = 0; i < m_count; ++i) {
+            m_objects[i]->UnlinkNode(std::addressof(m_nodes[i]));
+        }
+
+        // Set the waiting thread as not cancellable.
+        waiting_thread->ClearCancellable();
+
+        // Invoke the base cancel wait handler.
+        KThreadQueue::CancelWait(waiting_thread, wait_result, cancel_timer_task);
+    }
+};
+
+} // namespace
+
 void KSynchronizationObject::Finalize() {
     this->OnFinalizeSynchronizationObject();
     KAutoObject::Finalize();
@@ -25,11 +84,19 @@ ResultCode KSynchronizationObject::Wait(KernelCore& kernel_ctx, s32* out_index,
     std::vector<ThreadListNode> thread_nodes(num_objects);
 
     // Prepare for wait.
-    KThread* thread = kernel_ctx.CurrentScheduler()->GetCurrentThread();
+    KThread* thread = GetCurrentThreadPointer(kernel_ctx);
+    ThreadQueueImplForKSynchronizationObjectWait wait_queue(kernel_ctx, objects,
+                                                            thread_nodes.data(), num_objects);
 
     {
         // Setup the scheduling lock and sleep.
-        KScopedSchedulerLockAndSleep slp{kernel_ctx, thread, timeout};
+        KScopedSchedulerLockAndSleep slp(kernel_ctx, thread, timeout);
+
+        // Check if the thread should terminate.
+        if (thread->IsTerminationRequested()) {
+            slp.CancelSleep();
+            return ResultTerminationRequested;
+        }
 
         // Check if any of the objects are already signaled.
         for (auto i = 0; i < num_objects; ++i) {
@@ -48,12 +115,6 @@ ResultCode KSynchronizationObject::Wait(KernelCore& kernel_ctx, s32* out_index,
             return ResultTimedOut;
         }
 
-        // Check if the thread should terminate.
-        if (thread->IsTerminationRequested()) {
-            slp.CancelSleep();
-            return ResultTerminationRequested;
-        }
-
         // Check if waiting was canceled.
         if (thread->IsWaitCancelled()) {
             slp.CancelSleep();
@@ -66,73 +127,25 @@ ResultCode KSynchronizationObject::Wait(KernelCore& kernel_ctx, s32* out_index,
             thread_nodes[i].thread = thread;
             thread_nodes[i].next = nullptr;
 
-            if (objects[i]->thread_list_tail == nullptr) {
-                objects[i]->thread_list_head = std::addressof(thread_nodes[i]);
-            } else {
-                objects[i]->thread_list_tail->next = std::addressof(thread_nodes[i]);
-            }
-
-            objects[i]->thread_list_tail = std::addressof(thread_nodes[i]);
+            objects[i]->LinkNode(std::addressof(thread_nodes[i]));
         }
 
-        // For debugging only
-        thread->SetWaitObjectsForDebugging({objects, static_cast<std::size_t>(num_objects)});
-
-        // Mark the thread as waiting.
+        // Mark the thread as cancellable.
         thread->SetCancellable();
-        thread->SetSyncedObject(nullptr, ResultTimedOut);
-        thread->SetState(ThreadState::Waiting);
-        thread->SetWaitReasonForDebugging(ThreadWaitReasonForDebugging::Synchronization);
-    }
-
-    // The lock/sleep is done, so we should be able to get our result.
 
-    // Thread is no longer cancellable.
-    thread->ClearCancellable();
+        // Clear the thread's synced index.
+        thread->SetSyncedIndex(-1);
 
-    // For debugging only
-    thread->SetWaitObjectsForDebugging({});
+        // Wait for an object to be signaled.
+        thread->BeginWait(std::addressof(wait_queue));
+        thread->SetWaitReasonForDebugging(ThreadWaitReasonForDebugging::Synchronization);
+    }
 
-    // Cancel the timer as needed.
-    kernel_ctx.TimeManager().UnscheduleTimeEvent(thread);
+    // Set the output index.
+    *out_index = thread->GetSyncedIndex();
 
     // Get the wait result.
-    ResultCode wait_result{ResultSuccess};
-    s32 sync_index = -1;
-    {
-        KScopedSchedulerLock lock(kernel_ctx);
-        KSynchronizationObject* synced_obj;
-        wait_result = thread->GetWaitResult(std::addressof(synced_obj));
-
-        for (auto i = 0; i < num_objects; ++i) {
-            // Unlink the object from the list.
-            ThreadListNode* prev_ptr =
-                reinterpret_cast<ThreadListNode*>(std::addressof(objects[i]->thread_list_head));
-            ThreadListNode* prev_val = nullptr;
-            ThreadListNode *prev, *tail_prev;
-
-            do {
-                prev = prev_ptr;
-                prev_ptr = prev_ptr->next;
-                tail_prev = prev_val;
-                prev_val = prev_ptr;
-            } while (prev_ptr != std::addressof(thread_nodes[i]));
-
-            if (objects[i]->thread_list_tail == std::addressof(thread_nodes[i])) {
-                objects[i]->thread_list_tail = tail_prev;
-            }
-
-            prev->next = thread_nodes[i].next;
-
-            if (objects[i] == synced_obj) {
-                sync_index = i;
-            }
-        }
-    }
-
-    // Set output.
-    *out_index = sync_index;
-    return wait_result;
+    return thread->GetWaitResult();
 }
 
 KSynchronizationObject::KSynchronizationObject(KernelCore& kernel_)
@@ -141,7 +154,7 @@ KSynchronizationObject::KSynchronizationObject(KernelCore& kernel_)
 KSynchronizationObject::~KSynchronizationObject() = default;
 
 void KSynchronizationObject::NotifyAvailable(ResultCode result) {
-    KScopedSchedulerLock lock(kernel);
+    KScopedSchedulerLock sl(kernel);
 
     // If we're not signaled, we've nothing to notify.
     if (!this->IsSignaled()) {
@@ -150,11 +163,7 @@ void KSynchronizationObject::NotifyAvailable(ResultCode result) {
 
     // Iterate over each thread.
     for (auto* cur_node = thread_list_head; cur_node != nullptr; cur_node = cur_node->next) {
-        KThread* thread = cur_node->thread;
-        if (thread->GetState() == ThreadState::Waiting) {
-            thread->SetSyncedObject(this, result);
-            thread->SetState(ThreadState::Runnable);
-        }
+        cur_node->thread->NotifyAvailable(this, result);
     }
 }
 
diff --git a/src/core/hle/kernel/k_synchronization_object.h b/src/core/hle/kernel/k_synchronization_object.h
index 898e58e16..789ff4780 100644
--- a/src/core/hle/kernel/k_synchronization_object.h
+++ b/src/core/hle/kernel/k_synchronization_object.h
@@ -35,6 +35,38 @@ public:
 
     [[nodiscard]] std::vector<KThread*> GetWaitingThreadsForDebugging() const;
 
+    void LinkNode(ThreadListNode* node) {
+        // Link the node to the list.
+        if (thread_list_tail == nullptr) {
+            thread_list_head = node;
+        } else {
+            thread_list_tail->next = node;
+        }
+
+        thread_list_tail = node;
+    }
+
+    void UnlinkNode(ThreadListNode* node) {
+        // Unlink the node from the list.
+        ThreadListNode* prev_ptr =
+            reinterpret_cast<ThreadListNode*>(std::addressof(thread_list_head));
+        ThreadListNode* prev_val = nullptr;
+        ThreadListNode *prev, *tail_prev;
+
+        do {
+            prev = prev_ptr;
+            prev_ptr = prev_ptr->next;
+            tail_prev = prev_val;
+            prev_val = prev_ptr;
+        } while (prev_ptr != node);
+
+        if (thread_list_tail == node) {
+            thread_list_tail = tail_prev;
+        }
+
+        prev->next = node->next;
+    }
+
 protected:
     explicit KSynchronizationObject(KernelCore& kernel);
     ~KSynchronizationObject() override;
diff --git a/src/core/hle/kernel/k_thread.cpp b/src/core/hle/kernel/k_thread.cpp
index 41bf9a6bb..3331b4845 100644
--- a/src/core/hle/kernel/k_thread.cpp
+++ b/src/core/hle/kernel/k_thread.cpp
@@ -303,7 +303,7 @@ void KThread::Wakeup() {
 
     if (GetState() == ThreadState::Waiting) {
         if (sleeping_queue != nullptr) {
-            sleeping_queue->WakeupThread(this);
+            sleeping_queue->EndWait(this, ResultSuccess);
         } else {
             SetState(ThreadState::Runnable);
         }
@@ -331,7 +331,7 @@ void KThread::StartTermination() {
 
     // Signal.
     signaled = true;
-    NotifyAvailable();
+    KSynchronizationObject::NotifyAvailable();
 
     // Clear previous thread in KScheduler.
     KScheduler::ClearPreviousThread(kernel, this);
@@ -1026,6 +1026,44 @@ ResultCode KThread::Sleep(s64 timeout) {
     return ResultSuccess;
 }
 
+void KThread::BeginWait(KThreadQueue* queue) {
+    // Set our state as waiting.
+    SetState(ThreadState::Waiting);
+
+    // Set our wait queue.
+    sleeping_queue = queue;
+}
+
+void KThread::NotifyAvailable(KSynchronizationObject* signaled_object, ResultCode wait_result_) {
+    // Lock the scheduler.
+    KScopedSchedulerLock sl(kernel);
+
+    // If we're waiting, notify our queue that we're available.
+    if (GetState() == ThreadState::Waiting) {
+        sleeping_queue->NotifyAvailable(this, signaled_object, wait_result_);
+    }
+}
+
+void KThread::EndWait(ResultCode wait_result_) {
+    // Lock the scheduler.
+    KScopedSchedulerLock sl(kernel);
+
+    // If we're waiting, notify our queue that we're available.
+    if (GetState() == ThreadState::Waiting) {
+        sleeping_queue->EndWait(this, wait_result_);
+    }
+}
+
+void KThread::CancelWait(ResultCode wait_result_, bool cancel_timer_task) {
+    // Lock the scheduler.
+    KScopedSchedulerLock sl(kernel);
+
+    // If we're waiting, notify our queue that we're available.
+    if (GetState() == ThreadState::Waiting) {
+        sleeping_queue->CancelWait(this, wait_result_, cancel_timer_task);
+    }
+}
+
 void KThread::SetState(ThreadState state) {
     KScopedSchedulerLock sl{kernel};
 
diff --git a/src/core/hle/kernel/k_thread.h b/src/core/hle/kernel/k_thread.h
index e4c4c877d..bea5fd84d 100644
--- a/src/core/hle/kernel/k_thread.h
+++ b/src/core/hle/kernel/k_thread.h
@@ -202,6 +202,23 @@ public:
         wait_result = wait_res;
     }
 
+    constexpr void SetSyncedIndex(s32 index) {
+        synced_index = index;
+    }
+
+    constexpr s32 GetSyncedIndex() const {
+        return synced_index;
+    }
+
+    constexpr void SetWaitResult(ResultCode wait_res) {
+        wait_result = wait_res;
+        synced_object = nullptr;
+    }
+
+    constexpr ResultCode GetWaitResult() const {
+        return wait_result;
+    }
+
     [[nodiscard]] ResultCode GetWaitResult(KSynchronizationObject** out) const {
         *out = synced_object;
         return wait_result;
@@ -596,6 +613,15 @@ public:
         address_key_value = val;
     }
 
+    void ClearWaitQueue() {
+        sleeping_queue = nullptr;
+    }
+
+    void BeginWait(KThreadQueue* queue);
+    void NotifyAvailable(KSynchronizationObject* signaled_object, ResultCode wait_result_);
+    void EndWait(ResultCode wait_result_);
+    void CancelWait(ResultCode wait_result_, bool cancel_timer_task);
+
     [[nodiscard]] bool HasWaiters() const {
         return !waiter_list.empty();
     }
@@ -707,6 +733,7 @@ private:
     u32 address_key_value{};
     u32 suspend_request_flags{};
     u32 suspend_allowed_flags{};
+    s32 synced_index{};
     ResultCode wait_result{ResultSuccess};
     s32 base_priority{};
     s32 physical_ideal_core_id{};
diff --git a/src/core/hle/kernel/k_thread_queue.cpp b/src/core/hle/kernel/k_thread_queue.cpp
new file mode 100644
index 000000000..46f27172b
--- /dev/null
+++ b/src/core/hle/kernel/k_thread_queue.cpp
@@ -0,0 +1,51 @@
+// Copyright 2021 yuzu Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#pragma once
+
+#include "core/hle/kernel/k_thread_queue.h"
+#include "core/hle/kernel/kernel.h"
+#include "core/hle/kernel/time_manager.h"
+
+namespace Kernel {
+
+void KThreadQueue::NotifyAvailable([[maybe_unused]] KThread* waiting_thread,
+                                   [[maybe_unused]] KSynchronizationObject* signaled_object,
+                                   [[maybe_unused]] ResultCode wait_result) {}
+
+void KThreadQueue::EndWait(KThread* waiting_thread, ResultCode wait_result) {
+    // Set the thread's wait result.
+    waiting_thread->SetWaitResult(wait_result);
+
+    // Set the thread as runnable.
+    waiting_thread->SetState(ThreadState::Runnable);
+
+    // Clear the thread's wait queue.
+    waiting_thread->ClearWaitQueue();
+
+    // Cancel the thread task.
+    kernel.TimeManager().UnscheduleTimeEvent(waiting_thread);
+}
+
+void KThreadQueue::CancelWait(KThread* waiting_thread, ResultCode wait_result,
+                              bool cancel_timer_task) {
+    // Set the thread's wait result.
+    waiting_thread->SetWaitResult(wait_result);
+
+    // Set the thread as runnable.
+    waiting_thread->SetState(ThreadState::Runnable);
+
+    // Clear the thread's wait queue.
+    waiting_thread->ClearWaitQueue();
+
+    // Cancel the thread task.
+    if (cancel_timer_task) {
+        kernel.TimeManager().UnscheduleTimeEvent(waiting_thread);
+    }
+}
+
+void KThreadQueueWithoutEndWait::EndWait([[maybe_unused]] KThread* waiting_thread,
+                                         [[maybe_unused]] ResultCode wait_result) {}
+
+} // namespace Kernel
diff --git a/src/core/hle/kernel/k_thread_queue.h b/src/core/hle/kernel/k_thread_queue.h
index 35d471dc5..74e76e7cb 100644
--- a/src/core/hle/kernel/k_thread_queue.h
+++ b/src/core/hle/kernel/k_thread_queue.h
@@ -4,6 +4,7 @@
 
 #pragma once
 
+#include "core/hle/kernel/k_scheduler.h"
 #include "core/hle/kernel/k_thread.h"
 
 namespace Kernel {
@@ -11,7 +12,16 @@ namespace Kernel {
 class KThreadQueue {
 public:
     explicit KThreadQueue(KernelCore& kernel_) : kernel{kernel_} {}
+    virtual ~KThreadQueue(){};
 
+    virtual void NotifyAvailable(KThread* waiting_thread, KSynchronizationObject* signaled_object,
+                                 ResultCode wait_result);
+    virtual void EndWait(KThread* waiting_thread, ResultCode wait_result);
+    virtual void CancelWait(KThread* waiting_thread, ResultCode wait_result,
+                            bool cancel_timer_task);
+
+    // Deprecated, will be removed in subsequent commits.
+public:
     bool IsEmpty() const {
         return wait_list.empty();
     }
@@ -78,4 +88,11 @@ private:
     KThread::WaiterList wait_list{};
 };
 
+class KThreadQueueWithoutEndWait : public KThreadQueue {
+public:
+    explicit KThreadQueueWithoutEndWait(KernelCore& kernel_) : KThreadQueue(kernel_) {}
+
+    virtual void EndWait(KThread* waiting_thread, ResultCode wait_result) override final;
+};
+
 } // namespace Kernel