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