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