1//===-- SymbolStringPool.h -- Thread-safe pool for JIT symbols --*- 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// Contains a thread-safe string pool suitable for use with ORC.
10//
11//===----------------------------------------------------------------------===//
12
13#ifndef LLVM_EXECUTIONENGINE_ORC_SYMBOLSTRINGPOOL_H
14#define LLVM_EXECUTIONENGINE_ORC_SYMBOLSTRINGPOOL_H
15
16#include "llvm/ADT/DenseMap.h"
17#include "llvm/ADT/StringMap.h"
18#include <atomic>
19#include <mutex>
20
21namespace llvm {
22
23class raw_ostream;
24
25namespace orc {
26
27class SymbolStringPtrBase;
28class SymbolStringPtr;
29class NonOwningSymbolStringPtr;
30
31/// String pool for symbol names used by the JIT.
32class SymbolStringPool {
33  friend class SymbolStringPoolTest;
34  friend class SymbolStringPtrBase;
35  friend class SymbolStringPoolEntryUnsafe;
36
37  // Implemented in DebugUtils.h.
38  friend raw_ostream &operator<<(raw_ostream &OS, const SymbolStringPool &SSP);
39
40public:
41  /// Destroy a SymbolStringPool.
42  ~SymbolStringPool();
43
44  /// Create a symbol string pointer from the given string.
45  SymbolStringPtr intern(StringRef S);
46
47  /// Remove from the pool any entries that are no longer referenced.
48  void clearDeadEntries();
49
50  /// Returns true if the pool is empty.
51  bool empty() const;
52
53private:
54  size_t getRefCount(const SymbolStringPtrBase &S) const;
55
56  using RefCountType = std::atomic<size_t>;
57  using PoolMap = StringMap<RefCountType>;
58  using PoolMapEntry = StringMapEntry<RefCountType>;
59  mutable std::mutex PoolMutex;
60  PoolMap Pool;
61};
62
63/// Base class for both owning and non-owning symbol-string ptrs.
64///
65/// All symbol-string ptrs are convertible to bool, dereferenceable and
66/// comparable.
67///
68/// SymbolStringPtrBases are default-constructible and constructible
69/// from nullptr to enable comparison with these values.
70class SymbolStringPtrBase {
71  friend class SymbolStringPool;
72  friend struct DenseMapInfo<SymbolStringPtr>;
73  friend struct DenseMapInfo<NonOwningSymbolStringPtr>;
74
75public:
76  SymbolStringPtrBase() = default;
77  SymbolStringPtrBase(std::nullptr_t) {}
78
79  explicit operator bool() const { return S; }
80
81  StringRef operator*() const { return S->first(); }
82
83  friend bool operator==(SymbolStringPtrBase LHS, SymbolStringPtrBase RHS) {
84    return LHS.S == RHS.S;
85  }
86
87  friend bool operator!=(SymbolStringPtrBase LHS, SymbolStringPtrBase RHS) {
88    return !(LHS == RHS);
89  }
90
91  friend bool operator<(SymbolStringPtrBase LHS, SymbolStringPtrBase RHS) {
92    return LHS.S < RHS.S;
93  }
94
95#ifndef NDEBUG
96  // Returns true if the pool entry's ref count is above zero (or if the entry
97  // is an empty or tombstone value). Useful for debugging and testing -- this
98  // method can be used to identify SymbolStringPtrs and
99  // NonOwningSymbolStringPtrs that are pointing to abandoned pool entries.
100  bool poolEntryIsAlive() const {
101    return isRealPoolEntry(S) ? S->getValue() != 0 : true;
102  }
103#endif
104
105protected:
106  using PoolEntry = SymbolStringPool::PoolMapEntry;
107  using PoolEntryPtr = PoolEntry *;
108
109  SymbolStringPtrBase(PoolEntryPtr S) : S(S) {}
110
111  constexpr static uintptr_t EmptyBitPattern =
112      std::numeric_limits<uintptr_t>::max()
113      << PointerLikeTypeTraits<PoolEntryPtr>::NumLowBitsAvailable;
114
115  constexpr static uintptr_t TombstoneBitPattern =
116      (std::numeric_limits<uintptr_t>::max() - 1)
117      << PointerLikeTypeTraits<PoolEntryPtr>::NumLowBitsAvailable;
118
119  constexpr static uintptr_t InvalidPtrMask =
120      (std::numeric_limits<uintptr_t>::max() - 3)
121      << PointerLikeTypeTraits<PoolEntryPtr>::NumLowBitsAvailable;
122
123  // Returns false for null, empty, and tombstone values, true otherwise.
124  static bool isRealPoolEntry(PoolEntryPtr P) {
125    return ((reinterpret_cast<uintptr_t>(P) - 1) & InvalidPtrMask) !=
126           InvalidPtrMask;
127  }
128
129  size_t getRefCount() const {
130    return isRealPoolEntry(S) ? size_t(S->getValue()) : size_t(0);
131  }
132
133  PoolEntryPtr S = nullptr;
134};
135
136/// Pointer to a pooled string representing a symbol name.
137class SymbolStringPtr : public SymbolStringPtrBase {
138  friend class SymbolStringPool;
139  friend class SymbolStringPoolEntryUnsafe;
140  friend struct DenseMapInfo<SymbolStringPtr>;
141
142public:
143  SymbolStringPtr() = default;
144  SymbolStringPtr(std::nullptr_t) {}
145  SymbolStringPtr(const SymbolStringPtr &Other) : SymbolStringPtrBase(Other.S) {
146    incRef();
147  }
148
149  explicit SymbolStringPtr(NonOwningSymbolStringPtr Other);
150
151  SymbolStringPtr& operator=(const SymbolStringPtr &Other) {
152    decRef();
153    S = Other.S;
154    incRef();
155    return *this;
156  }
157
158  SymbolStringPtr(SymbolStringPtr &&Other) { std::swap(S, Other.S); }
159
160  SymbolStringPtr& operator=(SymbolStringPtr &&Other) {
161    decRef();
162    S = nullptr;
163    std::swap(S, Other.S);
164    return *this;
165  }
166
167  ~SymbolStringPtr() { decRef(); }
168
169private:
170  SymbolStringPtr(PoolEntryPtr S) : SymbolStringPtrBase(S) { incRef(); }
171
172  void incRef() {
173    if (isRealPoolEntry(S))
174      ++S->getValue();
175  }
176
177  void decRef() {
178    if (isRealPoolEntry(S)) {
179      assert(S->getValue() && "Releasing SymbolStringPtr with zero ref count");
180      --S->getValue();
181    }
182  }
183
184  static SymbolStringPtr getEmptyVal() {
185    return SymbolStringPtr(reinterpret_cast<PoolEntryPtr>(EmptyBitPattern));
186  }
187
188  static SymbolStringPtr getTombstoneVal() {
189    return SymbolStringPtr(reinterpret_cast<PoolEntryPtr>(TombstoneBitPattern));
190  }
191};
192
193/// Provides unsafe access to ownership operations on SymbolStringPtr.
194/// This class can be used to manage SymbolStringPtr instances from C.
195class SymbolStringPoolEntryUnsafe {
196public:
197  using PoolEntry = SymbolStringPool::PoolMapEntry;
198
199  SymbolStringPoolEntryUnsafe(PoolEntry *E) : E(E) {}
200
201  /// Create an unsafe pool entry ref without changing the ref-count.
202  static SymbolStringPoolEntryUnsafe from(const SymbolStringPtr &S) {
203    return S.S;
204  }
205
206  /// Consumes the given SymbolStringPtr without releasing the pool entry.
207  static SymbolStringPoolEntryUnsafe take(SymbolStringPtr &&S) {
208    PoolEntry *E = nullptr;
209    std::swap(E, S.S);
210    return E;
211  }
212
213  PoolEntry *rawPtr() { return E; }
214
215  /// Creates a SymbolStringPtr for this entry, with the SymbolStringPtr
216  /// retaining the entry as usual.
217  SymbolStringPtr copyToSymbolStringPtr() { return SymbolStringPtr(E); }
218
219  /// Creates a SymbolStringPtr for this entry *without* performing a retain
220  /// operation during construction.
221  SymbolStringPtr moveToSymbolStringPtr() {
222    SymbolStringPtr S;
223    std::swap(S.S, E);
224    return S;
225  }
226
227  void retain() { ++E->getValue(); }
228  void release() { --E->getValue(); }
229
230private:
231  PoolEntry *E = nullptr;
232};
233
234/// Non-owning SymbolStringPool entry pointer. Instances are comparable with
235/// SymbolStringPtr instances and guaranteed to have the same hash, but do not
236/// affect the ref-count of the pooled string (and are therefore cheaper to
237/// copy).
238///
239/// NonOwningSymbolStringPtrs are silently invalidated if the pool entry's
240/// ref-count drops to zero, so they should only be used in contexts where a
241/// corresponding SymbolStringPtr is known to exist (which will guarantee that
242/// the ref-count stays above zero). E.g. in a graph where nodes are
243/// represented by SymbolStringPtrs the edges can be represented by pairs of
244/// NonOwningSymbolStringPtrs and this will make the introduction of deletion
245/// of edges cheaper.
246class NonOwningSymbolStringPtr : public SymbolStringPtrBase {
247  friend struct DenseMapInfo<orc::NonOwningSymbolStringPtr>;
248
249public:
250  NonOwningSymbolStringPtr() = default;
251  explicit NonOwningSymbolStringPtr(const SymbolStringPtr &S)
252      : SymbolStringPtrBase(S) {}
253
254  using SymbolStringPtrBase::operator=;
255
256private:
257  NonOwningSymbolStringPtr(PoolEntryPtr S) : SymbolStringPtrBase(S) {}
258
259  static NonOwningSymbolStringPtr getEmptyVal() {
260    return NonOwningSymbolStringPtr(
261        reinterpret_cast<PoolEntryPtr>(EmptyBitPattern));
262  }
263
264  static NonOwningSymbolStringPtr getTombstoneVal() {
265    return NonOwningSymbolStringPtr(
266        reinterpret_cast<PoolEntryPtr>(TombstoneBitPattern));
267  }
268};
269
270inline SymbolStringPtr::SymbolStringPtr(NonOwningSymbolStringPtr Other)
271    : SymbolStringPtrBase(Other) {
272  assert(poolEntryIsAlive() &&
273         "SymbolStringPtr constructed from invalid non-owning pointer.");
274
275  if (isRealPoolEntry(S))
276    ++S->getValue();
277}
278
279inline SymbolStringPool::~SymbolStringPool() {
280#ifndef NDEBUG
281  clearDeadEntries();
282  assert(Pool.empty() && "Dangling references at pool destruction time");
283#endif // NDEBUG
284}
285
286inline SymbolStringPtr SymbolStringPool::intern(StringRef S) {
287  std::lock_guard<std::mutex> Lock(PoolMutex);
288  PoolMap::iterator I;
289  bool Added;
290  std::tie(I, Added) = Pool.try_emplace(S, 0);
291  return SymbolStringPtr(&*I);
292}
293
294inline void SymbolStringPool::clearDeadEntries() {
295  std::lock_guard<std::mutex> Lock(PoolMutex);
296  for (auto I = Pool.begin(), E = Pool.end(); I != E;) {
297    auto Tmp = I++;
298    if (Tmp->second == 0)
299      Pool.erase(Tmp);
300  }
301}
302
303inline bool SymbolStringPool::empty() const {
304  std::lock_guard<std::mutex> Lock(PoolMutex);
305  return Pool.empty();
306}
307
308inline size_t
309SymbolStringPool::getRefCount(const SymbolStringPtrBase &S) const {
310  return S.getRefCount();
311}
312
313} // end namespace orc
314
315template <>
316struct DenseMapInfo<orc::SymbolStringPtr> {
317
318  static orc::SymbolStringPtr getEmptyKey() {
319    return orc::SymbolStringPtr::getEmptyVal();
320  }
321
322  static orc::SymbolStringPtr getTombstoneKey() {
323    return orc::SymbolStringPtr::getTombstoneVal();
324  }
325
326  static unsigned getHashValue(const orc::SymbolStringPtrBase &V) {
327    return DenseMapInfo<orc::SymbolStringPtr::PoolEntryPtr>::getHashValue(V.S);
328  }
329
330  static bool isEqual(const orc::SymbolStringPtrBase &LHS,
331                      const orc::SymbolStringPtrBase &RHS) {
332    return LHS.S == RHS.S;
333  }
334};
335
336template <> struct DenseMapInfo<orc::NonOwningSymbolStringPtr> {
337
338  static orc::NonOwningSymbolStringPtr getEmptyKey() {
339    return orc::NonOwningSymbolStringPtr::getEmptyVal();
340  }
341
342  static orc::NonOwningSymbolStringPtr getTombstoneKey() {
343    return orc::NonOwningSymbolStringPtr::getTombstoneVal();
344  }
345
346  static unsigned getHashValue(const orc::SymbolStringPtrBase &V) {
347    return DenseMapInfo<
348        orc::NonOwningSymbolStringPtr::PoolEntryPtr>::getHashValue(V.S);
349  }
350
351  static bool isEqual(const orc::SymbolStringPtrBase &LHS,
352                      const orc::SymbolStringPtrBase &RHS) {
353    return LHS.S == RHS.S;
354  }
355};
356
357} // end namespace llvm
358
359#endif // LLVM_EXECUTIONENGINE_ORC_SYMBOLSTRINGPOOL_H
360