1//===- TrieRawHashMap.cpp -------------------------------------------------===//
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#include "llvm/ADT/TrieRawHashMap.h"
10#include "llvm/ADT/LazyAtomicPointer.h"
11#include "llvm/ADT/StringExtras.h"
12#include "llvm/ADT/TrieHashIndexGenerator.h"
13#include "llvm/Support/Allocator.h"
14#include "llvm/Support/Casting.h"
15#include "llvm/Support/Debug.h"
16#include "llvm/Support/ThreadSafeAllocator.h"
17#include "llvm/Support/TrailingObjects.h"
18#include "llvm/Support/raw_ostream.h"
19#include <memory>
20
21using namespace llvm;
22
23namespace {
24struct TrieNode {
25 const bool IsSubtrie = false;
26
27 TrieNode(bool IsSubtrie) : IsSubtrie(IsSubtrie) {}
28
29 static void *operator new(size_t Size) { return ::operator new(Size); }
30 void operator delete(void *Ptr) { ::operator delete(Ptr); }
31};
32
33struct TrieContent final : public TrieNode {
34 const uint8_t ContentOffset;
35 const uint8_t HashSize;
36 const uint8_t HashOffset;
37
38 void *getValuePointer() const {
39 auto *Content = reinterpret_cast<const uint8_t *>(this) + ContentOffset;
40 return const_cast<uint8_t *>(Content);
41 }
42
43 ArrayRef<uint8_t> getHash() const {
44 auto *Begin = reinterpret_cast<const uint8_t *>(this) + HashOffset;
45 return ArrayRef(Begin, Begin + HashSize);
46 }
47
48 TrieContent(size_t ContentOffset, size_t HashSize, size_t HashOffset)
49 : TrieNode(/*IsSubtrie=*/false), ContentOffset(ContentOffset),
50 HashSize(HashSize), HashOffset(HashOffset) {}
51
52 static bool classof(const TrieNode *TN) { return !TN->IsSubtrie; }
53};
54
55static_assert(sizeof(TrieContent) ==
56 ThreadSafeTrieRawHashMapBase::TrieContentBaseSize,
57 "Check header assumption!");
58
59class TrieSubtrie final
60 : public TrieNode,
61 private TrailingObjects<TrieSubtrie, LazyAtomicPointer<TrieNode>> {
62public:
63 using Slot = LazyAtomicPointer<TrieNode>;
64
65 Slot &get(size_t I) { return getTrailingObjects()[I]; }
66 TrieNode *load(size_t I) { return get(I).load(); }
67
68 unsigned size() const { return Size; }
69
70 TrieSubtrie *
71 sink(size_t I, TrieContent &Content, size_t NumSubtrieBits, size_t NewI,
72 function_ref<TrieSubtrie *(std::unique_ptr<TrieSubtrie>)> Saver);
73
74 static std::unique_ptr<TrieSubtrie> create(size_t StartBit, size_t NumBits);
75
76 explicit TrieSubtrie(size_t StartBit, size_t NumBits);
77
78 static bool classof(const TrieNode *TN) { return TN->IsSubtrie; }
79
80 static constexpr size_t sizeToAlloc(unsigned NumBits) {
81 assert(NumBits < 20 && "Tries should have fewer than ~1M slots");
82 unsigned Count = 1u << NumBits;
83 return totalSizeToAlloc<LazyAtomicPointer<TrieNode>>(Counts: Count);
84 }
85
86private:
87 // FIXME: Use a bitset to speed up access:
88 //
89 // std::array<std::atomic<uint64_t>, NumSlots/64> IsSet;
90 //
91 // This will avoid needing to visit sparsely filled slots in
92 // \a ThreadSafeTrieRawHashMapBase::destroyImpl() when there's a non-trivial
93 // destructor.
94 //
95 // It would also greatly speed up iteration, if we add that some day, and
96 // allow get() to return one level sooner.
97 //
98 // This would be the algorithm for updating IsSet (after updating Slots):
99 //
100 // std::atomic<uint64_t> &Bits = IsSet[I.High];
101 // const uint64_t NewBit = 1ULL << I.Low;
102 // uint64_t Old = 0;
103 // while (!Bits.compare_exchange_weak(Old, Old | NewBit))
104 // ;
105
106 // For debugging.
107 unsigned StartBit = 0;
108 unsigned NumBits = 0;
109 unsigned Size = 0;
110 friend class llvm::ThreadSafeTrieRawHashMapBase;
111 friend class TrailingObjects;
112
113public:
114 /// Linked list for ownership of tries. The pointer is owned by TrieSubtrie.
115 std::atomic<TrieSubtrie *> Next;
116};
117} // end namespace
118
119std::unique_ptr<TrieSubtrie> TrieSubtrie::create(size_t StartBit,
120 size_t NumBits) {
121 void *Memory = ::operator new(sizeToAlloc(NumBits));
122 TrieSubtrie *S = ::new (Memory) TrieSubtrie(StartBit, NumBits);
123 return std::unique_ptr<TrieSubtrie>(S);
124}
125
126TrieSubtrie::TrieSubtrie(size_t StartBit, size_t NumBits)
127 : TrieNode(true), StartBit(StartBit), NumBits(NumBits), Size(1u << NumBits),
128 Next(nullptr) {
129 for (unsigned I = 0; I < Size; ++I)
130 new (&get(I)) Slot(nullptr);
131
132 static_assert(
133 std::is_trivially_destructible<LazyAtomicPointer<TrieNode>>::value,
134 "Expected no work in destructor for TrieNode");
135}
136
137// Sink the nodes down sub-trie when the object being inserted collides with
138// the index of existing object in the trie. In this case, a new sub-trie needs
139// to be allocated to hold existing object.
140TrieSubtrie *TrieSubtrie::sink(
141 size_t I, TrieContent &Content, size_t NumSubtrieBits, size_t NewI,
142 function_ref<TrieSubtrie *(std::unique_ptr<TrieSubtrie>)> Saver) {
143 // Create a new sub-trie that points to the existing object with the new
144 // index for the next level.
145 assert(NumSubtrieBits > 0);
146 std::unique_ptr<TrieSubtrie> S = create(StartBit: StartBit + NumBits, NumBits: NumSubtrieBits);
147
148 assert(NewI < Size);
149 S->get(I: NewI).store(Value: &Content);
150
151 // Using compare_exchange to atomically add back the new sub-trie to the trie
152 // in the place of the exsiting object.
153 TrieNode *ExistingNode = &Content;
154 assert(I < Size);
155 if (get(I).compare_exchange_strong(ExistingValue&: ExistingNode, NewValue: S.get()))
156 return Saver(std::move(S));
157
158 // Another thread created a subtrie already. Return it and let "S" be
159 // destructed.
160 return cast<TrieSubtrie>(Val: ExistingNode);
161}
162
163class ThreadSafeTrieRawHashMapBase::ImplType final
164 : private TrailingObjects<ThreadSafeTrieRawHashMapBase::ImplType,
165 TrieSubtrie> {
166public:
167 static std::unique_ptr<ImplType> create(size_t StartBit, size_t NumBits) {
168 size_t Size = sizeof(ImplType) + TrieSubtrie::sizeToAlloc(NumBits);
169 void *Memory = ::operator new(Size);
170 ImplType *Impl = ::new (Memory) ImplType(StartBit, NumBits);
171 return std::unique_ptr<ImplType>(Impl);
172 }
173
174 // Save the Subtrie into the ownship list of the trie structure in a
175 // thread-safe way. The ownership transfer is done by compare_exchange the
176 // pointer value inside the unique_ptr.
177 TrieSubtrie *save(std::unique_ptr<TrieSubtrie> S) {
178 assert(!S->Next && "Expected S to a freshly-constructed leaf");
179
180 TrieSubtrie *CurrentHead = nullptr;
181 // Add ownership of "S" to front of the list, so that Root -> S ->
182 // Root.Next. This works by repeatedly setting S->Next to a candidate value
183 // of Root.Next (initially nullptr), then setting Root.Next to S once the
184 // candidate matches reality.
185 while (!getRoot()->Next.compare_exchange_weak(p1&: CurrentHead, p2: S.get()))
186 S->Next.exchange(p: CurrentHead);
187
188 // Ownership transferred to subtrie successfully. Release the unique_ptr.
189 return S.release();
190 }
191
192 // Get the root which is the trailing object.
193 TrieSubtrie *getRoot() { return getTrailingObjects(); }
194
195 static void *operator new(size_t Size) { return ::operator new(Size); }
196 void operator delete(void *Ptr) { ::operator delete(Ptr); }
197
198 /// FIXME: This should take a function that allocates and constructs the
199 /// content lazily (taking the hash as a separate parameter), in case of
200 /// collision.
201 ThreadSafeAllocator<BumpPtrAllocator> ContentAlloc;
202
203private:
204 friend class TrailingObjects;
205
206 ImplType(size_t StartBit, size_t NumBits) {
207 ::new (getRoot()) TrieSubtrie(StartBit, NumBits);
208 }
209};
210
211ThreadSafeTrieRawHashMapBase::ImplType &
212ThreadSafeTrieRawHashMapBase::getOrCreateImpl() {
213 if (ImplType *Impl = ImplPtr.load())
214 return *Impl;
215
216 // Create a new ImplType and store it if another thread doesn't do so first.
217 // If another thread wins this one is destroyed locally.
218 std::unique_ptr<ImplType> Impl = ImplType::create(StartBit: 0, NumBits: NumRootBits);
219 ImplType *ExistingImpl = nullptr;
220
221 // If the ownership transferred succesfully, release unique_ptr and return
222 // the pointer to the new ImplType.
223 if (ImplPtr.compare_exchange_strong(p1&: ExistingImpl, p2: Impl.get()))
224 return *Impl.release();
225
226 // Already created, return the existing ImplType.
227 return *ExistingImpl;
228}
229
230ThreadSafeTrieRawHashMapBase::PointerBase
231ThreadSafeTrieRawHashMapBase::find(ArrayRef<uint8_t> Hash) const {
232 assert(!Hash.empty() && "Uninitialized hash");
233
234 ImplType *Impl = ImplPtr.load();
235 if (!Impl)
236 return PointerBase();
237
238 TrieSubtrie *S = Impl->getRoot();
239 TrieHashIndexGenerator IndexGen{.NumRootBits: NumRootBits, .NumSubtrieBits: NumSubtrieBits, .Bytes: Hash};
240 size_t Index = IndexGen.next();
241 while (Index != IndexGen.end()) {
242 // Try to set the content.
243 TrieNode *Existing = S->get(I: Index);
244 if (!Existing)
245 return PointerBase(S, Index, *IndexGen.StartBit);
246
247 // Check for an exact match.
248 if (auto *ExistingContent = dyn_cast<TrieContent>(Val: Existing))
249 return ExistingContent->getHash() == Hash
250 ? PointerBase(ExistingContent->getValuePointer())
251 : PointerBase(S, Index, *IndexGen.StartBit);
252
253 Index = IndexGen.next();
254 S = cast<TrieSubtrie>(Val: Existing);
255 }
256 llvm_unreachable("failed to locate the node after consuming all hash bytes");
257}
258
259ThreadSafeTrieRawHashMapBase::PointerBase ThreadSafeTrieRawHashMapBase::insert(
260 PointerBase Hint, ArrayRef<uint8_t> Hash,
261 function_ref<const uint8_t *(void *Mem, ArrayRef<uint8_t> Hash)>
262 Constructor) {
263 assert(!Hash.empty() && "Uninitialized hash");
264
265 ImplType &Impl = getOrCreateImpl();
266 TrieSubtrie *S = Impl.getRoot();
267 TrieHashIndexGenerator IndexGen{.NumRootBits: NumRootBits, .NumSubtrieBits: NumSubtrieBits, .Bytes: Hash};
268 size_t Index;
269 if (Hint.isHint()) {
270 S = static_cast<TrieSubtrie *>(Hint.P);
271 Index = IndexGen.hint(Index: Hint.I, Bit: Hint.B);
272 } else {
273 Index = IndexGen.next();
274 }
275
276 while (Index != IndexGen.end()) {
277 // Load the node from the slot, allocating and calling the constructor if
278 // the slot is empty.
279 bool Generated = false;
280 TrieNode &Existing = S->get(I: Index).loadOrGenerate(Generator: [&]() {
281 Generated = true;
282
283 // Construct the value itself at the tail.
284 uint8_t *Memory = reinterpret_cast<uint8_t *>(
285 Impl.ContentAlloc.Allocate(Size: ContentAllocSize, Align: ContentAllocAlign));
286 const uint8_t *HashStorage = Constructor(Memory + ContentOffset, Hash);
287
288 // Construct the TrieContent header, passing in the offset to the hash.
289 TrieContent *Content = ::new (Memory)
290 TrieContent(ContentOffset, Hash.size(), HashStorage - Memory);
291 assert(Hash == Content->getHash() && "Hash not properly initialized");
292 return Content;
293 });
294 // If we just generated it, return it!
295 if (Generated)
296 return PointerBase(cast<TrieContent>(Val&: Existing).getValuePointer());
297
298 if (auto *ST = dyn_cast<TrieSubtrie>(Val: &Existing)) {
299 S = ST;
300 Index = IndexGen.next();
301 continue;
302 }
303
304 // Return the existing content if it's an exact match!
305 auto &ExistingContent = cast<TrieContent>(Val&: Existing);
306 if (ExistingContent.getHash() == Hash)
307 return PointerBase(ExistingContent.getValuePointer());
308
309 // Sink the existing content as long as the indexes match.
310 size_t NextIndex = IndexGen.next();
311 while (NextIndex != IndexGen.end()) {
312 size_t NewIndexForExistingContent =
313 IndexGen.getCollidingBits(CollidingBits: ExistingContent.getHash());
314 S = S->sink(I: Index, Content&: ExistingContent, NumSubtrieBits: IndexGen.getNumBits(),
315 NewI: NewIndexForExistingContent,
316 Saver: [&Impl](std::unique_ptr<TrieSubtrie> S) {
317 return Impl.save(S: std::move(S));
318 });
319 Index = NextIndex;
320
321 // Found the difference.
322 if (NextIndex != NewIndexForExistingContent)
323 break;
324
325 NextIndex = IndexGen.next();
326 }
327 }
328 llvm_unreachable("failed to insert the node after consuming all hash bytes");
329}
330
331ThreadSafeTrieRawHashMapBase::ThreadSafeTrieRawHashMapBase(
332 size_t ContentAllocSize, size_t ContentAllocAlign, size_t ContentOffset,
333 std::optional<size_t> NumRootBits, std::optional<size_t> NumSubtrieBits)
334 : ContentAllocSize(ContentAllocSize), ContentAllocAlign(ContentAllocAlign),
335 ContentOffset(ContentOffset),
336 NumRootBits(NumRootBits.value_or(u: DefaultNumRootBits)),
337 NumSubtrieBits(NumSubtrieBits.value_or(u: DefaultNumSubtrieBits)),
338 ImplPtr(nullptr) {
339 // Assertion checks for reasonable configuration. The settings below are not
340 // hard limits on most platforms, but a reasonable configuration should fall
341 // within those limits.
342 assert((!NumRootBits || *NumRootBits < 20) &&
343 "Root should have fewer than ~1M slots");
344 assert((!NumSubtrieBits || *NumSubtrieBits < 10) &&
345 "Subtries should have fewer than ~1K slots");
346}
347
348ThreadSafeTrieRawHashMapBase::ThreadSafeTrieRawHashMapBase(
349 ThreadSafeTrieRawHashMapBase &&RHS)
350 : ContentAllocSize(RHS.ContentAllocSize),
351 ContentAllocAlign(RHS.ContentAllocAlign),
352 ContentOffset(RHS.ContentOffset), NumRootBits(RHS.NumRootBits),
353 NumSubtrieBits(RHS.NumSubtrieBits) {
354 // Steal the root from RHS.
355 ImplPtr = RHS.ImplPtr.exchange(p: nullptr);
356}
357
358ThreadSafeTrieRawHashMapBase::~ThreadSafeTrieRawHashMapBase() {
359 assert(!ImplPtr.load() && "Expected subclass to call destroyImpl()");
360}
361
362void ThreadSafeTrieRawHashMapBase::destroyImpl(
363 function_ref<void(void *)> Destructor) {
364 std::unique_ptr<ImplType> Impl(ImplPtr.exchange(p: nullptr));
365 if (!Impl)
366 return;
367
368 // Destroy content nodes throughout trie. Avoid destroying any subtries since
369 // we need TrieNode::classof() to find the content nodes.
370 //
371 // FIXME: Once we have bitsets (see FIXME in TrieSubtrie class), use them
372 // facilitate sparse iteration here.
373 if (Destructor)
374 for (TrieSubtrie *Trie = Impl->getRoot(); Trie; Trie = Trie->Next.load())
375 for (unsigned I = 0; I < Trie->size(); ++I)
376 if (auto *Content = dyn_cast_or_null<TrieContent>(Val: Trie->load(I)))
377 Destructor(Content->getValuePointer());
378
379 // Destroy the subtries. Incidentally, this destroys them in the reverse order
380 // of saving.
381 TrieSubtrie *Trie = Impl->getRoot()->Next;
382 while (Trie) {
383 TrieSubtrie *Next = Trie->Next.exchange(p: nullptr);
384 delete Trie;
385 Trie = Next;
386 }
387}
388
389ThreadSafeTrieRawHashMapBase::PointerBase
390ThreadSafeTrieRawHashMapBase::getRoot() const {
391 ImplType *Impl = ImplPtr.load();
392 if (!Impl)
393 return PointerBase();
394 return PointerBase(Impl->getRoot());
395}
396
397unsigned ThreadSafeTrieRawHashMapBase::getStartBit(
398 ThreadSafeTrieRawHashMapBase::PointerBase P) const {
399 assert(!P.isHint() && "Not a valid trie");
400 if (!P.P)
401 return 0;
402 if (auto *S = dyn_cast<TrieSubtrie>(Val: (TrieNode *)P.P))
403 return S->StartBit;
404 return 0;
405}
406
407unsigned ThreadSafeTrieRawHashMapBase::getNumBits(
408 ThreadSafeTrieRawHashMapBase::PointerBase P) const {
409 assert(!P.isHint() && "Not a valid trie");
410 if (!P.P)
411 return 0;
412 if (auto *S = dyn_cast<TrieSubtrie>(Val: (TrieNode *)P.P))
413 return S->NumBits;
414 return 0;
415}
416
417unsigned ThreadSafeTrieRawHashMapBase::getNumSlotUsed(
418 ThreadSafeTrieRawHashMapBase::PointerBase P) const {
419 assert(!P.isHint() && "Not a valid trie");
420 if (!P.P)
421 return 0;
422 auto *S = dyn_cast<TrieSubtrie>(Val: (TrieNode *)P.P);
423 if (!S)
424 return 0;
425 unsigned Num = 0;
426 for (unsigned I = 0, E = S->size(); I < E; ++I)
427 if (S->load(I))
428 ++Num;
429 return Num;
430}
431
432std::string ThreadSafeTrieRawHashMapBase::getTriePrefixAsString(
433 ThreadSafeTrieRawHashMapBase::PointerBase P) const {
434 assert(!P.isHint() && "Not a valid trie");
435 if (!P.P)
436 return "";
437
438 auto *S = dyn_cast<TrieSubtrie>(Val: (TrieNode *)P.P);
439 if (!S || !S->IsSubtrie)
440 return "";
441
442 // Find a TrieContent node which has hash stored. Depth search following the
443 // first used slot until a TrieContent node is found.
444 TrieSubtrie *Current = S;
445 TrieContent *Node = nullptr;
446 while (Current) {
447 TrieSubtrie *Next = nullptr;
448 // Find first used slot in the trie.
449 for (unsigned I = 0, E = Current->size(); I < E; ++I) {
450 auto *S = Current->load(I);
451 if (!S)
452 continue;
453
454 if (auto *Content = dyn_cast<TrieContent>(Val: S))
455 Node = Content;
456 else if (auto *Sub = dyn_cast<TrieSubtrie>(Val: S))
457 Next = Sub;
458 break;
459 }
460
461 // Found the node.
462 if (Node)
463 break;
464
465 // Continue to the next level if the node is not found.
466 Current = Next;
467 }
468
469 assert(Node && "malformed trie, cannot find TrieContent on leaf node");
470 // The prefix for the current trie is the first `StartBit` of the content
471 // stored underneath this subtrie.
472 std::string Str;
473 raw_string_ostream SS(Str);
474
475 unsigned StartFullBytes = (S->StartBit + 1) / 8 - 1;
476 SS << toHex(Input: toStringRef(Input: Node->getHash()).take_front(N: StartFullBytes),
477 /*LowerCase=*/true);
478
479 // For the part of the prefix that doesn't fill a byte, print raw bit values.
480 std::string Bits;
481 for (unsigned I = StartFullBytes * 8, E = S->StartBit; I < E; ++I) {
482 unsigned Index = I / 8;
483 unsigned Offset = 7 - I % 8;
484 Bits.push_back(c: '0' + ((Node->getHash()[Index] >> Offset) & 1));
485 }
486
487 if (!Bits.empty())
488 SS << "[" << Bits << "]";
489
490 return SS.str();
491}
492
493unsigned ThreadSafeTrieRawHashMapBase::getNumTries() const {
494 ImplType *Impl = ImplPtr.load();
495 if (!Impl)
496 return 0;
497 unsigned Num = 0;
498 for (TrieSubtrie *Trie = Impl->getRoot(); Trie; Trie = Trie->Next.load())
499 ++Num;
500 return Num;
501}
502
503ThreadSafeTrieRawHashMapBase::PointerBase
504ThreadSafeTrieRawHashMapBase::getNextTrie(
505 ThreadSafeTrieRawHashMapBase::PointerBase P) const {
506 assert(!P.isHint() && "Not a valid trie");
507 if (!P.P)
508 return PointerBase();
509 auto *S = dyn_cast<TrieSubtrie>(Val: (TrieNode *)P.P);
510 if (!S)
511 return PointerBase();
512 if (auto *E = S->Next.load())
513 return PointerBase(E);
514 return PointerBase();
515}
516