fixup! Address review feedback for more comments

Author: cachemeifyoucan

llvm/lib/Support/TrieHashIndexGenerator.h

@@ -14,51 +14,82 @@
 
 namespace llvm {
 
+/// The utility class that helps computing the index of the object inside trie
+/// from its hash. The generator can be configured with the number of bits
+/// used for each level of trie structure with \c NumRootsBits and \c
+/// NumSubtrieBits.
+/// For example, try computing indexes for a 16-bit hash 0x1234 with 8-bit root
+/// and 4-bit sub-trie:
+///
+///   IndexGenerator IndexGen{8, 4, Hash};
+///   size_t index1 = IndexGen.next(); // index 18 in root node.
+///   size_t index2 = IndexGen.next(); // index 3 in sub-trie level 1.
+///   size_t index3 = IndexGen.next(); // index 4 in sub-tire level 2.
+///
+/// This is used by different trie implementation to figure out where to
+/// insert/find the object in the data structure.
 struct IndexGenerator {
   size_t NumRootBits;
   size_t NumSubtrieBits;
   ArrayRef<uint8_t> Bytes;
   std::optional<size_t> StartBit = std::nullopt;
 
+  // Get the number of bits used to generate current index.
   size_t getNumBits() const {
     assert(StartBit);
     size_t TotalNumBits = Bytes.size() * 8;
     assert(*StartBit <= TotalNumBits);
     return std::min(*StartBit ? NumSubtrieBits : NumRootBits,
                     TotalNumBits - *StartBit);
   }
+
+  // Get the index of the object in the next level of trie.
   size_t next() {
     size_t Index;
     if (!StartBit) {
+      // Compute index for root when StartBit is not set.
       StartBit = 0;
       Index = getIndex(Bytes, *StartBit, NumRootBits);
     } else {
+      // Compute index for sub-trie.
       *StartBit += *StartBit ? NumSubtrieBits : NumRootBits;
       assert((*StartBit - NumRootBits) % NumSubtrieBits == 0);
       Index = getIndex(Bytes, *StartBit, NumSubtrieBits);
     }
     return Index;
   }
 
+  // Provide a hint to speed up the index generation by providing the
+  // information of the hash in current level. For example, if the object is
+  // known to have \c Index on a level that already consumes first n \c Bits of
+  // the hash, it can start index generation from this level by calling \c hint
+  // function.
   size_t hint(unsigned Index, unsigned Bit) {
     assert(Bit < Bytes.size() * 8);
     assert(Bit == 0 || (Bit - NumRootBits) % NumSubtrieBits == 0);
     StartBit = Bit;
     return Index;
   }
 
+  // Utility funciton for looking up the index in the trie for an object that
+  // has colliding hash bits in the front as the hash of the object that is
+  // currently being computed.
   size_t getCollidingBits(ArrayRef<uint8_t> CollidingBits) const {
     assert(StartBit);
     return getIndex(CollidingBits, *StartBit, NumSubtrieBits);
   }
 
+  // Compute the index for the object from its hash, current start bits, and
+  // the number of bits used for current level.
   static size_t getIndex(ArrayRef<uint8_t> Bytes, size_t StartBit,
                          size_t NumBits) {
     assert(StartBit < Bytes.size() * 8);
-
+    // Drop all the bits before StartBit.
     Bytes = Bytes.drop_front(StartBit / 8u);
     StartBit %= 8u;
     size_t Index = 0;
+    // Compute the index using the bits in range [StartBit, StartBit + NumBits),
+    // note the range can spread across few `uint8_t` in the array.
     for (uint8_t Byte : Bytes) {
       size_t ByteStart = 0, ByteEnd = 8;
       if (StartBit) {

llvm/lib/Support/TrieRawHashMap.cpp

@@ -102,14 +102,15 @@ class TrieSubtrie final : public TrieNode {
 };
 } // end namespace
 
-static size_t getTrieTailSize(size_t StartBit, size_t NumBits) {
-  assert(NumBits < 20 && "Tries should have fewer than ~1M slots");
+// Compute the trailing object size in the trie node. This is the size of \c
+// Slots in TrieNodes that pointing to the children.
+static size_t getTrieTailSize(size_t NumBits) {
   return sizeof(TrieNode *) * (1u << NumBits);
 }
 
 std::unique_ptr<TrieSubtrie> TrieSubtrie::create(size_t StartBit,
                                                  size_t NumBits) {
-  size_t Size = sizeof(TrieSubtrie) + getTrieTailSize(StartBit, NumBits);
+  size_t Size = sizeof(TrieSubtrie) + getTrieTailSize(NumBits);
   void *Memory = ::malloc(Size);
   TrieSubtrie *S = ::new (Memory) TrieSubtrie(StartBit, NumBits);
   return std::unique_ptr<TrieSubtrie>(S);
@@ -128,15 +129,22 @@ TrieSubtrie::TrieSubtrie(size_t StartBit, size_t NumBits)
       "Expected no work in destructor for TrieNode");
 }
 
+// Sink the nodes down sub-trie when the object being inserted collides with
+// the index of existing object in the trie. In this case, a new sub-trie needs
+// to be allocated to hold existing object.
 TrieSubtrie *TrieSubtrie::sink(
     size_t I, TrieContent &Content, size_t NumSubtrieBits, size_t NewI,
     function_ref<TrieSubtrie *(std::unique_ptr<TrieSubtrie>)> Saver) {
+  // Create a new sub-trie that points to the existing object with the new
+  // index for the next level.
   assert(NumSubtrieBits > 0);
   std::unique_ptr<TrieSubtrie> S = create(StartBit + NumBits, NumSubtrieBits);
 
   assert(NewI < S->Slots.size());
   S->Slots[NewI].store(&Content);
 
+  // Using compare_exchange to atomically add back the new sub-trie to the trie
+  // in the place of the exsiting object.
   TrieNode *ExistingNode = &Content;
   assert(I < Slots.size());
   if (Slots[I].compare_exchange_strong(ExistingNode, S.get()))
@@ -149,12 +157,15 @@ TrieSubtrie *TrieSubtrie::sink(
 
 struct ThreadSafeTrieRawHashMapBase::ImplType {
   static std::unique_ptr<ImplType> create(size_t StartBit, size_t NumBits) {
-    size_t Size = sizeof(ImplType) + getTrieTailSize(StartBit, NumBits);
+    size_t Size = sizeof(ImplType) + getTrieTailSize(NumBits);
     void *Memory = ::malloc(Size);
     ImplType *Impl = ::new (Memory) ImplType(StartBit, NumBits);
     return std::unique_ptr<ImplType>(Impl);
   }
 
+  // Save the Subtrie into the ownship list of the trie structure in a
+  // thread-safe way. The ownership transfer is done by compare_exchange the
+  // pointer value inside the unique_ptr.
   TrieSubtrie *save(std::unique_ptr<TrieSubtrie> S) {
     assert(!S->Next && "Expected S to a freshly-constructed leaf");
 
@@ -306,12 +317,7 @@ ThreadSafeTrieRawHashMapBase::ThreadSafeTrieRawHashMapBase(
       ContentOffset(ContentOffset),
       NumRootBits(NumRootBits ? *NumRootBits : DefaultNumRootBits),
       NumSubtrieBits(NumSubtrieBits ? *NumSubtrieBits : DefaultNumSubtrieBits),
-      ImplPtr(nullptr) {
-  assert((!NumRootBits || *NumRootBits < 20) &&
-         "Root should have fewer than ~1M slots");
-  assert((!NumSubtrieBits || *NumSubtrieBits < 10) &&
-         "Subtries should have fewer than ~1K slots");
-}
+      ImplPtr(nullptr) {}
 
 ThreadSafeTrieRawHashMapBase::ThreadSafeTrieRawHashMapBase(
     ThreadSafeTrieRawHashMapBase &&RHS)
@@ -413,7 +419,7 @@ std::string ThreadSafeTrieRawHashMapBase::getTriePrefixAsString(
   TrieContent *Node = nullptr;
   while (Current) {
     TrieSubtrie *Next = nullptr;
-    // find first used slot in the trie.
+    // Find first used slot in the trie.
     for (unsigned I = 0, E = Current->Slots.size(); I < E; ++I) {
       auto *S = Current->get(I);
       if (!S)