1//===-- sanitizer_addrhashmap.h ---------------------------------*- 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// Concurrent uptr->T hashmap.
10//
11//===----------------------------------------------------------------------===//
12
13#ifndef SANITIZER_ADDRHASHMAP_H
14#define SANITIZER_ADDRHASHMAP_H
15
16#include "sanitizer_common.h"
17#include "sanitizer_mutex.h"
18#include "sanitizer_atomic.h"
19#include "sanitizer_allocator_internal.h"
20
21namespace __sanitizer {
22
23// Concurrent uptr->T hashmap.
24// T must be a POD type, kSize is preferably a prime but can be any number.
25// Usage example:
26//
27// typedef AddrHashMap<uptr, 11> Map;
28// Map m;
29// {
30//   Map::Handle h(&m, addr);
31//   use h.operator->() to access the data
32//   if h.created() then the element was just created, and the current thread
33//     has exclusive access to it
34//   otherwise the current thread has only read access to the data
35// }
36// {
37//   Map::Handle h(&m, addr, true);
38//   this will remove the data from the map in Handle dtor
39//   the current thread has exclusive access to the data
40//   if !h.exists() then the element never existed
41// }
42// {
43//   Map::Handle h(&m, addr, false, true);
44//   this will create a new element or return a handle to an existing element
45//   if !h.created() this thread does *not* have exclusive access to the data
46// }
47template<typename T, uptr kSize>
48class AddrHashMap {
49 private:
50  struct Cell {
51    atomic_uintptr_t addr;
52    T                val;
53  };
54
55  struct AddBucket {
56    uptr cap;
57    uptr size;
58    Cell cells[1];  // variable len
59  };
60
61  static const uptr kBucketSize = 3;
62
63  struct Bucket {
64    Mutex mtx;
65    atomic_uintptr_t add;
66    Cell             cells[kBucketSize];
67  };
68
69 public:
70  AddrHashMap();
71
72  class Handle {
73   public:
74    Handle(AddrHashMap<T, kSize> *map, uptr addr);
75    Handle(AddrHashMap<T, kSize> *map, uptr addr, bool remove);
76    Handle(AddrHashMap<T, kSize> *map, uptr addr, bool remove, bool create);
77
78    ~Handle();
79    T *operator->();
80    T &operator*();
81    const T &operator*() const;
82    bool created() const;
83    bool exists() const;
84
85   private:
86    friend AddrHashMap<T, kSize>;
87    AddrHashMap<T, kSize> *map_;
88    Bucket                *bucket_;
89    Cell                  *cell_;
90    uptr                   addr_;
91    uptr                   addidx_;
92    bool                   created_;
93    bool                   remove_;
94    bool                   create_;
95  };
96
97  typedef void (*ForEachCallback)(const uptr key, const T &val, void *arg);
98  // ForEach acquires a lock on each bucket while iterating over
99  // elements. Note that this only ensures that the structure of the hashmap is
100  // unchanged, there may be a data race to the element itself.
101  void ForEach(ForEachCallback cb, void *arg);
102
103 private:
104  friend class Handle;
105  Bucket *table_;
106
107  void acquire(Handle *h);
108  void release(Handle *h);
109  uptr calcHash(uptr addr);
110};
111
112template <typename T, uptr kSize>
113void AddrHashMap<T, kSize>::ForEach(ForEachCallback cb, void *arg) {
114  for (uptr n = 0; n < kSize; n++) {
115    Bucket *bucket = &table_[n];
116
117    ReadLock lock(&bucket->mtx);
118
119    for (uptr i = 0; i < kBucketSize; i++) {
120      Cell *c = &bucket->cells[i];
121      uptr addr1 = atomic_load(&c->addr, memory_order_acquire);
122      if (addr1 != 0)
123        cb(addr1, c->val, arg);
124    }
125
126    // Iterate over any additional cells.
127    if (AddBucket *add =
128            (AddBucket *)atomic_load(&bucket->add, memory_order_acquire)) {
129      for (uptr i = 0; i < add->size; i++) {
130        Cell *c = &add->cells[i];
131        uptr addr1 = atomic_load(&c->addr, memory_order_acquire);
132        if (addr1 != 0)
133          cb(addr1, c->val, arg);
134      }
135    }
136  }
137}
138
139template<typename T, uptr kSize>
140AddrHashMap<T, kSize>::Handle::Handle(AddrHashMap<T, kSize> *map, uptr addr) {
141  map_ = map;
142  addr_ = addr;
143  remove_ = false;
144  create_ = true;
145  map_->acquire(this);
146}
147
148template<typename T, uptr kSize>
149AddrHashMap<T, kSize>::Handle::Handle(AddrHashMap<T, kSize> *map, uptr addr,
150    bool remove) {
151  map_ = map;
152  addr_ = addr;
153  remove_ = remove;
154  create_ = true;
155  map_->acquire(this);
156}
157
158template<typename T, uptr kSize>
159AddrHashMap<T, kSize>::Handle::Handle(AddrHashMap<T, kSize> *map, uptr addr,
160    bool remove, bool create) {
161  map_ = map;
162  addr_ = addr;
163  remove_ = remove;
164  create_ = create;
165  map_->acquire(this);
166}
167
168template<typename T, uptr kSize>
169AddrHashMap<T, kSize>::Handle::~Handle() {
170  map_->release(this);
171}
172
173template <typename T, uptr kSize>
174T *AddrHashMap<T, kSize>::Handle::operator->() {
175  return &cell_->val;
176}
177
178template <typename T, uptr kSize>
179const T &AddrHashMap<T, kSize>::Handle::operator*() const {
180  return cell_->val;
181}
182
183template <typename T, uptr kSize>
184T &AddrHashMap<T, kSize>::Handle::operator*() {
185  return cell_->val;
186}
187
188template<typename T, uptr kSize>
189bool AddrHashMap<T, kSize>::Handle::created() const {
190  return created_;
191}
192
193template<typename T, uptr kSize>
194bool AddrHashMap<T, kSize>::Handle::exists() const {
195  return cell_ != nullptr;
196}
197
198template<typename T, uptr kSize>
199AddrHashMap<T, kSize>::AddrHashMap() {
200  table_ = (Bucket*)MmapOrDie(kSize * sizeof(table_[0]), "AddrHashMap");
201}
202
203template <typename T, uptr kSize>
204void AddrHashMap<T, kSize>::acquire(Handle *h)
205    SANITIZER_NO_THREAD_SAFETY_ANALYSIS {
206  uptr addr = h->addr_;
207  uptr hash = calcHash(addr);
208  Bucket *b = &table_[hash];
209
210  h->created_ = false;
211  h->addidx_ = -1U;
212  h->bucket_ = b;
213  h->cell_ = nullptr;
214
215  // If we want to remove the element, we need exclusive access to the bucket,
216  // so skip the lock-free phase.
217  if (h->remove_)
218    goto locked;
219
220 retry:
221  // First try to find an existing element w/o read mutex.
222  CHECK(!h->remove_);
223  // Check the embed cells.
224  for (uptr i = 0; i < kBucketSize; i++) {
225    Cell *c = &b->cells[i];
226    uptr addr1 = atomic_load(&c->addr, memory_order_acquire);
227    if (addr1 == addr) {
228      h->cell_ = c;
229      return;
230    }
231  }
232
233  // Check the add cells with read lock.
234  if (atomic_load(&b->add, memory_order_relaxed)) {
235    b->mtx.ReadLock();
236    AddBucket *add = (AddBucket*)atomic_load(&b->add, memory_order_relaxed);
237    for (uptr i = 0; i < add->size; i++) {
238      Cell *c = &add->cells[i];
239      uptr addr1 = atomic_load(&c->addr, memory_order_relaxed);
240      if (addr1 == addr) {
241        h->addidx_ = i;
242        h->cell_ = c;
243        return;
244      }
245    }
246    b->mtx.ReadUnlock();
247  }
248
249 locked:
250  // Re-check existence under write lock.
251  // Embed cells.
252  b->mtx.Lock();
253  for (uptr i = 0; i < kBucketSize; i++) {
254    Cell *c = &b->cells[i];
255    uptr addr1 = atomic_load(&c->addr, memory_order_relaxed);
256    if (addr1 == addr) {
257      if (h->remove_) {
258        h->cell_ = c;
259        return;
260      }
261      b->mtx.Unlock();
262      goto retry;
263    }
264  }
265
266  // Add cells.
267  AddBucket *add = (AddBucket*)atomic_load(&b->add, memory_order_relaxed);
268  if (add) {
269    for (uptr i = 0; i < add->size; i++) {
270      Cell *c = &add->cells[i];
271      uptr addr1 = atomic_load(&c->addr, memory_order_relaxed);
272      if (addr1 == addr) {
273        if (h->remove_) {
274          h->addidx_ = i;
275          h->cell_ = c;
276          return;
277        }
278        b->mtx.Unlock();
279        goto retry;
280      }
281    }
282  }
283
284  // The element does not exist, no need to create it if we want to remove.
285  if (h->remove_ || !h->create_) {
286    b->mtx.Unlock();
287    return;
288  }
289
290  // Now try to create it under the mutex.
291  h->created_ = true;
292  // See if we have a free embed cell.
293  for (uptr i = 0; i < kBucketSize; i++) {
294    Cell *c = &b->cells[i];
295    uptr addr1 = atomic_load(&c->addr, memory_order_relaxed);
296    if (addr1 == 0) {
297      h->cell_ = c;
298      return;
299    }
300  }
301
302  // Store in the add cells.
303  if (!add) {
304    // Allocate a new add array.
305    const uptr kInitSize = 64;
306    add = (AddBucket*)InternalAlloc(kInitSize);
307    internal_memset(add, 0, kInitSize);
308    add->cap = (kInitSize - sizeof(*add)) / sizeof(add->cells[0]) + 1;
309    add->size = 0;
310    atomic_store(&b->add, (uptr)add, memory_order_relaxed);
311  }
312  if (add->size == add->cap) {
313    // Grow existing add array.
314    uptr oldsize = sizeof(*add) + (add->cap - 1) * sizeof(add->cells[0]);
315    uptr newsize = oldsize * 2;
316    AddBucket *add1 = (AddBucket*)InternalAlloc(newsize);
317    internal_memset(add1, 0, newsize);
318    add1->cap = (newsize - sizeof(*add)) / sizeof(add->cells[0]) + 1;
319    add1->size = add->size;
320    internal_memcpy(add1->cells, add->cells, add->size * sizeof(add->cells[0]));
321    InternalFree(add);
322    atomic_store(&b->add, (uptr)add1, memory_order_relaxed);
323    add = add1;
324  }
325  // Store.
326  uptr i = add->size++;
327  Cell *c = &add->cells[i];
328  CHECK_EQ(atomic_load(&c->addr, memory_order_relaxed), 0);
329  h->addidx_ = i;
330  h->cell_ = c;
331 }
332
333 template <typename T, uptr kSize>
334 void AddrHashMap<T, kSize>::release(Handle *h)
335     SANITIZER_NO_THREAD_SAFETY_ANALYSIS {
336   if (!h->cell_)
337     return;
338   Bucket *b = h->bucket_;
339   Cell *c = h->cell_;
340   uptr addr1 = atomic_load(&c->addr, memory_order_relaxed);
341   if (h->created_) {
342     // Denote completion of insertion.
343     CHECK_EQ(addr1, 0);
344     // After the following store, the element becomes available
345     // for lock-free reads.
346     atomic_store(&c->addr, h->addr_, memory_order_release);
347     b->mtx.Unlock();
348   } else if (h->remove_) {
349     // Denote that the cell is empty now.
350     CHECK_EQ(addr1, h->addr_);
351     atomic_store(&c->addr, 0, memory_order_release);
352     // See if we need to compact the bucket.
353     AddBucket *add = (AddBucket *)atomic_load(&b->add, memory_order_relaxed);
354     if (h->addidx_ == -1U) {
355       // Removed from embed array, move an add element into the freed cell.
356       if (add && add->size != 0) {
357         uptr last = --add->size;
358         Cell *c1 = &add->cells[last];
359         c->val = c1->val;
360         uptr addr1 = atomic_load(&c1->addr, memory_order_relaxed);
361         atomic_store(&c->addr, addr1, memory_order_release);
362         atomic_store(&c1->addr, 0, memory_order_release);
363       }
364     } else {
365       // Removed from add array, compact it.
366       uptr last = --add->size;
367       Cell *c1 = &add->cells[last];
368       if (c != c1) {
369         *c = *c1;
370         atomic_store(&c1->addr, 0, memory_order_relaxed);
371       }
372     }
373     if (add && add->size == 0) {
374       // FIXME(dvyukov): free add?
375     }
376     b->mtx.Unlock();
377   } else {
378     CHECK_EQ(addr1, h->addr_);
379     if (h->addidx_ != -1U)
380       b->mtx.ReadUnlock();
381   }
382 }
383
384template<typename T, uptr kSize>
385uptr AddrHashMap<T, kSize>::calcHash(uptr addr) {
386  addr += addr << 10;
387  addr ^= addr >> 6;
388  return addr % kSize;
389}
390
391} // namespace __sanitizer
392
393#endif // SANITIZER_ADDRHASHMAP_H
394