1 //===- LazyAtomicPointer.----------------------------------------*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #ifndef LLVM_ADT_LAZYATOMICPOINTER_H
10 #define LLVM_ADT_LAZYATOMICPOINTER_H
11 
12 #include "llvm/ADT/STLFunctionalExtras.h"
13 #include "llvm/Support/Compiler.h"
14 #include <assert.h>
15 #include <atomic>
16 
17 namespace llvm {
18 
19 /// Atomic pointer that's lock-free, but that can coordinate concurrent writes
20 /// from a lazy generator. Should be reserved for cases where concurrent uses of
21 /// a generator for the same storage is unlikely.
22 ///
23 /// The laziness comes in with \a loadOrGenerate(), which lazily calls the
24 /// provided generator ONLY when the value is currently \c nullptr. With
25 /// concurrent calls, only one generator is called and the rest see that value.
26 ///
27 /// Most other APIs treat an in-flight \a loadOrGenerate() as if \c nullptr
28 /// were stored. APIs that are required to write a value will spin.
29 ///
30 /// The underlying storage is \a std::atomic<uintptr_t>.
31 ///
32 /// TODO: In C++20, use std::atomic<T>::wait() instead of spinning and call
33 /// std::atomic<T>::notify_all() in \a loadOrGenerate().
34 template <class T> class LazyAtomicPointer {
getNull()35   static constexpr uintptr_t getNull() { return 0; }
getBusy()36   static constexpr uintptr_t getBusy() { return -1ULL; }
37 
makePointer(uintptr_t Value)38   static T *makePointer(uintptr_t Value) {
39     assert(Value != getBusy());
40     return Value ? reinterpret_cast<T *>(Value) : nullptr;
41   }
makeRaw(T * Value)42   static uintptr_t makeRaw(T *Value) {
43     uintptr_t Raw = Value ? reinterpret_cast<uintptr_t>(Value) : getNull();
44     assert(Raw != getBusy());
45     return Raw;
46   }
47 
48 public:
49   /// Store a value. Waits for concurrent \a loadOrGenerate() calls.
store(T * Value)50   void store(T *Value) { return (void)exchange(Value); }
51 
52   /// Set a value. Return the old value. Waits for concurrent \a
53   /// loadOrGenerate() calls.
exchange(T * Value)54   T *exchange(T *Value) {
55     // Note: the call to compare_exchange_weak() fails "spuriously" if the
56     // current value is \a getBusy(), causing the loop to spin.
57     T *Old = nullptr;
58     while (!compare_exchange_weak(Old, Value)) {
59     }
60     return Old;
61   }
62 
63   /// Compare-exchange. Returns \c false if there is a concurrent \a
64   /// loadOrGenerate() call, setting \p ExistingValue to \c nullptr.
compare_exchange_weak(T * & ExistingValue,T * NewValue)65   bool compare_exchange_weak(T *&ExistingValue, T *NewValue) {
66     uintptr_t RawExistingValue = makeRaw(ExistingValue);
67     if (Storage.compare_exchange_weak(RawExistingValue, makeRaw(NewValue)))
68       return true;
69 
70     /// Report the existing value as "None" if busy.
71     if (RawExistingValue == getBusy())
72       ExistingValue = nullptr;
73     else
74       ExistingValue = makePointer(RawExistingValue);
75     return false;
76   }
77 
78   /// Compare-exchange. Keeps trying if there is a concurrent
79   /// \a loadOrGenerate() call.
compare_exchange_strong(T * & ExistingValue,T * NewValue)80   bool compare_exchange_strong(T *&ExistingValue, T *NewValue) {
81     uintptr_t RawExistingValue = makeRaw(ExistingValue);
82     const uintptr_t OriginalRawExistingValue = RawExistingValue;
83     if (Storage.compare_exchange_strong(RawExistingValue, makeRaw(NewValue)))
84       return true;
85 
86     /// Keep trying as long as it's busy.
87     if (LLVM_UNLIKELY(RawExistingValue == getBusy())) {
88       while (RawExistingValue == getBusy()) {
89         RawExistingValue = OriginalRawExistingValue;
90         if (Storage.compare_exchange_weak(RawExistingValue, makeRaw(NewValue)))
91           return true;
92       }
93     }
94     ExistingValue = makePointer(RawExistingValue);
95     return false;
96   }
97 
98   /// Return the current stored value. Returns \a None if there is a concurrent
99   /// \a loadOrGenerate() in flight.
load()100   T *load() const {
101     uintptr_t RawValue = Storage.load();
102     return RawValue == getBusy() ? nullptr : makePointer(RawValue);
103   }
104 
105   /// Get the current value, or call \p Generator to generate a value.
106   /// Guarantees that only one thread's \p Generator will run.
107   ///
108   /// \pre \p Generator doesn't return \c nullptr.
loadOrGenerate(function_ref<T * ()> Generator)109   T &loadOrGenerate(function_ref<T *()> Generator) {
110     // Return existing value, if already set.
111     uintptr_t Raw = Storage.load();
112     if (Raw != getNull() && Raw != getBusy())
113       return *makePointer(Raw);
114 
115     // Try to mark as busy, then generate and store a new value.
116     if (LLVM_LIKELY(Raw == getNull() &&
117                     Storage.compare_exchange_strong(Raw, getBusy()))) {
118       Raw = makeRaw(Generator());
119       assert(Raw != getNull() && "Expected non-null from generator");
120       Storage.store(Raw);
121       return *makePointer(Raw);
122     }
123 
124     // Contended with another generator. Wait for it to complete.
125     while (Raw == getBusy())
126       Raw = Storage.load();
127     assert(Raw != getNull() && "Expected non-null from competing generator");
128     return *makePointer(Raw);
129   }
130 
131   explicit operator bool() const { return load(); }
132   operator T *() const { return load(); }
133 
134   T &operator*() const {
135     T *P = load();
136     assert(P && "Unexpected null dereference");
137     return *P;
138   }
139   T *operator->() const { return &operator*(); }
140 
LazyAtomicPointer()141   LazyAtomicPointer() : Storage(0) {}
LazyAtomicPointer(std::nullptr_t)142   LazyAtomicPointer(std::nullptr_t) : Storage(0) {}
LazyAtomicPointer(T * Value)143   LazyAtomicPointer(T *Value) : Storage(makeRaw(Value)) {}
LazyAtomicPointer(const LazyAtomicPointer & RHS)144   LazyAtomicPointer(const LazyAtomicPointer &RHS)
145       : Storage(makeRaw(RHS.load())) {}
146 
147   LazyAtomicPointer &operator=(std::nullptr_t) {
148     store(nullptr);
149     return *this;
150   }
151   LazyAtomicPointer &operator=(T *RHS) {
152     store(RHS);
153     return *this;
154   }
155   LazyAtomicPointer &operator=(const LazyAtomicPointer &RHS) {
156     store(RHS.load());
157     return *this;
158   }
159 
160 private:
161   std::atomic<uintptr_t> Storage;
162 };
163 
164 } // end namespace llvm
165 
166 #endif // LLVM_ADT_LAZYATOMICPOINTER_H
167