194235fac57043b263302bc7f81a2bfc36bfb201
[oota-llvm.git] / include / llvm / ADT / StringMap.h
1 //===--- StringMap.h - String Hash table map interface ----------*- C++ -*-===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // This file defines the StringMap class.
11 //
12 //===----------------------------------------------------------------------===//
13
14 #ifndef LLVM_ADT_STRINGMAP_H
15 #define LLVM_ADT_STRINGMAP_H
16
17 #include "llvm/ADT/StringRef.h"
18 #include "llvm/Support/Allocator.h"
19 #include <cstring>
20 #include <utility>
21
22 namespace llvm {
23   template<typename ValueT>
24   class StringMapConstIterator;
25   template<typename ValueT>
26   class StringMapIterator;
27   template<typename ValueTy>
28   class StringMapEntry;
29
30 /// StringMapEntryBase - Shared base class of StringMapEntry instances.
31 class StringMapEntryBase {
32   unsigned StrLen;
33
34 public:
35   explicit StringMapEntryBase(unsigned Len) : StrLen(Len) {}
36
37   unsigned getKeyLength() const { return StrLen; }
38 };
39
40 /// StringMapImpl - This is the base class of StringMap that is shared among
41 /// all of its instantiations.
42 class StringMapImpl {
43 protected:
44   // Array of NumBuckets pointers to entries, null pointers are holes.
45   // TheTable[NumBuckets] contains a sentinel value for easy iteration. Followed
46   // by an array of the actual hash values as unsigned integers.
47   StringMapEntryBase **TheTable;
48   unsigned NumBuckets;
49   unsigned NumItems;
50   unsigned NumTombstones;
51   unsigned ItemSize;
52
53 protected:
54   explicit StringMapImpl(unsigned itemSize)
55       : TheTable(nullptr),
56         // Initialize the map with zero buckets to allocation.
57         NumBuckets(0), NumItems(0), NumTombstones(0), ItemSize(itemSize) {}
58   StringMapImpl(StringMapImpl &&RHS)
59       : TheTable(RHS.TheTable), NumBuckets(RHS.NumBuckets),
60         NumItems(RHS.NumItems), NumTombstones(RHS.NumTombstones),
61         ItemSize(RHS.ItemSize) {
62     RHS.TheTable = nullptr;
63     RHS.NumBuckets = 0;
64     RHS.NumItems = 0;
65     RHS.NumTombstones = 0;
66   }
67
68   StringMapImpl(unsigned InitSize, unsigned ItemSize);
69   unsigned RehashTable(unsigned BucketNo = 0);
70
71   /// LookupBucketFor - Look up the bucket that the specified string should end
72   /// up in.  If it already exists as a key in the map, the Item pointer for the
73   /// specified bucket will be non-null.  Otherwise, it will be null.  In either
74   /// case, the FullHashValue field of the bucket will be set to the hash value
75   /// of the string.
76   unsigned LookupBucketFor(StringRef Key);
77
78   /// FindKey - Look up the bucket that contains the specified key. If it exists
79   /// in the map, return the bucket number of the key.  Otherwise return -1.
80   /// This does not modify the map.
81   int FindKey(StringRef Key) const;
82
83   /// RemoveKey - Remove the specified StringMapEntry from the table, but do not
84   /// delete it.  This aborts if the value isn't in the table.
85   void RemoveKey(StringMapEntryBase *V);
86
87   /// RemoveKey - Remove the StringMapEntry for the specified key from the
88   /// table, returning it.  If the key is not in the table, this returns null.
89   StringMapEntryBase *RemoveKey(StringRef Key);
90
91 private:
92   void init(unsigned Size);
93
94 public:
95   static StringMapEntryBase *getTombstoneVal() {
96     return (StringMapEntryBase*)-1;
97   }
98
99   unsigned getNumBuckets() const { return NumBuckets; }
100   unsigned getNumItems() const { return NumItems; }
101
102   bool empty() const { return NumItems == 0; }
103   unsigned size() const { return NumItems; }
104
105   void swap(StringMapImpl &Other) {
106     std::swap(TheTable, Other.TheTable);
107     std::swap(NumBuckets, Other.NumBuckets);
108     std::swap(NumItems, Other.NumItems);
109     std::swap(NumTombstones, Other.NumTombstones);
110   }
111 };
112
113 /// StringMapEntry - This is used to represent one value that is inserted into
114 /// a StringMap.  It contains the Value itself and the key: the string length
115 /// and data.
116 template<typename ValueTy>
117 class StringMapEntry : public StringMapEntryBase {
118   StringMapEntry(StringMapEntry &E) = delete;
119
120 public:
121   ValueTy second;
122
123   explicit StringMapEntry(unsigned strLen)
124     : StringMapEntryBase(strLen), second() {}
125   template <class InitTy>
126   StringMapEntry(unsigned strLen, InitTy &&V)
127       : StringMapEntryBase(strLen), second(std::forward<InitTy>(V)) {}
128
129   StringRef getKey() const {
130     return StringRef(getKeyData(), getKeyLength());
131   }
132
133   const ValueTy &getValue() const { return second; }
134   ValueTy &getValue() { return second; }
135
136   void setValue(const ValueTy &V) { second = V; }
137
138   /// getKeyData - Return the start of the string data that is the key for this
139   /// value.  The string data is always stored immediately after the
140   /// StringMapEntry object.
141   const char *getKeyData() const {return reinterpret_cast<const char*>(this+1);}
142
143   StringRef first() const { return StringRef(getKeyData(), getKeyLength()); }
144
145   /// Create - Create a StringMapEntry for the specified key and default
146   /// construct the value.
147   template <typename AllocatorTy, typename InitType>
148   static StringMapEntry *Create(StringRef Key, AllocatorTy &Allocator,
149                                 InitType &&InitVal) {
150     unsigned KeyLength = Key.size();
151
152     // Allocate a new item with space for the string at the end and a null
153     // terminator.
154     unsigned AllocSize = static_cast<unsigned>(sizeof(StringMapEntry))+
155       KeyLength+1;
156     unsigned Alignment = alignOf<StringMapEntry>();
157
158     StringMapEntry *NewItem =
159       static_cast<StringMapEntry*>(Allocator.Allocate(AllocSize,Alignment));
160
161     // Default construct the value.
162     new (NewItem) StringMapEntry(KeyLength, std::forward<InitType>(InitVal));
163
164     // Copy the string information.
165     char *StrBuffer = const_cast<char*>(NewItem->getKeyData());
166     if (KeyLength > 0)
167       memcpy(StrBuffer, Key.data(), KeyLength);
168     StrBuffer[KeyLength] = 0;  // Null terminate for convenience of clients.
169     return NewItem;
170   }
171
172   template<typename AllocatorTy>
173   static StringMapEntry *Create(StringRef Key, AllocatorTy &Allocator) {
174     return Create(Key, Allocator, ValueTy());
175   }
176
177   /// Create - Create a StringMapEntry with normal malloc/free.
178   template<typename InitType>
179   static StringMapEntry *Create(StringRef Key, InitType &&InitVal) {
180     MallocAllocator A;
181     return Create(Key, A, std::forward<InitType>(InitVal));
182   }
183
184   static StringMapEntry *Create(StringRef Key) {
185     return Create(Key, ValueTy());
186   }
187
188   /// GetStringMapEntryFromKeyData - Given key data that is known to be embedded
189   /// into a StringMapEntry, return the StringMapEntry itself.
190   static StringMapEntry &GetStringMapEntryFromKeyData(const char *KeyData) {
191     char *Ptr = const_cast<char*>(KeyData) - sizeof(StringMapEntry<ValueTy>);
192     return *reinterpret_cast<StringMapEntry*>(Ptr);
193   }
194
195   /// Destroy - Destroy this StringMapEntry, releasing memory back to the
196   /// specified allocator.
197   template<typename AllocatorTy>
198   void Destroy(AllocatorTy &Allocator) {
199     // Free memory referenced by the item.
200     unsigned AllocSize =
201         static_cast<unsigned>(sizeof(StringMapEntry)) + getKeyLength() + 1;
202     this->~StringMapEntry();
203     Allocator.Deallocate(static_cast<void *>(this), AllocSize);
204   }
205
206   /// Destroy this object, releasing memory back to the malloc allocator.
207   void Destroy() {
208     MallocAllocator A;
209     Destroy(A);
210   }
211 };
212
213 /// StringMap - This is an unconventional map that is specialized for handling
214 /// keys that are "strings", which are basically ranges of bytes. This does some
215 /// funky memory allocation and hashing things to make it extremely efficient,
216 /// storing the string data *after* the value in the map.
217 template<typename ValueTy, typename AllocatorTy = MallocAllocator>
218 class StringMap : public StringMapImpl {
219   AllocatorTy Allocator;
220
221 public:
222   typedef StringMapEntry<ValueTy> MapEntryTy;
223
224   StringMap() : StringMapImpl(static_cast<unsigned>(sizeof(MapEntryTy))) {}
225   explicit StringMap(unsigned InitialSize)
226     : StringMapImpl(InitialSize, static_cast<unsigned>(sizeof(MapEntryTy))) {}
227
228   explicit StringMap(AllocatorTy A)
229     : StringMapImpl(static_cast<unsigned>(sizeof(MapEntryTy))), Allocator(A) {}
230
231   StringMap(unsigned InitialSize, AllocatorTy A)
232     : StringMapImpl(InitialSize, static_cast<unsigned>(sizeof(MapEntryTy))),
233       Allocator(A) {}
234
235   StringMap(StringMap &&RHS)
236       : StringMapImpl(std::move(RHS)), Allocator(std::move(RHS.Allocator)) {}
237
238   StringMap &operator=(StringMap RHS) {
239     StringMapImpl::swap(RHS);
240     std::swap(Allocator, RHS.Allocator);
241     return *this;
242   }
243
244   // FIXME: Implement copy operations if/when they're needed.
245
246   AllocatorTy &getAllocator() { return Allocator; }
247   const AllocatorTy &getAllocator() const { return Allocator; }
248
249   typedef const char* key_type;
250   typedef ValueTy mapped_type;
251   typedef StringMapEntry<ValueTy> value_type;
252   typedef size_t size_type;
253
254   typedef StringMapConstIterator<ValueTy> const_iterator;
255   typedef StringMapIterator<ValueTy> iterator;
256
257   iterator begin() {
258     return iterator(TheTable, NumBuckets == 0);
259   }
260   iterator end() {
261     return iterator(TheTable+NumBuckets, true);
262   }
263   const_iterator begin() const {
264     return const_iterator(TheTable, NumBuckets == 0);
265   }
266   const_iterator end() const {
267     return const_iterator(TheTable+NumBuckets, true);
268   }
269
270   iterator find(StringRef Key) {
271     int Bucket = FindKey(Key);
272     if (Bucket == -1) return end();
273     return iterator(TheTable+Bucket, true);
274   }
275
276   const_iterator find(StringRef Key) const {
277     int Bucket = FindKey(Key);
278     if (Bucket == -1) return end();
279     return const_iterator(TheTable+Bucket, true);
280   }
281
282   /// lookup - Return the entry for the specified key, or a default
283   /// constructed value if no such entry exists.
284   ValueTy lookup(StringRef Key) const {
285     const_iterator it = find(Key);
286     if (it != end())
287       return it->second;
288     return ValueTy();
289   }
290
291   ValueTy &operator[](StringRef Key) {
292     return insert(std::make_pair(Key, ValueTy())).first->second;
293   }
294
295   /// count - Return 1 if the element is in the map, 0 otherwise.
296   size_type count(StringRef Key) const {
297     return find(Key) == end() ? 0 : 1;
298   }
299
300   /// insert - Insert the specified key/value pair into the map.  If the key
301   /// already exists in the map, return false and ignore the request, otherwise
302   /// insert it and return true.
303   bool insert(MapEntryTy *KeyValue) {
304     unsigned BucketNo = LookupBucketFor(KeyValue->getKey());
305     StringMapEntryBase *&Bucket = TheTable[BucketNo];
306     if (Bucket && Bucket != getTombstoneVal())
307       return false;  // Already exists in map.
308
309     if (Bucket == getTombstoneVal())
310       --NumTombstones;
311     Bucket = KeyValue;
312     ++NumItems;
313     assert(NumItems + NumTombstones <= NumBuckets);
314
315     RehashTable();
316     return true;
317   }
318
319   /// insert - Inserts the specified key/value pair into the map if the key
320   /// isn't already in the map. The bool component of the returned pair is true
321   /// if and only if the insertion takes place, and the iterator component of
322   /// the pair points to the element with key equivalent to the key of the pair.
323   std::pair<iterator, bool> insert(std::pair<StringRef, ValueTy> KV) {
324     unsigned BucketNo = LookupBucketFor(KV.first);
325     StringMapEntryBase *&Bucket = TheTable[BucketNo];
326     if (Bucket && Bucket != getTombstoneVal())
327       return std::make_pair(iterator(TheTable + BucketNo, false),
328                             false); // Already exists in map.
329
330     if (Bucket == getTombstoneVal())
331       --NumTombstones;
332     Bucket =
333         MapEntryTy::Create(KV.first, Allocator, std::move(KV.second));
334     ++NumItems;
335     assert(NumItems + NumTombstones <= NumBuckets);
336
337     BucketNo = RehashTable(BucketNo);
338     return std::make_pair(iterator(TheTable + BucketNo, false), true);
339   }
340
341   // clear - Empties out the StringMap
342   void clear() {
343     if (empty()) return;
344
345     // Zap all values, resetting the keys back to non-present (not tombstone),
346     // which is safe because we're removing all elements.
347     for (unsigned I = 0, E = NumBuckets; I != E; ++I) {
348       StringMapEntryBase *&Bucket = TheTable[I];
349       if (Bucket && Bucket != getTombstoneVal()) {
350         static_cast<MapEntryTy*>(Bucket)->Destroy(Allocator);
351       }
352       Bucket = nullptr;
353     }
354
355     NumItems = 0;
356     NumTombstones = 0;
357   }
358
359   /// remove - Remove the specified key/value pair from the map, but do not
360   /// erase it.  This aborts if the key is not in the map.
361   void remove(MapEntryTy *KeyValue) {
362     RemoveKey(KeyValue);
363   }
364
365   void erase(iterator I) {
366     MapEntryTy &V = *I;
367     remove(&V);
368     V.Destroy(Allocator);
369   }
370
371   bool erase(StringRef Key) {
372     iterator I = find(Key);
373     if (I == end()) return false;
374     erase(I);
375     return true;
376   }
377
378   ~StringMap() {
379     // Delete all the elements in the map, but don't reset the elements
380     // to default values.  This is a copy of clear(), but avoids unnecessary
381     // work not required in the destructor.
382     if (!empty()) {
383       for (unsigned I = 0, E = NumBuckets; I != E; ++I) {
384         StringMapEntryBase *Bucket = TheTable[I];
385         if (Bucket && Bucket != getTombstoneVal()) {
386           static_cast<MapEntryTy*>(Bucket)->Destroy(Allocator);
387         }
388       }
389     }
390     free(TheTable);
391   }
392 };
393
394 template <typename ValueTy> class StringMapConstIterator {
395 protected:
396   StringMapEntryBase **Ptr;
397
398 public:
399   typedef StringMapEntry<ValueTy> value_type;
400
401   StringMapConstIterator() : Ptr(nullptr) { }
402
403   explicit StringMapConstIterator(StringMapEntryBase **Bucket,
404                                   bool NoAdvance = false)
405   : Ptr(Bucket) {
406     if (!NoAdvance) AdvancePastEmptyBuckets();
407   }
408
409   const value_type &operator*() const {
410     return *static_cast<StringMapEntry<ValueTy>*>(*Ptr);
411   }
412   const value_type *operator->() const {
413     return static_cast<StringMapEntry<ValueTy>*>(*Ptr);
414   }
415
416   bool operator==(const StringMapConstIterator &RHS) const {
417     return Ptr == RHS.Ptr;
418   }
419   bool operator!=(const StringMapConstIterator &RHS) const {
420     return Ptr != RHS.Ptr;
421   }
422
423   inline StringMapConstIterator& operator++() {   // Preincrement
424     ++Ptr;
425     AdvancePastEmptyBuckets();
426     return *this;
427   }
428   StringMapConstIterator operator++(int) {        // Postincrement
429     StringMapConstIterator tmp = *this; ++*this; return tmp;
430   }
431
432 private:
433   void AdvancePastEmptyBuckets() {
434     while (*Ptr == nullptr || *Ptr == StringMapImpl::getTombstoneVal())
435       ++Ptr;
436   }
437 };
438
439 template<typename ValueTy>
440 class StringMapIterator : public StringMapConstIterator<ValueTy> {
441 public:
442   StringMapIterator() {}
443   explicit StringMapIterator(StringMapEntryBase **Bucket,
444                              bool NoAdvance = false)
445     : StringMapConstIterator<ValueTy>(Bucket, NoAdvance) {
446   }
447   StringMapEntry<ValueTy> &operator*() const {
448     return *static_cast<StringMapEntry<ValueTy>*>(*this->Ptr);
449   }
450   StringMapEntry<ValueTy> *operator->() const {
451     return static_cast<StringMapEntry<ValueTy>*>(*this->Ptr);
452   }
453 };
454 }
455
456 #endif