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
17namespace 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().
34template <class T> class LazyAtomicPointer {
35  static constexpr uintptr_t getNull() { return 0; }
36  static constexpr uintptr_t getBusy() { return -1ULL; }
37
38  static T *makePointer(uintptr_t Value) {
39    assert(Value != getBusy());
40    return Value ? reinterpret_cast<T *>(Value) : nullptr;
41  }
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
48public:
49  /// Store a value. Waits for concurrent \a loadOrGenerate() calls.
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.
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.
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.
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.
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.
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
141  LazyAtomicPointer() : Storage(0) {}
142  LazyAtomicPointer(std::nullptr_t) : Storage(0) {}
143  LazyAtomicPointer(T *Value) : Storage(makeRaw(Value)) {}
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
160private:
161  std::atomic<uintptr_t> Storage;
162};
163
164} // end namespace llvm
165
166#endif // LLVM_ADT_LAZYATOMICPOINTER_H
167