1// -*- C++ -*-
2//===----------------------------------------------------------------------===//
3//
4// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
5// See https://llvm.org/LICENSE.txt for license information.
6// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
7//
8//===----------------------------------------------------------------------===//
9
10#ifndef _LIBCPP___STOP_TOKEN_ATOMIC_UNIQUE_LOCK_H
11#define _LIBCPP___STOP_TOKEN_ATOMIC_UNIQUE_LOCK_H
12
13#include <__atomic/atomic.h>
14#include <__atomic/memory_order.h>
15#include <__bit/has_single_bit.h>
16#include <__config>
17
18#if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
19# pragma GCC system_header
20#endif
21
22#if _LIBCPP_STD_VER >= 20
23
24_LIBCPP_BEGIN_NAMESPACE_STD
25
26// This class implements an RAII unique_lock without a mutex.
27// It uses std::atomic<State>,
28// where State contains a lock bit and might contain other data,
29// and LockedBit is the value of State when the lock bit is set, e.g 1 << 2
30template <class _State, _State _LockedBit>
31class __atomic_unique_lock {
32 static_assert(std::has_single_bit(t: static_cast<unsigned long long>(_LockedBit)),
33 "LockedBit must be an integer where only one bit is set");
34
35 std::atomic<_State>& __state_;
36 bool __is_locked_;
37
38public:
39 _LIBCPP_HIDE_FROM_ABI explicit __atomic_unique_lock(std::atomic<_State>& __state) noexcept
40 : __state_(__state), __is_locked_(true) {
41 __lock();
42 }
43
44 template <class _Pred>
45 _LIBCPP_HIDE_FROM_ABI __atomic_unique_lock(std::atomic<_State>& __state, _Pred&& __give_up_locking) noexcept
46 : __state_(__state), __is_locked_(false) {
47 __is_locked_ = __lock_impl(__give_up_locking, __set_locked_bit, std::memory_order_acquire);
48 }
49
50 template <class _Pred, class _UnaryFunction>
51 _LIBCPP_HIDE_FROM_ABI __atomic_unique_lock(
52 std::atomic<_State>& __state,
53 _Pred&& __give_up_locking,
54 _UnaryFunction&& __state_after_lock,
55 std::memory_order __locked_ordering) noexcept
56 : __state_(__state), __is_locked_(false) {
57 __is_locked_ = __lock_impl(__give_up_locking, __state_after_lock, __locked_ordering);
58 }
59
60 __atomic_unique_lock(const __atomic_unique_lock&) = delete;
61 __atomic_unique_lock(__atomic_unique_lock&&) = delete;
62 __atomic_unique_lock& operator=(const __atomic_unique_lock&) = delete;
63 __atomic_unique_lock& operator=(__atomic_unique_lock&&) = delete;
64
65 _LIBCPP_HIDE_FROM_ABI ~__atomic_unique_lock() {
66 if (__is_locked_) {
67 __unlock();
68 }
69 }
70
71 _LIBCPP_HIDE_FROM_ABI bool __owns_lock() const noexcept { return __is_locked_; }
72
73 _LIBCPP_HIDE_FROM_ABI void __lock() noexcept {
74 const auto __never_give_up_locking = [](_State) { return false; };
75 // std::memory_order_acquire because we'd like to make sure that all the read operations after the lock can read the
76 // up-to-date values.
77 __lock_impl(__never_give_up_locking, __set_locked_bit, std::memory_order_acquire);
78 __is_locked_ = true;
79 }
80
81 _LIBCPP_HIDE_FROM_ABI void __unlock() noexcept {
82 // unset the _LockedBit. `memory_order_release` because we need to make sure all the write operations before calling
83 // `__unlock` will be made visible to other threads
84 __state_.fetch_and(static_cast<_State>(~_LockedBit), std::memory_order_release);
85 __state_.notify_all();
86 __is_locked_ = false;
87 }
88
89private:
90 template <class _Pred, class _UnaryFunction>
91 _LIBCPP_HIDE_FROM_ABI bool
92 __lock_impl(_Pred&& __give_up_locking, // while trying to lock the state, if the predicate returns true, give up
93 // locking and return
94 _UnaryFunction&& __state_after_lock,
95 std::memory_order __locked_ordering) noexcept {
96 // At this stage, until we exit the inner while loop, other than the atomic state, we are not reading any order
97 // dependent values that is written on other threads, or writing anything that needs to be seen on other threads.
98 // Therefore `memory_order_relaxed` is enough.
99 _State __current_state = __state_.load(std::memory_order_relaxed);
100 do {
101 while (true) {
102 if (__give_up_locking(__current_state)) {
103 // user provided early return condition. fail to lock
104 return false;
105 } else if ((__current_state & _LockedBit) != 0) {
106 // another thread has locked the state, we need to wait
107 __state_.wait(__current_state, std::memory_order_relaxed);
108 // when it is woken up by notifyAll or spuriously, the __state_
109 // might have changed. reload the state
110 // Note that the new state's _LockedBit may or may not equal to 0
111 __current_state = __state_.load(std::memory_order_relaxed);
112 } else {
113 // at least for now, it is not locked. we can try `compare_exchange_weak` to lock it.
114 // Note that the variable `__current_state`'s lock bit has to be 0 at this point.
115 break;
116 }
117 }
118 } while (!__state_.compare_exchange_weak(
119 __current_state, // if __state_ has the same value of __current_state, lock bit must be zero before exchange and
120 // we are good to lock/exchange and return. If _state has a different value, because other
121 // threads locked it between the `break` statement above and this statement, exchange will fail
122 // and go back to the inner while loop above.
123 __state_after_lock(__current_state), // state after lock. Usually it should be __current_state | _LockedBit.
124 // Some use cases need to set other bits at the same time as an atomic
125 // operation therefore we accept a function
126 __locked_ordering, // sucessful exchange order. Usually it should be std::memory_order_acquire.
127 // Some use cases need more strict ordering therefore we accept it as a parameter
128 std::memory_order_relaxed // fail to exchange order. We don't need any ordering as we are going back to the
129 // inner while loop
130 ));
131 return true;
132 }
133
134 _LIBCPP_HIDE_FROM_ABI static constexpr auto __set_locked_bit = [](_State __state) { return __state | _LockedBit; };
135};
136
137_LIBCPP_END_NAMESPACE_STD
138
139#endif // _LIBCPP_STD_VER >= 20
140
141#endif // _LIBCPP___STOP_TOKEN_ATOMIC_UNIQUE_LOCK_H
142