From 7929779df6a449c3f21b635ddc78c5db49abdfec Mon Sep 17 00:00:00 2001
From: flysand7 <thebumboni@gmail.com>
Date: Mon, 28 Aug 2023 01:36:46 +1100
Subject: [PATCH] Fix mutex

---
 os/linux/tinyrt.c    |  8 ++++++--
 src/cia-sync/mutex.c | 37 ++++++++++++++++++++++++++-----------
 2 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/os/linux/tinyrt.c b/os/linux/tinyrt.c
index 630bd0a..b79ea88 100644
--- a/os/linux/tinyrt.c
+++ b/os/linux/tinyrt.c
@@ -168,7 +168,9 @@ static _RT_Status _rt_sync_wake_one(u32 *addr, u32 *n_woken) {
     if(result < 0) {
         return _RT_ERROR_GENERIC;
     }
-    *n_woken = (u32)result;
+    if(n_woken != NULL) {
+        *n_woken = (u32)result;
+    }
     return _RT_STATUS_OK;
 }
 
@@ -177,6 +179,8 @@ static _RT_Status _rt_sync_wake_all(u32 *addr, u32 *n_woken) {
     if(result < 0) {
         return _RT_ERROR_GENERIC;
     }
-    *n_woken = (u32)result;
+    if(n_woken != NULL) {
+        *n_woken = (u32)result;
+    }
     return _RT_STATUS_OK;
 }
diff --git a/src/cia-sync/mutex.c b/src/cia-sync/mutex.c
index d4254a8..3dd9b48 100644
--- a/src/cia-sync/mutex.c
+++ b/src/cia-sync/mutex.c
@@ -1,21 +1,36 @@
 
-#define _CIA_MUTEX_TAG_UNLOCKED 0x00000000
-#define _CIA_MUTEX_TAG_LOCKED   0x00000001
+#define _CIA_MUTEX_FREE 0
+#define _CIA_MUTEX_LOCK 1
+#define _CIA_MUTEX_CONT 2
 
 void cia_mutex_init(Cia_Mutex *mutex) {
-    mutex->tag = _CIA_MUTEX_TAG_UNLOCKED;
+    mutex->tag = _CIA_MUTEX_FREE;
 }
 
 void cia_mutex_lock(Cia_Mutex *mutex) {
-    do {
-        _rt_sync_wait(&mutex->tag, _CIA_MUTEX_TAG_LOCKED, _RT_SYNC_WAIT_INFINITE);
-    } while(mutex->tag == _CIA_MUTEX_TAG_LOCKED);
-    mutex->tag = _CIA_MUTEX_TAG_LOCKED;
+    u64 prev_tag;
+    for(;;) {
+        prev_tag = __sync_val_compare_and_swap(&mutex->tag, _CIA_MUTEX_FREE, _CIA_MUTEX_LOCK);
+        // We got the mutex, lets bail
+        if(prev_tag == _CIA_MUTEX_FREE) {
+            break;
+        }
+        // We should wait if:
+        //   (1) the mutex is contested
+        //   (2) this thread locking the mutex makes it contested
+        bool should_wait = 0;
+        should_wait |= (prev_tag == _CIA_MUTEX_CONT);
+        should_wait |= (__sync_val_compare_and_swap(&mutex->tag, _CIA_MUTEX_LOCK, _CIA_MUTEX_CONT) != _CIA_MUTEX_FREE);
+        // We wait while its contested
+        if(should_wait) {
+            _rt_sync_wait(&mutex->tag, _CIA_MUTEX_CONT, _RT_SYNC_WAIT_INFINITE);
+        }
+    }
 }
 
 void cia_mutex_unlock(Cia_Mutex *mutex) {
-    mutex->tag = _CIA_MUTEX_TAG_UNLOCKED;
-    u32 woken = 0;
-    _rt_sync_wake_one(&mutex->tag, &woken);
+    // TODO: add error when we unlock a free mutex
+    // TODO: support recursive muteces
+    mutex->tag = _CIA_MUTEX_FREE;
+    _rt_sync_wake_one(&mutex->tag, NULL);
 }
-