waitN: Fix index bug when there are handles before UEvents in the array

nx/include/switch/kernel/wait.h

@@ -35,8 +35,8 @@ typedef struct {
         UEvent* parent_event;
         UTimer* parent_timer;
     };
-    size_t idx;
+    s32 idx;
-    size_t* idx_out;
+    s32* idx_out;
 } WaiterNode;
 
 struct Waitable {

nx/source/kernel/uevent.c

@@ -36,10 +36,8 @@ void _ueventTryAutoClear(UEvent* e)
     mutexUnlock(&e->waitable.mutex);
 }
 
-bool _ueventAddListener(UEvent* e, WaiterNode* w, size_t idx, size_t* idx_out, Handle thread)
+bool _ueventAddListener(UEvent* e, WaiterNode* w, s32 idx, s32* idx_out, Handle thread)
 {
-    _waiterNodeCreate(w, WaiterNodeType_Event, &e->waitable, thread, idx, idx_out);
-
     mutexLock(&e->waitable.mutex);
 
     bool signalled = e->signal;
@@ -47,8 +45,10 @@ bool _ueventAddListener(UEvent* e, WaiterNode* w, size_t idx, size_t* idx_out, H
     if (signalled) {
         if (e->auto_clear)
             e->signal = false;
-    } else
+    } else {
+        _waiterNodeCreate(w, WaiterNodeType_Event, &e->waitable, thread, idx, idx_out);
         _waiterNodeAddToWaitable(w, &e->waitable);
+    }
 
     mutexUnlock(&e->waitable.mutex);
     return !signalled;

nx/source/kernel/uevent.h

@@ -3,4 +3,4 @@
 #include "kernel/uevent.h"
 
 void _ueventTryAutoClear(UEvent* e);
-bool _ueventAddListener(UEvent* e, WaiterNode* w, size_t idx, size_t* idx_out, Handle thread);
+bool _ueventAddListener(UEvent* e, WaiterNode* w, s32 idx, s32* idx_out, Handle thread);

nx/source/kernel/utimer.c

@@ -75,7 +75,7 @@ u64 _utimerGetNextTick(UTimer* t)
     return ret;
 }
 
-void _utimerAddListener(UTimer* t, WaiterNode* w, size_t idx, size_t* idx_out, Handle thread)
+void _utimerAddListener(UTimer* t, WaiterNode* w, s32 idx, s32* idx_out, Handle thread)
 {
     _waiterNodeCreate(w, WaiterNodeType_Timer, &t->waitable, thread, idx, idx_out);
 

nx/source/kernel/utimer.h

@@ -4,4 +4,4 @@
 
 void _utimerRecalculate(UTimer* t, u64 old_tick);
 u64  _utimerGetNextTick(UTimer* t);
-void _utimerAddListener(UTimer* t, WaiterNode* w, size_t idx, size_t* idx_out, Handle thread);
+void _utimerAddListener(UTimer* t, WaiterNode* w, s32 idx, s32* idx_out, Handle thread);

nx/source/kernel/wait.c

@@ -26,8 +26,8 @@ static Result waitImpl(s32* idx_out, Waiter* objects, size_t num_objects, u64 ti
     Handle handles[num_objects];
     u64 cur_tick = armGetSystemTick();
 
-    size_t triggered_idx = -1;
+    s32 triggered_idx = -1;
-    size_t num_waiters = 0;
+    u64 waiters_added = 0;
     WaiterNode waiters[num_objects];
 
     u64 end_tick = UINT64_MAX;
@@ -67,16 +67,16 @@ static Result waitImpl(s32* idx_out, Waiter* objects, size_t num_objects, u64 ti
                 // Always add a listener on the timer,
                 // If the timer is started/stopped we want to detect that.
                 _utimerAddListener(
-                    obj->timer, &waiters[num_waiters], num_waiters, &triggered_idx,
+                    obj->timer, &waiters[i], i, &triggered_idx,
                     own_thread_handle);
 
-                num_waiters++;
+                waiters_added |= 1ULL << i;
                 break;
 
             case WaiterType_UEvent:
                 // Try to add a listener to the event, if it hasn't already signalled.
                 added = _ueventAddListener(
-                    obj->event, &waiters[num_waiters], num_waiters, &triggered_idx,
+                    obj->event, &waiters[i], i, &triggered_idx,
                     own_thread_handle);
 
                 // If the event already happened, we're done.
@@ -87,7 +87,7 @@ static Result waitImpl(s32* idx_out, Waiter* objects, size_t num_objects, u64 ti
                 }
 
                 // If the event hasn't signalled, we added a listener.
-                num_waiters++;
+                waiters_added |= 1ULL << i;
                 break;
 
             case WaiterType_Handle:
@@ -138,8 +138,9 @@ static Result waitImpl(s32* idx_out, Waiter* objects, size_t num_objects, u64 ti
 
 clean_up:
     // Remove listeners.
-    for (i = 0; i < num_waiters; i ++)
+    for (i = 0; i < num_objects; i ++)
-        _waiterNodeFree(&waiters[i]);
+        if (waiters_added & (1ULL << i))
+            _waiterNodeFree(&waiters[i]);
 
     return rc;
 }

nx/source/kernel/wait.h

@@ -21,7 +21,7 @@ static inline void _waitableSignalAllListeners(Waitable* ww)
 
         // Try to swap -1 => idx on the waiter thread.
         // If another waitable signals simultaneously only one will win the race and insert its own idx.
-        size_t minus_one = -1;
+        s32 minus_one = -1;
         bool sent_idx = __atomic_compare_exchange_n(
             w->idx_out, &minus_one, w->idx, true, __ATOMIC_SEQ_CST, __ATOMIC_SEQ_CST);
 
@@ -32,7 +32,7 @@ static inline void _waitableSignalAllListeners(Waitable* ww)
 
 static inline void _waiterNodeCreate(
     WaiterNode* w, WaiterNodeType type, Waitable* parent, Handle thread,
-    size_t idx, size_t* idx_out)
+    s32 idx, s32* idx_out)
 {
     w->type = type;
     w->parent = parent;