Add support for removing elements out of StringMap.
[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   
27 /// StringMapEntryBase - Shared base class of StringMapEntry instances.
28 class StringMapEntryBase {
29   unsigned StrLen;
30 public:
31   StringMapEntryBase(unsigned Len) : StrLen(Len) {}
32   
33   unsigned getKeyLength() const { return StrLen; }
34 };
35   
36 /// StringMapImpl - This is the base class of StringMap that is shared among
37 /// all of its instantiations.
38 class StringMapImpl {
39 public:
40   /// ItemBucket - The hash table consists of an array of these.  If Item is
41   /// non-null, this is an extant entry, otherwise, it is a hole.
42   struct ItemBucket {
43     /// FullHashValue - This remembers the full hash value of the key for
44     /// easy scanning.
45     unsigned FullHashValue;
46     
47     /// Item - This is a pointer to the actual item object.
48     StringMapEntryBase *Item;
49   };
50   
51 protected:
52   ItemBucket *TheTable;
53   unsigned NumBuckets;
54   unsigned NumItems;
55   unsigned NumTombstones;
56   unsigned ItemSize;
57 protected:
58   StringMapImpl(unsigned InitSize, unsigned ItemSize);
59   void RehashTable();
60   
61   /// LookupBucketFor - Look up the bucket that the specified string should end
62   /// up in.  If it already exists as a key in the map, the Item pointer for the
63   /// specified bucket will be non-null.  Otherwise, it will be null.  In either
64   /// case, the FullHashValue field of the bucket will be set to the hash value
65   /// of the string.
66   unsigned LookupBucketFor(const char *KeyStart, const char *KeyEnd);
67   
68   /// FindKey - Look up the bucket that contains the specified key. If it exists
69   /// in the map, return the bucket number of the key.  Otherwise return -1.
70   /// This does not modify the map.
71   int FindKey(const char *KeyStart, const char *KeyEnd) const;
72
73   /// RemoveKey - Remove the specified StringMapEntry from the table, but do not
74   /// delete it.  This aborts if the value isn't in the table.
75   void RemoveKey(StringMapEntryBase *V);
76
77   /// RemoveKey - Remove the StringMapEntry for the specified key from the
78   /// table, returning it.  If the key is not in the table, this returns null.
79   StringMapEntryBase *RemoveKey(const char *KeyStart, const char *KeyEnd);
80   
81 public:
82   static StringMapEntryBase *getTombstoneVal() {
83     return (StringMapEntryBase*)-1;
84   }
85   
86   unsigned getNumBuckets() const { return NumBuckets; }
87   unsigned getNumItems() const { return NumItems; }
88
89   bool empty() const { return NumItems == 0; }
90   unsigned size() const { return NumItems; }
91 };
92
93 /// StringMapEntry - This is used to represent one value that is inserted into
94 /// a StringMap.  It contains the Value itself and the key: the string length
95 /// and data.
96 template<typename ValueTy>
97 class StringMapEntry : public StringMapEntryBase {
98   ValueTy Val;
99 public:
100   StringMapEntry(unsigned StrLen)
101     : StringMapEntryBase(StrLen), Val() {}
102   StringMapEntry(unsigned StrLen, const ValueTy &V)
103     : StringMapEntryBase(StrLen), Val(V) {}
104
105   const ValueTy &getValue() const { return Val; }
106   ValueTy &getValue() { return Val; }
107   
108   void setValue(const ValueTy &V) { Val = V; }
109   
110   /// getKeyData - Return the start of the string data that is the key for this
111   /// value.  The string data is always stored immediately after the
112   /// StringMapEntry object.
113   const char *getKeyData() const {return reinterpret_cast<const char*>(this+1);}
114   
115   /// Create - Create a StringMapEntry for the specified key and default
116   /// construct the value.
117   template<typename AllocatorTy>
118   static StringMapEntry *Create(const char *KeyStart, const char *KeyEnd,
119                                 AllocatorTy &Allocator) {
120     unsigned KeyLength = KeyEnd-KeyStart;
121     
122     // Okay, the item doesn't already exist, and 'Bucket' is the bucket to fill
123     // in.  Allocate a new item with space for the string at the end and a null
124     // terminator.
125     unsigned AllocSize = sizeof(StringMapEntry)+KeyLength+1;
126     
127 #ifdef __GNUC__
128     unsigned Alignment = __alignof__(StringMapEntry);
129 #else
130     // FIXME: ugly.
131     unsigned Alignment = 8;
132 #endif
133     StringMapEntry *NewItem = 
134       static_cast<StringMapEntry*>(Allocator.Allocate(AllocSize, Alignment));
135     
136     // Default construct the value.
137     new (NewItem) StringMapEntry(KeyLength);
138     
139     // Copy the string information.
140     char *StrBuffer = const_cast<char*>(NewItem->getKeyData());
141     memcpy(StrBuffer, KeyStart, KeyLength);
142     StrBuffer[KeyLength] = 0;  // Null terminate for convenience of clients.
143     return NewItem;
144   }
145   
146   /// Create - Create a StringMapEntry with normal malloc/free.
147   static StringMapEntry *Create(const char *KeyStart, const char *KeyEnd) {
148     MallocAllocator A;
149     return Create(KeyStart, KeyEnd, A);
150   }
151
152   /// Destroy - Destroy this StringMapEntry, releasing memory back to the
153   /// specified allocator.
154   template<typename AllocatorTy>
155   void Destroy(AllocatorTy &Allocator) {
156     // Free memory referenced by the item.
157     this->~StringMapEntry();
158     Allocator.Deallocate(this);
159   }
160   
161   /// Destroy this object, releasing memory back to the malloc allocator.
162   void Destroy() {
163     MallocAllocator A;
164     Destroy(A);
165   }
166 };
167
168
169 /// StringMap - This is an unconventional map that is specialized for handling
170 /// keys that are "strings", which are basically ranges of bytes. This does some
171 /// funky memory allocation and hashing things to make it extremely efficient,
172 /// storing the string data *after* the value in the map.
173 template<typename ValueTy, typename AllocatorTy = MallocAllocator>
174 class StringMap : public StringMapImpl {
175   AllocatorTy Allocator;
176   typedef StringMapEntry<ValueTy> MapEntryTy;
177 public:
178   StringMap(unsigned InitialSize = 0)
179     : StringMapImpl(InitialSize, sizeof(MapEntryTy)) {}
180   
181   AllocatorTy &getAllocator() { return Allocator; }
182   const AllocatorTy &getAllocator() const { return Allocator; }
183
184   typedef StringMapConstIterator<ValueTy> const_iterator;
185   typedef StringMapIterator<ValueTy> iterator;
186   
187   iterator begin() { return iterator(TheTable); }
188   iterator end() { return iterator(TheTable+NumBuckets); }
189   const_iterator begin() const { return const_iterator(TheTable); }
190   const_iterator end() const { return const_iterator(TheTable+NumBuckets); }
191   
192   
193   iterator find(const char *KeyStart, const char *KeyEnd) {
194     int Bucket = FindKey(KeyStart, KeyEnd);
195     if (Bucket == -1) return end();
196     return iterator(TheTable+Bucket);
197   }
198
199   const_iterator find(const char *KeyStart, const char *KeyEnd) const {
200     int Bucket = FindKey(KeyStart, KeyEnd);
201     if (Bucket == -1) return end();
202     return const_iterator(TheTable+Bucket);
203   }
204   
205   /// insert - Insert the specified key/value pair into the map.  If the key
206   /// already exists in the map, return false and ignore the request, otherwise
207   /// insert it and return true.
208   bool insert(MapEntryTy *KeyValue) {
209     unsigned BucketNo =
210       LookupBucketFor(KeyValue->getKeyData(),
211                       KeyValue->getKeyData()+KeyValue->getKeyLength());
212     ItemBucket &Bucket = TheTable[BucketNo];
213     if (Bucket.Item && Bucket.Item != getTombstoneVal()) 
214       return false;  // Already exists in map.
215     
216     if (Bucket.Item == getTombstoneVal())
217       --NumTombstones;
218     Bucket.Item = KeyValue;
219     ++NumItems;
220     
221     // If the hash table is now more than 3/4 full, rehash into a larger table.
222     if (NumItems > NumBuckets*3/4)
223       RehashTable();
224     return true;
225   }
226   
227   /// GetOrCreateValue - Look up the specified key in the table.  If a value
228   /// exists, return it.  Otherwise, default construct a value, insert it, and
229   /// return.
230   StringMapEntry<ValueTy> &GetOrCreateValue(const char *KeyStart, 
231                                             const char *KeyEnd) {
232     unsigned BucketNo = LookupBucketFor(KeyStart, KeyEnd);
233     ItemBucket &Bucket = TheTable[BucketNo];
234     if (Bucket.Item && Bucket.Item != getTombstoneVal())
235       return *static_cast<MapEntryTy*>(Bucket.Item);
236     
237     MapEntryTy *NewItem = MapEntryTy::Create(KeyStart, KeyEnd, Allocator);
238     
239     if (Bucket.Item == getTombstoneVal())
240       --NumTombstones;
241     ++NumItems;
242     
243     // Fill in the bucket for the hash table.  The FullHashValue was already
244     // filled in by LookupBucketFor.
245     Bucket.Item = NewItem;
246     
247     // If the hash table is now more than 3/4 full, rehash into a larger table.
248     if (NumItems > NumBuckets*3/4)
249       RehashTable();
250     return *NewItem;
251   }
252   
253   /// remove - Remove the specified key/value pair from the map, but do not
254   /// erase it.  This aborts if the key is not in the map.
255   void remove(MapEntryTy *KeyValue) {
256     RemoveKey(KeyValue);
257   }
258   
259   ~StringMap() {
260     for (ItemBucket *I = TheTable, *E = TheTable+NumBuckets; I != E; ++I) {
261       if (MapEntryTy *Id = static_cast<MapEntryTy*>(I->Item))
262         Id->Destroy(Allocator);
263     }
264     delete [] TheTable;
265   }
266 };
267   
268
269 template<typename ValueTy>
270 class StringMapConstIterator {
271 protected:
272   StringMapImpl::ItemBucket *Ptr;
273 public:
274   StringMapConstIterator(StringMapImpl::ItemBucket *Bucket) : Ptr(Bucket) {
275     AdvancePastEmptyBuckets();
276   }
277   
278   const StringMapEntry<ValueTy> &operator*() const {
279     return *static_cast<StringMapEntry<ValueTy>*>(Ptr->Item);
280   }
281   const StringMapEntry<ValueTy> *operator->() const {
282     return static_cast<StringMapEntry<ValueTy>*>(Ptr->Item);
283   }
284   
285   bool operator==(const StringMapConstIterator &RHS) const {
286     return Ptr == RHS.Ptr;
287   }
288   bool operator!=(const StringMapConstIterator &RHS) const {
289     return Ptr != RHS.Ptr;
290   }
291   
292   inline StringMapConstIterator& operator++() {          // Preincrement
293     ++Ptr;
294     AdvancePastEmptyBuckets();
295     return *this;
296   }
297   StringMapConstIterator operator++(int) {        // Postincrement
298     StringMapConstIterator tmp = *this; ++*this; return tmp;
299   }
300   
301 private:
302   void AdvancePastEmptyBuckets() {
303     while (Ptr->Item == 0 || Ptr->Item == StringMapImpl::getTombstoneVal())
304       ++Ptr;
305   }
306 };
307
308 template<typename ValueTy>
309 class StringMapIterator : public StringMapConstIterator<ValueTy> {
310 public:  
311   StringMapIterator(StringMapImpl::ItemBucket *Bucket)
312     : StringMapConstIterator<ValueTy>(Bucket) {
313   }
314   StringMapEntry<ValueTy> &operator*() const {
315     return *static_cast<StringMapEntry<ValueTy>*>(this->Ptr->Item);
316   }
317   StringMapEntry<ValueTy> *operator->() const {
318     return static_cast<StringMapEntry<ValueTy>*>(this->Ptr->Item);
319   }
320 };
321
322 }
323
324 #endif
325