Add StringMap::lookup.
[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    /// lookup - Return the entry for the specified key, or a default
307   /// constructed value if no such entry exists.
308   ValueTy lookup(const char *KeyStart, const char *KeyEnd) const {
309     const_iterator it = find(KeyStart, KeyEnd);
310     if (it != end())
311       return it->second;
312     return ValueTy();
313   }
314   ValueTy lookup(const char *Key) const {
315     const_iterator it = find(Key);
316     if (it != end())
317       return it->second;
318     return ValueTy();
319   }
320   ValueTy lookup(const std::string &Key) const {
321     const_iterator it = find(Key);
322     if (it != end())
323       return it->second;
324     return ValueTy();
325   }
326
327   ValueTy& operator[](const char *Key) {
328     return GetOrCreateValue(Key, Key + strlen(Key)).getValue();
329   }
330   ValueTy& operator[](const std::string &Key) {
331     const char* key_start = (Key.empty() ? NULL : &Key[0]);
332     return GetOrCreateValue(key_start, key_start + Key.size()).getValue();
333   }
334
335   size_type count(const char *KeyStart, const char *KeyEnd) const {
336     return find(KeyStart, KeyEnd) == end() ? 0 : 1;
337   }
338   size_type count(const char *Key) const {
339     return count(Key, Key + strlen(Key));
340   }
341   size_type count(const std::string &Key) const {
342     const char* key_start = (Key.empty() ? NULL : &Key[0]);
343     return count(key_start, key_start + Key.size());
344   }
345
346   /// insert - Insert the specified key/value pair into the map.  If the key
347   /// already exists in the map, return false and ignore the request, otherwise
348   /// insert it and return true.
349   bool insert(MapEntryTy *KeyValue) {
350     unsigned BucketNo =
351       LookupBucketFor(KeyValue->getKeyData(),
352                       KeyValue->getKeyData()+KeyValue->getKeyLength());
353     ItemBucket &Bucket = TheTable[BucketNo];
354     if (Bucket.Item && Bucket.Item != getTombstoneVal())
355       return false;  // Already exists in map.
356
357     if (Bucket.Item == getTombstoneVal())
358       --NumTombstones;
359     Bucket.Item = KeyValue;
360     ++NumItems;
361
362     if (ShouldRehash())
363       RehashTable();
364     return true;
365   }
366
367   // clear - Empties out the StringMap
368   void clear() {
369     if (empty()) return;
370
371     // Zap all values, resetting the keys back to non-present (not tombstone),
372     // which is safe because we're removing all elements.
373     for (ItemBucket *I = TheTable, *E = TheTable+NumBuckets; I != E; ++I) {
374       if (I->Item && I->Item != getTombstoneVal()) {
375         static_cast<MapEntryTy*>(I->Item)->Destroy(Allocator);
376         I->Item = 0;
377       }
378     }
379
380     NumItems = 0;
381   }
382
383   /// GetOrCreateValue - Look up the specified key in the table.  If a value
384   /// exists, return it.  Otherwise, default construct a value, insert it, and
385   /// return.
386   template <typename InitTy>
387   StringMapEntry<ValueTy> &GetOrCreateValue(const char *KeyStart,
388                                             const char *KeyEnd,
389                                             InitTy Val) {
390     unsigned BucketNo = LookupBucketFor(KeyStart, KeyEnd);
391     ItemBucket &Bucket = TheTable[BucketNo];
392     if (Bucket.Item && Bucket.Item != getTombstoneVal())
393       return *static_cast<MapEntryTy*>(Bucket.Item);
394
395     MapEntryTy *NewItem = MapEntryTy::Create(KeyStart, KeyEnd, Allocator, Val);
396
397     if (Bucket.Item == getTombstoneVal())
398       --NumTombstones;
399     ++NumItems;
400
401     // Fill in the bucket for the hash table.  The FullHashValue was already
402     // filled in by LookupBucketFor.
403     Bucket.Item = NewItem;
404
405     if (ShouldRehash())
406       RehashTable();
407     return *NewItem;
408   }
409
410   StringMapEntry<ValueTy> &GetOrCreateValue(const char *KeyStart,
411                                             const char *KeyEnd) {
412     return GetOrCreateValue(KeyStart, KeyEnd, ValueTy());
413   }
414
415   /// remove - Remove the specified key/value pair from the map, but do not
416   /// erase it.  This aborts if the key is not in the map.
417   void remove(MapEntryTy *KeyValue) {
418     RemoveKey(KeyValue);
419   }
420
421   void erase(iterator I) {
422     MapEntryTy &V = *I;
423     remove(&V);
424     V.Destroy(Allocator);
425   }
426
427   bool erase(const char *Key) {
428     iterator I = find(Key);
429     if (I == end()) return false;
430     erase(I);
431     return true;
432   }
433
434   bool erase(const std::string &Key) {
435     iterator I = find(Key);
436     if (I == end()) return false;
437     erase(I);
438     return true;
439   }
440
441   ~StringMap() {
442     clear();
443     free(TheTable);
444   }
445 };
446
447
448 template<typename ValueTy>
449 class StringMapConstIterator {
450 protected:
451   StringMapImpl::ItemBucket *Ptr;
452 public:
453   typedef StringMapEntry<ValueTy> value_type;
454
455   explicit StringMapConstIterator(StringMapImpl::ItemBucket *Bucket,
456                                   bool NoAdvance = false)
457   : Ptr(Bucket) {
458     if (!NoAdvance) AdvancePastEmptyBuckets();
459   }
460
461   const value_type &operator*() const {
462     return *static_cast<StringMapEntry<ValueTy>*>(Ptr->Item);
463   }
464   const value_type *operator->() const {
465     return static_cast<StringMapEntry<ValueTy>*>(Ptr->Item);
466   }
467
468   bool operator==(const StringMapConstIterator &RHS) const {
469     return Ptr == RHS.Ptr;
470   }
471   bool operator!=(const StringMapConstIterator &RHS) const {
472     return Ptr != RHS.Ptr;
473   }
474
475   inline StringMapConstIterator& operator++() {          // Preincrement
476     ++Ptr;
477     AdvancePastEmptyBuckets();
478     return *this;
479   }
480   StringMapConstIterator operator++(int) {        // Postincrement
481     StringMapConstIterator tmp = *this; ++*this; return tmp;
482   }
483
484 private:
485   void AdvancePastEmptyBuckets() {
486     while (Ptr->Item == 0 || Ptr->Item == StringMapImpl::getTombstoneVal())
487       ++Ptr;
488   }
489 };
490
491 template<typename ValueTy>
492 class StringMapIterator : public StringMapConstIterator<ValueTy> {
493 public:
494   explicit StringMapIterator(StringMapImpl::ItemBucket *Bucket,
495                              bool NoAdvance = false)
496     : StringMapConstIterator<ValueTy>(Bucket, NoAdvance) {
497   }
498   StringMapEntry<ValueTy> &operator*() const {
499     return *static_cast<StringMapEntry<ValueTy>*>(this->Ptr->Item);
500   }
501   StringMapEntry<ValueTy> *operator->() const {
502     return static_cast<StringMapEntry<ValueTy>*>(this->Ptr->Item);
503   }
504 };
505
506 }
507
508 #endif