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