Replace the ugly FindValue method with STL-like find methods.
[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 ItemSize;
56 protected:
57   StringMapImpl(unsigned InitSize, unsigned ItemSize);
58   void RehashTable();
59   
60   /// LookupBucketFor - Look up the bucket that the specified string should end
61   /// up in.  If it already exists as a key in the map, the Item pointer for the
62   /// specified bucket will be non-null.  Otherwise, it will be null.  In either
63   /// case, the FullHashValue field of the bucket will be set to the hash value
64   /// of the string.
65   unsigned LookupBucketFor(const char *KeyStart, const char *KeyEnd);
66   
67   /// FindKey - Look up the bucket that contains the specified key. If it exists
68   /// in the map, return the bucket number of the key.  Otherwise return -1.
69   /// This does not modify the map.
70   int FindKey(const char *KeyStart, const char *KeyEnd) const;
71   
72 public:
73   static StringMapEntryBase *getTombstoneVal() {
74     return (StringMapEntryBase*)-1;
75   }
76   
77   unsigned getNumBuckets() const { return NumBuckets; }
78   unsigned getNumItems() const { return NumItems; }
79
80   bool empty() const { return NumItems == 0; }
81   unsigned size() const { return NumItems; }
82 };
83
84 /// StringMapEntry - This is used to represent one value that is inserted into
85 /// a StringMap.  It contains the Value itself and the key: the string length
86 /// and data.
87 template<typename ValueTy>
88 class StringMapEntry : public StringMapEntryBase {
89   ValueTy Val;
90 public:
91   StringMapEntry(unsigned StrLen)
92     : StringMapEntryBase(StrLen), Val() {}
93   StringMapEntry(unsigned StrLen, const ValueTy &V)
94     : StringMapEntryBase(StrLen), Val(V) {}
95
96   const ValueTy &getValue() const { return Val; }
97   ValueTy &getValue() { return Val; }
98   
99   void setValue(const ValueTy &V) { Val = V; }
100   
101   /// getKeyData - Return the start of the string data that is the key for this
102   /// value.  The string data is always stored immediately after the
103   /// StringMapEntry object.
104   const char *getKeyData() const {return reinterpret_cast<const char*>(this+1);}
105   
106   /// Create - Create a StringMapEntry for the specified key and default
107   /// construct the value.
108   template<typename AllocatorTy>
109   static StringMapEntry *Create(const char *KeyStart, const char *KeyEnd,
110                                 AllocatorTy &Allocator) {
111     unsigned KeyLength = KeyEnd-KeyStart;
112     
113     // Okay, the item doesn't already exist, and 'Bucket' is the bucket to fill
114     // in.  Allocate a new item with space for the string at the end and a null
115     // terminator.
116     unsigned AllocSize = sizeof(StringMapEntry)+KeyLength+1;
117     
118 #ifdef __GNUC__
119     unsigned Alignment = __alignof__(StringMapEntry);
120 #else
121     // FIXME: ugly.
122     unsigned Alignment = 8;
123 #endif
124     StringMapEntry *NewItem = 
125       static_cast<StringMapEntry*>(Allocator.Allocate(AllocSize, Alignment));
126     
127     // Default construct the value.
128     new (NewItem) StringMapEntry(KeyLength);
129     
130     // Copy the string information.
131     char *StrBuffer = const_cast<char*>(NewItem->getKeyData());
132     memcpy(StrBuffer, KeyStart, KeyLength);
133     StrBuffer[KeyLength] = 0;  // Null terminate for convenience of clients.
134     return NewItem;
135   }
136   
137   /// Create - Create a StringMapEntry with normal malloc/free.
138   static StringMapEntry *Create(const char *KeyStart, const char *KeyEnd) {
139     MallocAllocator A;
140     return Create(KeyStart, KeyEnd, A);
141   }
142
143   /// Destroy - Destroy this StringMapEntry, releasing memory back to the
144   /// specified allocator.
145   template<typename AllocatorTy>
146   void Destroy(AllocatorTy &Allocator) {
147     // Free memory referenced by the item.
148     this->~StringMapEntry();
149     Allocator.Deallocate(this);
150   }
151   
152   /// Destroy this object, releasing memory back to the malloc allocator.
153   void Destroy() {
154     MallocAllocator A;
155     Destroy(A);
156   }
157 };
158
159
160 /// StringMap - This is an unconventional map that is specialized for handling
161 /// keys that are "strings", which are basically ranges of bytes. This does some
162 /// funky memory allocation and hashing things to make it extremely efficient,
163 /// storing the string data *after* the value in the map.
164 template<typename ValueTy, typename AllocatorTy = MallocAllocator>
165 class StringMap : public StringMapImpl {
166   AllocatorTy Allocator;
167   typedef StringMapEntry<ValueTy> MapEntryTy;
168 public:
169   StringMap(unsigned InitialSize = 0)
170     : StringMapImpl(InitialSize, sizeof(MapEntryTy)) {}
171   
172   AllocatorTy &getAllocator() { return Allocator; }
173   const AllocatorTy &getAllocator() const { return Allocator; }
174
175   typedef StringMapConstIterator<ValueTy> const_iterator;
176   typedef StringMapIterator<ValueTy> iterator;
177   
178   iterator begin() { return iterator(TheTable); }
179   iterator end() { return iterator(TheTable+NumBuckets); }
180   const_iterator begin() const { return const_iterator(TheTable); }
181   const_iterator end() const { return const_iterator(TheTable+NumBuckets); }
182   
183   
184   iterator find(const char *KeyStart, const char *KeyEnd) {
185     int Bucket = FindKey(KeyStart, KeyEnd);
186     if (Bucket == -1) return end();
187     return iterator(TheTable+Bucket);
188   }
189
190   const_iterator find(const char *KeyStart, const char *KeyEnd) const {
191     int Bucket = FindKey(KeyStart, KeyEnd);
192     if (Bucket == -1) return end();
193     return const_iterator(TheTable+Bucket);
194   }
195   
196   /// GetOrCreateValue - Look up the specified key in the table.  If a value
197   /// exists, return it.  Otherwise, default construct a value, insert it, and
198   /// return.
199   StringMapEntry<ValueTy> &GetOrCreateValue(const char *KeyStart, 
200                                             const char *KeyEnd) {
201     unsigned BucketNo = LookupBucketFor(KeyStart, KeyEnd);
202     ItemBucket &Bucket = TheTable[BucketNo];
203     if (Bucket.Item)
204       return *static_cast<MapEntryTy*>(Bucket.Item);
205     
206     MapEntryTy *NewItem = MapEntryTy::Create(KeyStart, KeyEnd, Allocator);
207     ++NumItems;
208     
209     // Fill in the bucket for the hash table.  The FullHashValue was already
210     // filled in by LookupBucketFor.
211     Bucket.Item = NewItem;
212     
213     // If the hash table is now more than 3/4 full, rehash into a larger table.
214     if (NumItems > NumBuckets*3/4)
215       RehashTable();
216     return *NewItem;
217   }
218   
219   ~StringMap() {
220     for (ItemBucket *I = TheTable, *E = TheTable+NumBuckets; I != E; ++I) {
221       if (MapEntryTy *Id = static_cast<MapEntryTy*>(I->Item))
222         Id->Destroy(Allocator);
223     }
224     delete [] TheTable;
225   }
226 };
227   
228
229 template<typename ValueTy>
230 class StringMapConstIterator {
231   StringMapImpl::ItemBucket *Ptr;
232 public:
233   StringMapConstIterator(StringMapImpl::ItemBucket *Bucket) : Ptr(Bucket) {
234     AdvancePastEmptyBuckets();
235   }
236   
237   const StringMapEntry<ValueTy> &operator*() const {
238     return *static_cast<StringMapEntry<ValueTy>*>(Ptr->Item);
239   }
240   const StringMapEntry<ValueTy> *operator->() const {
241     return static_cast<StringMapEntry<ValueTy>*>(Ptr->Item);
242   }
243   
244   bool operator==(const StringMapConstIterator &RHS) const {
245     return Ptr == RHS.Ptr;
246   }
247   bool operator!=(const StringMapConstIterator &RHS) const {
248     return Ptr != RHS.Ptr;
249   }
250   
251   inline StringMapConstIterator& operator++() {          // Preincrement
252     ++Ptr;
253     AdvancePastEmptyBuckets();
254     return *this;
255   }
256   StringMapConstIterator operator++(int) {        // Postincrement
257     StringMapConstIterator tmp = *this; ++*this; return tmp;
258   }
259   
260 private:
261   void AdvancePastEmptyBuckets() {
262     while (Ptr->Item == 0 || Ptr->Item == StringMapImpl::getTombstoneVal())
263       ++Ptr;
264   }
265 };
266
267 template<typename ValueTy>
268 class StringMapIterator : public StringMapConstIterator<ValueTy> {
269 public:  
270   StringMapIterator(StringMapImpl::ItemBucket *Bucket)
271     : StringMapConstIterator<ValueTy>(Bucket) {
272   }
273   StringMapEntry<ValueTy> &operator*() const {
274     return *static_cast<StringMapEntry<ValueTy>*>(this->Ptr->Item);
275   }
276   StringMapEntry<ValueTy> *operator->() const {
277     return static_cast<StringMapEntry<ValueTy>*>(this->Ptr->Item);
278   }
279 };
280
281 }
282
283 #endif
284