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