1 //===------- string_pool.h - Thread-safe pool for strings -------*- 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. Strings are ref-counted, but not
10 // automatically deallocated. Unused entries can be cleared by calling
11 // StringPool::clearDeadEntries.
12 //
13 //===----------------------------------------------------------------------===//
14 
15 #ifndef ORC_RT_STRING_POOL_H
16 #define ORC_RT_STRING_POOL_H
17 
18 #include <atomic>
19 #include <cassert>
20 #include <functional>
21 #include <mutex>
22 #include <string>
23 #include <unordered_map>
24 
25 namespace __orc_rt {
26 
27 class PooledStringPtr;
28 
29 /// String pool for strings names used by the ORC runtime.
30 class StringPool {
31   friend class PooledStringPtr;
32 
33 public:
34   /// Destroy a StringPool.
35   ~StringPool();
36 
37   /// Create a string pointer from the given string.
38   PooledStringPtr intern(std::string S);
39 
40   /// Remove from the pool any entries that are no longer referenced.
41   void clearDeadEntries();
42 
43   /// Returns true if the pool is empty.
44   bool empty() const;
45 
46 private:
47   using RefCountType = std::atomic<size_t>;
48   using PoolMap = std::unordered_map<std::string, RefCountType>;
49   using PoolMapEntry = PoolMap::value_type;
50   mutable std::mutex PoolMutex;
51   PoolMap Pool;
52 };
53 
54 /// Pointer to a pooled string.
55 class PooledStringPtr {
56   friend class StringPool;
57   friend struct std::hash<PooledStringPtr>;
58 
59 public:
60   PooledStringPtr() = default;
61   PooledStringPtr(std::nullptr_t) {}
62   PooledStringPtr(const PooledStringPtr &Other) : S(Other.S) {
63     if (S)
64       ++S->second;
65   }
66 
67   PooledStringPtr &operator=(const PooledStringPtr &Other) {
68     if (S) {
69       assert(S->second && "Releasing PooledStringPtr with zero ref count");
70       --S->second;
71     }
72     S = Other.S;
73     if (S)
74       ++S->second;
75     return *this;
76   }
77 
78   PooledStringPtr(PooledStringPtr &&Other) : S(nullptr) {
79     std::swap(S, Other.S);
80   }
81 
82   PooledStringPtr &operator=(PooledStringPtr &&Other) {
83     if (S) {
84       assert(S->second && "Releasing PooledStringPtr with zero ref count");
85       --S->second;
86     }
87     S = nullptr;
88     std::swap(S, Other.S);
89     return *this;
90   }
91 
92   ~PooledStringPtr() {
93     if (S) {
94       assert(S->second && "Releasing PooledStringPtr with zero ref count");
95       --S->second;
96     }
97   }
98 
99   explicit operator bool() const { return S; }
100 
101   const std::string &operator*() const { return S->first; }
102 
103   friend bool operator==(const PooledStringPtr &LHS,
104                          const PooledStringPtr &RHS) {
105     return LHS.S == RHS.S;
106   }
107 
108   friend bool operator!=(const PooledStringPtr &LHS,
109                          const PooledStringPtr &RHS) {
110     return !(LHS == RHS);
111   }
112 
113   friend bool operator<(const PooledStringPtr &LHS,
114                         const PooledStringPtr &RHS) {
115     return LHS.S < RHS.S;
116   }
117 
118 private:
119   using PoolEntry = StringPool::PoolMapEntry;
120   using PoolEntryPtr = PoolEntry *;
121 
122   PooledStringPtr(StringPool::PoolMapEntry *S) : S(S) {
123     if (S)
124       ++S->second;
125   }
126 
127   PoolEntryPtr S = nullptr;
128 };
129 
130 inline StringPool::~StringPool() {
131 #ifndef NDEBUG
132   clearDeadEntries();
133   assert(Pool.empty() && "Dangling references at pool destruction time");
134 #endif // NDEBUG
135 }
136 
137 inline PooledStringPtr StringPool::intern(std::string S) {
138   std::lock_guard<std::mutex> Lock(PoolMutex);
139   PoolMap::iterator I;
140   bool Added;
141   std::tie(I, Added) = Pool.try_emplace(std::move(S), 0);
142   return PooledStringPtr(&*I);
143 }
144 
145 inline void StringPool::clearDeadEntries() {
146   std::lock_guard<std::mutex> Lock(PoolMutex);
147   for (auto I = Pool.begin(), E = Pool.end(); I != E;) {
148     auto Tmp = I++;
149     if (Tmp->second == 0)
150       Pool.erase(Tmp);
151   }
152 }
153 
154 inline bool StringPool::empty() const {
155   std::lock_guard<std::mutex> Lock(PoolMutex);
156   return Pool.empty();
157 }
158 
159 } // end namespace __orc_rt
160 
161 namespace std {
162 
163 // Make PooledStringPtrs hashable.
164 template <> struct hash<__orc_rt::PooledStringPtr> {
165   size_t operator()(const __orc_rt::PooledStringPtr &A) const {
166     return hash<__orc_rt::PooledStringPtr::PoolEntryPtr>()(A.S);
167   }
168 };
169 
170 } // namespace std
171 
172 #endif // ORC_RT_REF_COUNTED_STRING_POOL_H
173