llvm/ADT/LazyAtomicPointer.h

0001 //===- LazyAtomicPointer.----------------------------------------*- C++ -*-===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008
0009 #ifndef LLVM_ADT_LAZYATOMICPOINTER_H
0010 #define LLVM_ADT_LAZYATOMICPOINTER_H
0011
0012 #include "llvm/ADT/STLFunctionalExtras.h"
0013 #include "llvm/Support/Compiler.h"
0014 #include <assert.h>
0015 #include <atomic>
0016
0017 namespace llvm {
0018
0019 /// Atomic pointer that's lock-free, but that can coordinate concurrent writes
0020 /// from a lazy generator. Should be reserved for cases where concurrent uses of
0021 /// a generator for the same storage is unlikely.
0022 ///
0023 /// The laziness comes in with \a loadOrGenerate(), which lazily calls the
0024 /// provided generator ONLY when the value is currently \c nullptr. With
0025 /// concurrent calls, only one generator is called and the rest see that value.
0026 ///
0027 /// Most other APIs treat an in-flight \a loadOrGenerate() as if \c nullptr
0028 /// were stored. APIs that are required to write a value will spin.
0029 ///
0030 /// The underlying storage is \a std::atomic<uintptr_t>.
0031 ///
0032 /// TODO: In C++20, use std::atomic<T>::wait() instead of spinning and call
0033 /// std::atomic<T>::notify_all() in \a loadOrGenerate().
0034 template <class T> class LazyAtomicPointer {
0035   static constexpr uintptr_t getNull() { return 0; }
0036   static constexpr uintptr_t getBusy() { return UINTPTR_MAX; }
0037
0038   static T *makePointer(uintptr_t Value) {
0039     assert(Value != getBusy());
0040     return Value ? reinterpret_cast<T *>(Value) : nullptr;
0041   }
0042   static uintptr_t makeRaw(T *Value) {
0043     uintptr_t Raw = Value ? reinterpret_cast<uintptr_t>(Value) : getNull();
0044     assert(Raw != getBusy());
0045     return Raw;
0046   }
0047
0048 public:
0049   /// Store a value. Waits for concurrent \a loadOrGenerate() calls.
0050   void store(T *Value) { return (void)exchange(Value); }
0051
0052   /// Set a value. Return the old value. Waits for concurrent \a
0053   /// loadOrGenerate() calls.
0054   T *exchange(T *Value) {
0055     // Note: the call to compare_exchange_weak() fails "spuriously" if the
0056     // current value is \a getBusy(), causing the loop to spin.
0057     T *Old = nullptr;
0058     while (!compare_exchange_weak(Old, Value)) {
0059     }
0060     return Old;
0061   }
0062
0063   /// Compare-exchange. Returns \c false if there is a concurrent \a
0064   /// loadOrGenerate() call, setting \p ExistingValue to \c nullptr.
0065   bool compare_exchange_weak(T *&ExistingValue, T *NewValue) {
0066     uintptr_t RawExistingValue = makeRaw(ExistingValue);
0067     if (Storage.compare_exchange_weak(RawExistingValue, makeRaw(NewValue)))
0068       return true;
0069
0070     /// Report the existing value as "None" if busy.
0071     if (RawExistingValue == getBusy())
0072       ExistingValue = nullptr;
0073     else
0074       ExistingValue = makePointer(RawExistingValue);
0075     return false;
0076   }
0077
0078   /// Compare-exchange. Keeps trying if there is a concurrent
0079   /// \a loadOrGenerate() call.
0080   bool compare_exchange_strong(T *&ExistingValue, T *NewValue) {
0081     uintptr_t RawExistingValue = makeRaw(ExistingValue);
0082     const uintptr_t OriginalRawExistingValue = RawExistingValue;
0083     if (Storage.compare_exchange_strong(RawExistingValue, makeRaw(NewValue)))
0084       return true;
0085
0086     /// Keep trying as long as it's busy.
0087     if (LLVM_UNLIKELY(RawExistingValue == getBusy())) {
0088       while (RawExistingValue == getBusy()) {
0089         RawExistingValue = OriginalRawExistingValue;
0090         if (Storage.compare_exchange_weak(RawExistingValue, makeRaw(NewValue)))
0091           return true;
0092       }
0093     }
0094     ExistingValue = makePointer(RawExistingValue);
0095     return false;
0096   }
0097
0098   /// Return the current stored value. Returns \a None if there is a concurrent
0099   /// \a loadOrGenerate() in flight.
0100   T *load() const {
0101     uintptr_t RawValue = Storage.load();
0102     return RawValue == getBusy() ? nullptr : makePointer(RawValue);
0103   }
0104
0105   /// Get the current value, or call \p Generator to generate a value.
0106   /// Guarantees that only one thread's \p Generator will run.
0107   ///
0108   /// \pre \p Generator doesn't return \c nullptr.
0109   T &loadOrGenerate(function_ref<T *()> Generator) {
0110     // Return existing value, if already set.
0111     uintptr_t Raw = Storage.load();
0112     if (Raw != getNull() && Raw != getBusy())
0113       return *makePointer(Raw);
0114
0115     // Try to mark as busy, then generate and store a new value.
0116     if (LLVM_LIKELY(Raw == getNull() &&
0117                     Storage.compare_exchange_strong(Raw, getBusy()))) {
0118       Raw = makeRaw(Generator());
0119       assert(Raw != getNull() && "Expected non-null from generator");
0120       Storage.store(Raw);
0121       return *makePointer(Raw);
0122     }
0123
0124     // Contended with another generator. Wait for it to complete.
0125     while (Raw == getBusy())
0126       Raw = Storage.load();
0127     assert(Raw != getNull() && "Expected non-null from competing generator");
0128     return *makePointer(Raw);
0129   }
0130
0131   explicit operator bool() const { return load(); }
0132   operator T *() const { return load(); }
0133
0134   T &operator*() const {
0135     T *P = load();
0136     assert(P && "Unexpected null dereference");
0137     return *P;
0138   }
0139   T *operator->() const { return &operator*(); }
0140
0141   LazyAtomicPointer() : Storage(0) {}
0142   LazyAtomicPointer(std::nullptr_t) : Storage(0) {}
0143   LazyAtomicPointer(T *Value) : Storage(makeRaw(Value)) {}
0144   LazyAtomicPointer(const LazyAtomicPointer &RHS)
0145       : Storage(makeRaw(RHS.load())) {}
0146
0147   LazyAtomicPointer &operator=(std::nullptr_t) {
0148     store(nullptr);
0149     return *this;
0150   }
0151   LazyAtomicPointer &operator=(T *RHS) {
0152     store(RHS);
0153     return *this;
0154   }
0155   LazyAtomicPointer &operator=(const LazyAtomicPointer &RHS) {
0156     store(RHS.load());
0157     return *this;
0158   }
0159
0160 private:
0161   std::atomic<uintptr_t> Storage;
0162 };
0163
0164 } // end namespace llvm
0165
0166 #endif // LLVM_ADT_LAZYATOMICPOINTER_H