#include <cppcoro/async_mutex.hpp>

#include <cassert>

cppcoro::async_mutex::async_mutex() noexcept

: m_state(not_locked)

, m_waiters(nullptr)

{}

cppcoro::async_mutex::~async_mutex()

{

[[maybe_unused]] auto state = m_state.load(std::memory_order_relaxed);

assert(state == not_locked || state == locked_no_waiters);

assert(m_waiters == nullptr);

}

bool cppcoro::async_mutex::try_lock() noexcept

{

// Try to atomically transition from nullptr (not-locked) -> this (locked-no-waiters).

auto oldState = not_locked;

return m_state.compare_exchange_strong(

oldState,

locked_no_waiters,

std::memory_order_acquire,

std::memory_order_relaxed);

}

cppcoro::async_mutex_lock_operation cppcoro::async_mutex::lock_async() noexcept

{

return async_mutex_lock_operation{ *this };

}

cppcoro::async_mutex_scoped_lock_operation cppcoro::async_mutex::scoped_lock_async() noexcept

{

return async_mutex_scoped_lock_operation{ *this };

}

void cppcoro::async_mutex::unlock()

{

assert(m_state.load(std::memory_order_relaxed) != not_locked);

async_mutex_lock_operation* waitersHead = m_waiters;

if (waitersHead == nullptr)

{

auto oldState = locked_no_waiters;

const bool releasedLock = m_state.compare_exchange_strong(

oldState,

not_locked,

std::memory_order_release,

std::memory_order_relaxed);

if (releasedLock)

{

return;

}

// At least one new waiter.

// Acquire the list of new waiter operations atomically.

oldState = m_state.exchange(locked_no_waiters, std::memory_order_acquire);

assert(oldState != locked_no_waiters && oldState != not_locked);

// Transfer the list to m_waiters, reversing the list in the process so

// that the head of the list is the first to be resumed.

auto* next = reinterpret_cast<async_mutex_lock_operation*>(oldState);

do

{

auto* temp = next->m_next;

next->m_next = waitersHead;

waitersHead = next;

next = temp;

} while (next != nullptr);

}

assert(waitersHead != nullptr);

m_waiters = waitersHead->m_next;

// Resume the waiter.

// This will pass the ownership of the lock on to that operation/coroutine.

waitersHead->m_awaiter.resume();

}

bool cppcoro::async_mutex_lock_operation::await_suspend(cppcoro::coroutine_handle<> awaiter) noexcept

{

m_awaiter = awaiter;

std::uintptr_t oldState = m_mutex.m_state.load(std::memory_order_acquire);

while (true)

{

if (oldState == async_mutex::not_locked)

{

if (m_mutex.m_state.compare_exchange_weak(

oldState,

async_mutex::locked_no_waiters,

std::memory_order_acquire,

std::memory_order_relaxed))

{

// Acquired lock, don't suspend.

return false;

}

}

else

{

// Try to push this operation onto the head of the waiter stack.

m_next = reinterpret_cast<async_mutex_lock_operation*>(oldState);

if (m_mutex.m_state.compare_exchange_weak(

oldState,

reinterpret_cast<std::uintptr_t>(this),

std::memory_order_release,

std::memory_order_relaxed))

{

// Queued operation to waiters list, suspend now.

return true;

}

}

}

}