include/llvm/ADT/StringMap.h

   1 //===- StringMap.h - String Hash table map interface ------------*- C++ -*-===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file defines the StringMap class.
  10 //
  11 //===----------------------------------------------------------------------===//
  12
  13 #ifndef LLVM_ADT_STRINGMAP_H
  14 #define LLVM_ADT_STRINGMAP_H
  15
  16 #include "llvm/ADT/StringRef.h"
  17 #include "llvm/ADT/iterator.h"
  18 #include "llvm/ADT/iterator_range.h"
  19 #include "llvm/Support/Allocator.h"
  20 #include "llvm/Support/PointerLikeTypeTraits.h"
  21 #include "llvm/Support/ErrorHandling.h"
  22 #include <algorithm>
  23 #include <cassert>
  24 #include <cstdint>
  25 #include <cstdlib>
  26 #include <cstring>
  27 #include <initializer_list>
  28 #include <iterator>
  29 #include <utility>
  30
  31 namespace llvm {
  32
  33 template<typename ValueTy> class StringMapConstIterator;
  34 template<typename ValueTy> class StringMapIterator;
  35 template<typename ValueTy> class StringMapKeyIterator;
  36
  37 /// StringMapEntryBase - Shared base class of StringMapEntry instances.
  38 class StringMapEntryBase {
  39   size_t StrLen;
  40
  41 public:
  42   explicit StringMapEntryBase(size_t Len) : StrLen(Len) {}
  43
  44   size_t 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 protected:
  51   // Array of NumBuckets pointers to entries, null pointers are holes.
  52   // TheTable[NumBuckets] contains a sentinel value for easy iteration. Followed
  53   // by an array of the actual hash values as unsigned integers.
  54   StringMapEntryBase **TheTable = nullptr;
  55   unsigned NumBuckets = 0;
  56   unsigned NumItems = 0;
  57   unsigned NumTombstones = 0;
  58   unsigned ItemSize;
  59
  60 protected:
  61   explicit StringMapImpl(unsigned itemSize)
  62       : ItemSize(itemSize) {}
  63   StringMapImpl(StringMapImpl &&RHS)
  64       : TheTable(RHS.TheTable), NumBuckets(RHS.NumBuckets),
  65         NumItems(RHS.NumItems), NumTombstones(RHS.NumTombstones),
  66         ItemSize(RHS.ItemSize) {
  67     RHS.TheTable = nullptr;
  68     RHS.NumBuckets = 0;
  69     RHS.NumItems = 0;
  70     RHS.NumTombstones = 0;
  71   }
  72
  73   StringMapImpl(unsigned InitSize, unsigned ItemSize);
  74   unsigned RehashTable(unsigned BucketNo = 0);
  75
  76   /// LookupBucketFor - Look up the bucket that the specified string should end
  77   /// up in.  If it already exists as a key in the map, the Item pointer for the
  78   /// specified bucket will be non-null.  Otherwise, it will be null.  In either
  79   /// case, the FullHashValue field of the bucket will be set to the hash value
  80   /// of the string.
  81   unsigned LookupBucketFor(StringRef Key);
  82
  83   /// FindKey - Look up the bucket that contains the specified key. If it exists
  84   /// in the map, return the bucket number of the key.  Otherwise return -1.
  85   /// This does not modify the map.
  86   int FindKey(StringRef Key) const;
  87
  88   /// RemoveKey - Remove the specified StringMapEntry from the table, but do not
  89   /// delete it.  This aborts if the value isn't in the table.
  90   void RemoveKey(StringMapEntryBase *V);
  91
  92   /// RemoveKey - Remove the StringMapEntry for the specified key from the
  93   /// table, returning it.  If the key is not in the table, this returns null.
  94   StringMapEntryBase *RemoveKey(StringRef Key);
  95
  96   /// Allocate the table with the specified number of buckets and otherwise
  97   /// setup the map as empty.
  98   void init(unsigned Size);
  99
 100 public:
 101   static StringMapEntryBase *getTombstoneVal() {
 102     uintptr_t Val = static_cast<uintptr_t>(-1);
 103     Val <<= PointerLikeTypeTraits<StringMapEntryBase *>::NumLowBitsAvailable;
 104     return reinterpret_cast<StringMapEntryBase *>(Val);
 105   }
 106
 107   unsigned getNumBuckets() const { return NumBuckets; }
 108   unsigned getNumItems() const { return NumItems; }
 109
 110   bool empty() const { return NumItems == 0; }
 111   unsigned size() const { return NumItems; }
 112
 113   void swap(StringMapImpl &Other) {
 114     std::swap(TheTable, Other.TheTable);
 115     std::swap(NumBuckets, Other.NumBuckets);
 116     std::swap(NumItems, Other.NumItems);
 117     std::swap(NumTombstones, Other.NumTombstones);
 118   }
 119 };
 120
 121 /// StringMapEntryStorage - Holds the value in a StringMapEntry.
 122 ///
 123 /// Factored out into a separate base class to make it easier to specialize.
 124 /// This is primarily intended to support StringSet, which doesn't need a value
 125 /// stored at all.
 126 template<typename ValueTy>
 127 class StringMapEntryStorage : public StringMapEntryBase {
 128 public:
 129   ValueTy second;
 130
 131   explicit StringMapEntryStorage(size_t strLen)
 132     : StringMapEntryBase(strLen), second() {}
 133   template <typename... InitTy>
 134   StringMapEntryStorage(size_t strLen, InitTy &&... InitVals)
 135       : StringMapEntryBase(strLen), second(std::forward<InitTy>(InitVals)...) {}
 136   StringMapEntryStorage(StringMapEntryStorage &E) = delete;
 137
 138   const ValueTy &getValue() const { return second; }
 139   ValueTy &getValue() { return second; }
 140
 141   void setValue(const ValueTy &V) { second = V; }
 142 };
 143
 144 template<>
 145 class StringMapEntryStorage<NoneType> : public StringMapEntryBase {
 146 public:
 147   explicit StringMapEntryStorage(size_t strLen, NoneType none = None)
 148     : StringMapEntryBase(strLen) {}
 149   StringMapEntryStorage(StringMapEntryStorage &E) = delete;
 150
 151   NoneType getValue() const { return None; }
 152 };
 153
 154 /// StringMapEntry - This is used to represent one value that is inserted into
 155 /// a StringMap.  It contains the Value itself and the key: the string length
 156 /// and data.
 157 template<typename ValueTy>
 158 class StringMapEntry final : public StringMapEntryStorage<ValueTy> {
 159 public:
 160   using StringMapEntryStorage<ValueTy>::StringMapEntryStorage;
 161
 162   StringRef getKey() const {
 163     return StringRef(getKeyData(), this->getKeyLength());
 164   }
 165
 166   /// getKeyData - Return the start of the string data that is the key for this
 167   /// value.  The string data is always stored immediately after the
 168   /// StringMapEntry object.
 169   const char *getKeyData() const {return reinterpret_cast<const char*>(this+1);}
 170
 171   StringRef first() const {
 172     return StringRef(getKeyData(), this->getKeyLength());
 173   }
 174
 175   /// Create a StringMapEntry for the specified key construct the value using
 176   /// \p InitiVals.
 177   template <typename AllocatorTy, typename... InitTy>
 178   static StringMapEntry *Create(StringRef Key, AllocatorTy &Allocator,
 179                                 InitTy &&... InitVals) {
 180     size_t KeyLength = Key.size();
 181
 182     // Allocate a new item with space for the string at the end and a null
 183     // terminator.
 184     size_t AllocSize = sizeof(StringMapEntry) + KeyLength + 1;
 185     size_t Alignment = alignof(StringMapEntry);
 186
 187     StringMapEntry *NewItem =
 188       static_cast<StringMapEntry*>(Allocator.Allocate(AllocSize,Alignment));
 189     assert(NewItem && "Unhandled out-of-memory");
 190
 191     // Construct the value.
 192     new (NewItem) StringMapEntry(KeyLength, std::forward<InitTy>(InitVals)...);
 193
 194     // Copy the string information.
 195     char *StrBuffer = const_cast<char*>(NewItem->getKeyData());
 196     if (KeyLength > 0)
 197       memcpy(StrBuffer, Key.data(), KeyLength);
 198     StrBuffer[KeyLength] = 0;  // Null terminate for convenience of clients.
 199     return NewItem;
 200   }
 201
 202   /// Create - Create a StringMapEntry with normal malloc/free.
 203   template <typename... InitType>
 204   static StringMapEntry *Create(StringRef Key, InitType &&... InitVal) {
 205     MallocAllocator A;
 206     return Create(Key, A, std::forward<InitType>(InitVal)...);
 207   }
 208
 209   static StringMapEntry *Create(StringRef Key) {
 210     return Create(Key, ValueTy());
 211   }
 212
 213   /// GetStringMapEntryFromKeyData - Given key data that is known to be embedded
 214   /// into a StringMapEntry, return the StringMapEntry itself.
 215   static StringMapEntry &GetStringMapEntryFromKeyData(const char *KeyData) {
 216     char *Ptr = const_cast<char*>(KeyData) - sizeof(StringMapEntry<ValueTy>);
 217     return *reinterpret_cast<StringMapEntry*>(Ptr);
 218   }
 219
 220   /// Destroy - Destroy this StringMapEntry, releasing memory back to the
 221   /// specified allocator.
 222   template<typename AllocatorTy>
 223   void Destroy(AllocatorTy &Allocator) {
 224     // Free memory referenced by the item.
 225     size_t AllocSize = sizeof(StringMapEntry) + this->getKeyLength() + 1;
 226     this->~StringMapEntry();
 227     Allocator.Deallocate(static_cast<void *>(this), AllocSize);
 228   }
 229
 230   /// Destroy this object, releasing memory back to the malloc allocator.
 231   void Destroy() {
 232     MallocAllocator A;
 233     Destroy(A);
 234   }
 235 };
 236
 237 /// StringMap - This is an unconventional map that is specialized for handling
 238 /// keys that are "strings", which are basically ranges of bytes. This does some
 239 /// funky memory allocation and hashing things to make it extremely efficient,
 240 /// storing the string data *after* the value in the map.
 241 template<typename ValueTy, typename AllocatorTy = MallocAllocator>
 242 class StringMap : public StringMapImpl {
 243   AllocatorTy Allocator;
 244
 245 public:
 246   using MapEntryTy = StringMapEntry<ValueTy>;
 247
 248   StringMap() : StringMapImpl(static_cast<unsigned>(sizeof(MapEntryTy))) {}
 249
 250   explicit StringMap(unsigned InitialSize)
 251     : StringMapImpl(InitialSize, static_cast<unsigned>(sizeof(MapEntryTy))) {}
 252
 253   explicit StringMap(AllocatorTy A)
 254     : StringMapImpl(static_cast<unsigned>(sizeof(MapEntryTy))), Allocator(A) {}
 255
 256   StringMap(unsigned InitialSize, AllocatorTy A)
 257     : StringMapImpl(InitialSize, static_cast<unsigned>(sizeof(MapEntryTy))),
 258       Allocator(A) {}
 259
 260   StringMap(std::initializer_list<std::pair<StringRef, ValueTy>> List)
 261       : StringMapImpl(List.size(), static_cast<unsigned>(sizeof(MapEntryTy))) {
 262     for (const auto &P : List) {
 263       insert(P);
 264     }
 265   }
 266
 267   StringMap(StringMap &&RHS)
 268       : StringMapImpl(std::move(RHS)), Allocator(std::move(RHS.Allocator)) {}
 269
 270   StringMap(const StringMap &RHS) :
 271     StringMapImpl(static_cast<unsigned>(sizeof(MapEntryTy))),
 272     Allocator(RHS.Allocator) {
 273     if (RHS.empty())
 274       return;
 275
 276     // Allocate TheTable of the same size as RHS's TheTable, and set the
 277     // sentinel appropriately (and NumBuckets).
 278     init(RHS.NumBuckets);
 279     unsigned *HashTable = (unsigned *)(TheTable + NumBuckets + 1),
 280              *RHSHashTable = (unsigned *)(RHS.TheTable + NumBuckets + 1);
 281
 282     NumItems = RHS.NumItems;
 283     NumTombstones = RHS.NumTombstones;
 284     for (unsigned I = 0, E = NumBuckets; I != E; ++I) {
 285       StringMapEntryBase *Bucket = RHS.TheTable[I];
 286       if (!Bucket || Bucket == getTombstoneVal()) {
 287         TheTable[I] = Bucket;
 288         continue;
 289       }
 290
 291       TheTable[I] = MapEntryTy::Create(
 292           static_cast<MapEntryTy *>(Bucket)->getKey(), Allocator,
 293           static_cast<MapEntryTy *>(Bucket)->getValue());
 294       HashTable[I] = RHSHashTable[I];
 295     }
 296
 297     // Note that here we've copied everything from the RHS into this object,
 298     // tombstones included. We could, instead, have re-probed for each key to
 299     // instantiate this new object without any tombstone buckets. The
 300     // assumption here is that items are rarely deleted from most StringMaps,
 301     // and so tombstones are rare, so the cost of re-probing for all inputs is
 302     // not worthwhile.
 303   }
 304
 305   StringMap &operator=(StringMap RHS) {
 306     StringMapImpl::swap(RHS);
 307     std::swap(Allocator, RHS.Allocator);
 308     return *this;
 309   }
 310
 311   ~StringMap() {
 312     // Delete all the elements in the map, but don't reset the elements
 313     // to default values.  This is a copy of clear(), but avoids unnecessary
 314     // work not required in the destructor.
 315     if (!empty()) {
 316       for (unsigned I = 0, E = NumBuckets; I != E; ++I) {
 317         StringMapEntryBase *Bucket = TheTable[I];
 318         if (Bucket && Bucket != getTombstoneVal()) {
 319           static_cast<MapEntryTy*>(Bucket)->Destroy(Allocator);
 320         }
 321       }
 322     }
 323     free(TheTable);
 324   }
 325
 326   AllocatorTy &getAllocator() { return Allocator; }
 327   const AllocatorTy &getAllocator() const { return Allocator; }
 328
 329   using key_type = const char*;
 330   using mapped_type = ValueTy;
 331   using value_type = StringMapEntry<ValueTy>;
 332   using size_type = size_t;
 333
 334   using const_iterator = StringMapConstIterator<ValueTy>;
 335   using iterator = StringMapIterator<ValueTy>;
 336
 337   iterator begin() {
 338     return iterator(TheTable, NumBuckets == 0);
 339   }
 340   iterator end() {
 341     return iterator(TheTable+NumBuckets, true);
 342   }
 343   const_iterator begin() const {
 344     return const_iterator(TheTable, NumBuckets == 0);
 345   }
 346   const_iterator end() const {
 347     return const_iterator(TheTable+NumBuckets, true);
 348   }
 349
 350   iterator_range<StringMapKeyIterator<ValueTy>> keys() const {
 351     return make_range(StringMapKeyIterator<ValueTy>(begin()),
 352                       StringMapKeyIterator<ValueTy>(end()));
 353   }
 354
 355   iterator find(StringRef Key) {
 356     int Bucket = FindKey(Key);
 357     if (Bucket == -1) return end();
 358     return iterator(TheTable+Bucket, true);
 359   }
 360
 361   const_iterator find(StringRef Key) const {
 362     int Bucket = FindKey(Key);
 363     if (Bucket == -1) return end();
 364     return const_iterator(TheTable+Bucket, true);
 365   }
 366
 367   /// lookup - Return the entry for the specified key, or a default
 368   /// constructed value if no such entry exists.
 369   ValueTy lookup(StringRef Key) const {
 370     const_iterator it = find(Key);
 371     if (it != end())
 372       return it->second;
 373     return ValueTy();
 374   }
 375
 376   /// Lookup the ValueTy for the \p Key, or create a default constructed value
 377   /// if the key is not in the map.
 378   ValueTy &operator[](StringRef Key) { return try_emplace(Key).first->second; }
 379
 380   /// count - Return 1 if the element is in the map, 0 otherwise.
 381   size_type count(StringRef Key) const {
 382     return find(Key) == end() ? 0 : 1;
 383   }
 384
 385   template <typename InputTy>
 386   size_type count(const StringMapEntry<InputTy> &MapEntry) const {
 387     return count(MapEntry.getKey());
 388   }
 389
 390   /// insert - Insert the specified key/value pair into the map.  If the key
 391   /// already exists in the map, return false and ignore the request, otherwise
 392   /// insert it and return true.
 393   bool insert(MapEntryTy *KeyValue) {
 394     unsigned BucketNo = LookupBucketFor(KeyValue->getKey());
 395     StringMapEntryBase *&Bucket = TheTable[BucketNo];
 396     if (Bucket && Bucket != getTombstoneVal())
 397       return false;  // Already exists in map.
 398
 399     if (Bucket == getTombstoneVal())
 400       --NumTombstones;
 401     Bucket = KeyValue;
 402     ++NumItems;
 403     assert(NumItems + NumTombstones <= NumBuckets);
 404
 405     RehashTable();
 406     return true;
 407   }
 408
 409   /// insert - Inserts the specified key/value pair into the map if the key
 410   /// isn't already in the map. The bool component of the returned pair is true
 411   /// if and only if the insertion takes place, and the iterator component of
 412   /// the pair points to the element with key equivalent to the key of the pair.
 413   std::pair<iterator, bool> insert(std::pair<StringRef, ValueTy> KV) {
 414     return try_emplace(KV.first, std::move(KV.second));
 415   }
 416
 417   /// Inserts an element or assigns to the current element if the key already
 418   /// exists. The return type is the same as try_emplace.
 419   template <typename V>
 420   std::pair<iterator, bool> insert_or_assign(StringRef Key, V &&Val) {
 421     auto Ret = try_emplace(Key, std::forward<V>(Val));
 422     if (!Ret.second)
 423       Ret.first->second = std::forward<V>(Val);
 424     return Ret;
 425   }
 426
 427   /// Emplace a new element for the specified key into the map if the key isn't
 428   /// already in the map. The bool component of the returned pair is true
 429   /// if and only if the insertion takes place, and the iterator component of
 430   /// the pair points to the element with key equivalent to the key of the pair.
 431   template <typename... ArgsTy>
 432   std::pair<iterator, bool> try_emplace(StringRef Key, ArgsTy &&... Args) {
 433     unsigned BucketNo = LookupBucketFor(Key);
 434     StringMapEntryBase *&Bucket = TheTable[BucketNo];
 435     if (Bucket && Bucket != getTombstoneVal())
 436       return std::make_pair(iterator(TheTable + BucketNo, false),
 437                             false); // Already exists in map.
 438
 439     if (Bucket == getTombstoneVal())
 440       --NumTombstones;
 441     Bucket = MapEntryTy::Create(Key, Allocator, std::forward<ArgsTy>(Args)...);
 442     ++NumItems;
 443     assert(NumItems + NumTombstones <= NumBuckets);
 444
 445     BucketNo = RehashTable(BucketNo);
 446     return std::make_pair(iterator(TheTable + BucketNo, false), true);
 447   }
 448
 449   // clear - Empties out the StringMap
 450   void clear() {
 451     if (empty()) return;
 452
 453     // Zap all values, resetting the keys back to non-present (not tombstone),
 454     // which is safe because we're removing all elements.
 455     for (unsigned I = 0, E = NumBuckets; I != E; ++I) {
 456       StringMapEntryBase *&Bucket = TheTable[I];
 457       if (Bucket && Bucket != getTombstoneVal()) {
 458         static_cast<MapEntryTy*>(Bucket)->Destroy(Allocator);
 459       }
 460       Bucket = nullptr;
 461     }
 462
 463     NumItems = 0;
 464     NumTombstones = 0;
 465   }
 466
 467   /// remove - Remove the specified key/value pair from the map, but do not
 468   /// erase it.  This aborts if the key is not in the map.
 469   void remove(MapEntryTy *KeyValue) {
 470     RemoveKey(KeyValue);
 471   }
 472
 473   void erase(iterator I) {
 474     MapEntryTy &V = *I;
 475     remove(&V);
 476     V.Destroy(Allocator);
 477   }
 478
 479   bool erase(StringRef Key) {
 480     iterator I = find(Key);
 481     if (I == end()) return false;
 482     erase(I);
 483     return true;
 484   }
 485 };
 486
 487 template <typename DerivedTy, typename ValueTy>
 488 class StringMapIterBase
 489     : public iterator_facade_base<DerivedTy, std::forward_iterator_tag,
 490                                   ValueTy> {
 491 protected:
 492   StringMapEntryBase **Ptr = nullptr;
 493
 494 public:
 495   StringMapIterBase() = default;
 496
 497   explicit StringMapIterBase(StringMapEntryBase **Bucket,
 498                              bool NoAdvance = false)
 499       : Ptr(Bucket) {
 500     if (!NoAdvance) AdvancePastEmptyBuckets();
 501   }
 502
 503   DerivedTy &operator=(const DerivedTy &Other) {
 504     Ptr = Other.Ptr;
 505     return static_cast<DerivedTy &>(*this);
 506   }
 507
 508   bool operator==(const DerivedTy &RHS) const { return Ptr == RHS.Ptr; }
 509
 510   DerivedTy &operator++() { // Preincrement
 511     ++Ptr;
 512     AdvancePastEmptyBuckets();
 513     return static_cast<DerivedTy &>(*this);
 514   }
 515
 516   DerivedTy operator++(int) { // Post-increment
 517     DerivedTy Tmp(Ptr);
 518     ++*this;
 519     return Tmp;
 520   }
 521
 522 private:
 523   void AdvancePastEmptyBuckets() {
 524     while (*Ptr == nullptr || *Ptr == StringMapImpl::getTombstoneVal())
 525       ++Ptr;
 526   }
 527 };
 528
 529 template <typename ValueTy>
 530 class StringMapConstIterator
 531     : public StringMapIterBase<StringMapConstIterator<ValueTy>,
 532                                const StringMapEntry<ValueTy>> {
 533   using base = StringMapIterBase<StringMapConstIterator<ValueTy>,
 534                                  const StringMapEntry<ValueTy>>;
 535
 536 public:
 537   StringMapConstIterator() = default;
 538   explicit StringMapConstIterator(StringMapEntryBase **Bucket,
 539                                   bool NoAdvance = false)
 540       : base(Bucket, NoAdvance) {}
 541
 542   const StringMapEntry<ValueTy> &operator*() const {
 543     return *static_cast<const StringMapEntry<ValueTy> *>(*this->Ptr);
 544   }
 545 };
 546
 547 template <typename ValueTy>
 548 class StringMapIterator : public StringMapIterBase<StringMapIterator<ValueTy>,
 549                                                    StringMapEntry<ValueTy>> {
 550   using base =
 551       StringMapIterBase<StringMapIterator<ValueTy>, StringMapEntry<ValueTy>>;
 552
 553 public:
 554   StringMapIterator() = default;
 555   explicit StringMapIterator(StringMapEntryBase **Bucket,
 556                              bool NoAdvance = false)
 557       : base(Bucket, NoAdvance) {}
 558
 559   StringMapEntry<ValueTy> &operator*() const {
 560     return *static_cast<StringMapEntry<ValueTy> *>(*this->Ptr);
 561   }
 562
 563   operator StringMapConstIterator<ValueTy>() const {
 564     return StringMapConstIterator<ValueTy>(this->Ptr, true);
 565   }
 566 };
 567
 568 template <typename ValueTy>
 569 class StringMapKeyIterator
 570     : public iterator_adaptor_base<StringMapKeyIterator<ValueTy>,
 571                                    StringMapConstIterator<ValueTy>,
 572                                    std::forward_iterator_tag, StringRef> {
 573   using base = iterator_adaptor_base<StringMapKeyIterator<ValueTy>,
 574                                      StringMapConstIterator<ValueTy>,
 575                                      std::forward_iterator_tag, StringRef>;
 576
 577 public:
 578   StringMapKeyIterator() = default;
 579   explicit StringMapKeyIterator(StringMapConstIterator<ValueTy> Iter)
 580       : base(std::move(Iter)) {}
 581
 582   StringRef &operator*() {
 583     Key = this->wrapped()->getKey();
 584     return Key;
 585   }
 586
 587 private:
 588   StringRef Key;
 589 };
 590
 591 } // end namespace llvm
 592
 593 #endif // LLVM_ADT_STRINGMAP_H