4 files changed, 67 insertions, 6 deletions

diff --git a/src/core/hle/kernel/address_arbiter.cpp b/src/core/hle/kernel/address_arbiter.cpp
index cfd2c1590..367c4520d 100644
--- a/src/core/hle/kernel/address_arbiter.cpp
+++ b/src/core/hle/kernel/address_arbiter.cpp
@@ -5,14 +5,31 @@
 #include "common/assert.h"
 #include "common/common_funcs.h"
 #include "common/common_types.h"
+#include "core/core.h"
+#include "core/hle/kernel/errors.h"
 #include "core/hle/kernel/kernel.h"
 #include "core/hle/kernel/process.h"
 #include "core/hle/kernel/thread.h"
+#include "core/hle/lock.h"
 #include "core/memory.h"
 
 namespace Kernel {
     namespace AddressArbiter {
 
+        // Performs actual address waiting logic.
+        ResultCode WaitForAddress(VAddr address, s64 timeout) {
+            SharedPtr<Thread> current_thread = GetCurrentThread();
+            current_thread->arb_wait_address = address;
+            current_thread->arb_wait_result = RESULT_TIMEOUT;
+            current_thread->status = THREADSTATUS_WAIT_ARB;
+            current_thread->wakeup_callback = nullptr;
+
+            current_thread->WakeAfterDelay(timeout);
+
+            Core::System::GetInstance().CpuCore(current_thread->processor_id).PrepareReschedule();
+            return RESULT_SUCCESS;
+        }
+
         // Signals an address being waited on.
         ResultCode SignalToAddress(VAddr address, s32 value, s32 num_to_wake) {
             // TODO
@@ -33,14 +50,48 @@ namespace Kernel {
 
         // Waits on an address if the value passed is less than the argument value, optionally decrementing.
         ResultCode WaitForAddressIfLessThan(VAddr address, s32 value, s64 timeout, bool should_decrement) {
-            // TODO
-            return RESULT_SUCCESS;
+            // Ensure that we can read the address.
+            if (!Memory::IsValidVirtualAddress(address)) {
+                return ERR_INVALID_ADDRESS_STATE;
+            }
+
+            s32 cur_value;
+            // Get value, decrementing if less than
+            {
+                // Decrement if less than must be an atomic operation.
+                std::lock_guard<std::recursive_mutex> lock(HLE::g_hle_lock);
+                cur_value = (s32)Memory::Read32(address);
+                if (cur_value < value) {
+                    Memory::Write32(address, (u32)(cur_value - 1));
+                }
+            }
+            if (cur_value >= value) {
+                return ERR_INVALID_STATE;
+            }
+            // Short-circuit without rescheduling, if timeout is zero.
+            if (timeout == 0) {
+                return RESULT_TIMEOUT;
+            }
+
+            return WaitForAddress(address, timeout);
         }
 
         // Waits on an address if the value passed is equal to the argument value.
         ResultCode WaitForAddressIfEqual(VAddr address, s32 value, s64 timeout) {
-            // TODO
-            return RESULT_SUCCESS;
+            // Ensure that we can read the address.
+            if (!Memory::IsValidVirtualAddress(address)) {
+                return ERR_INVALID_ADDRESS_STATE;
+            }
+            // Only wait for the address if equal.
+            if ((s32)Memory::Read32(address) != value) {
+                return ERR_INVALID_STATE;
+            }
+            // Short-circuit without rescheduling, if timeout is zero.
+            if (timeout == 0) {
+                return RESULT_TIMEOUT;
+            }
+
+            return WaitForAddress(address, timeout);
         }
     } // namespace AddressArbiter
 } // namespace Kernel
\ No newline at end of file
diff --git a/src/core/hle/kernel/errors.h b/src/core/hle/kernel/errors.h
index 7ac960042..221cb1bb5 100644
--- a/src/core/hle/kernel/errors.h
+++ b/src/core/hle/kernel/errors.h
@@ -29,6 +29,7 @@ enum {
     SynchronizationCanceled = 118,
     TooLarge = 119,
     InvalidEnumValue = 120,
+    InvalidState = 125,
 };
 }
 
@@ -49,6 +50,7 @@ constexpr ResultCode ERR_OUT_OF_MEMORY(-1);
 constexpr ResultCode ERR_INVALID_ADDRESS(ErrorModule::Kernel, ErrCodes::InvalidAddress);
 constexpr ResultCode ERR_INVALID_ADDRESS_STATE(ErrorModule::Kernel, ErrCodes::InvalidMemoryState);
 constexpr ResultCode ERR_INVALID_HANDLE(ErrorModule::Kernel, ErrCodes::InvalidHandle);
+constexpr ResultCode ERR_INVALID_STATE(ErrorModule::Kernel, ErrCodes::InvalidState);
 constexpr ResultCode ERR_INVALID_POINTER(-1);
 constexpr ResultCode ERR_INVALID_OBJECT_ADDR(-1);
 constexpr ResultCode ERR_NOT_AUTHORIZED(-1);
diff --git a/src/core/hle/kernel/thread.cpp b/src/core/hle/kernel/thread.cpp
index cffa7ca83..2f333ec34 100644
--- a/src/core/hle/kernel/thread.cpp
+++ b/src/core/hle/kernel/thread.cpp
@@ -140,6 +140,11 @@ static void ThreadWakeupCallback(u64 thread_handle, int cycles_late) {
         }
     }
 
+    if (thread->arb_wait_address != 0) {
+        ASSERT(thread->status == THREADSTATUS_WAIT_ARB);
+        thread->arb_wait_address = 0;
+    }
+
     if (resume)
         thread->ResumeFromWait();
 }
@@ -179,6 +184,7 @@ void Thread::ResumeFromWait() {
     case THREADSTATUS_WAIT_SLEEP:
     case THREADSTATUS_WAIT_IPC:
     case THREADSTATUS_WAIT_MUTEX:
+    case THREADSTATUS_WAIT_ARB:
         break;
 
     case THREADSTATUS_READY:
diff --git a/src/core/hle/kernel/thread.h b/src/core/hle/kernel/thread.h
index 023c9dbe9..79e5d6e5c 100644
--- a/src/core/hle/kernel/thread.h
+++ b/src/core/hle/kernel/thread.h
@@ -228,8 +228,10 @@ public:
 
     // If waiting on a ConditionVariable, this is the ConditionVariable  address
     VAddr condvar_wait_address;
-    VAddr mutex_wait_address; ///< If waiting on a Mutex, this is the mutex address
-    Handle wait_handle;       ///< The handle used to wait for the mutex.
+    VAddr mutex_wait_address;   ///< If waiting on a Mutex, this is the mutex address
+    Handle wait_handle;         ///< The handle used to wait for the mutex.
+    VAddr arb_wait_address;     ///< If waiting for an AddressArbiter, this is the address
+    ResultCode arb_wait_result; ///< If waiting for an AddressArbiter, this is the result that will be returned.
 
     std::string name;