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