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