1276789Sdim//===-- sanitizer_stackdepotbase.h ------------------------------*- C++ -*-===//
2276789Sdim//
3353358Sdim// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4353358Sdim// See https://llvm.org/LICENSE.txt for license information.
5353358Sdim// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6276789Sdim//
7276789Sdim//===----------------------------------------------------------------------===//
8276789Sdim//
9276789Sdim// Implementation of a mapping from arbitrary values to unique 32-bit
10276789Sdim// identifiers.
11276789Sdim//===----------------------------------------------------------------------===//
12296417Sdim
13276789Sdim#ifndef SANITIZER_STACKDEPOTBASE_H
14276789Sdim#define SANITIZER_STACKDEPOTBASE_H
15276789Sdim
16276789Sdim#include "sanitizer_internal_defs.h"
17276789Sdim#include "sanitizer_mutex.h"
18276789Sdim#include "sanitizer_atomic.h"
19276789Sdim#include "sanitizer_persistent_allocator.h"
20276789Sdim
21276789Sdimnamespace __sanitizer {
22276789Sdim
23276789Sdimtemplate <class Node, int kReservedBits, int kTabSizeLog>
24276789Sdimclass StackDepotBase {
25276789Sdim public:
26276789Sdim  typedef typename Node::args_type args_type;
27276789Sdim  typedef typename Node::handle_type handle_type;
28276789Sdim  // Maps stack trace to an unique id.
29296417Sdim  handle_type Put(args_type args, bool *inserted = nullptr);
30276789Sdim  // Retrieves a stored stack trace by the id.
31276789Sdim  args_type Get(u32 id);
32276789Sdim
33276789Sdim  StackDepotStats *GetStats() { return &stats; }
34276789Sdim
35276789Sdim  void LockAll();
36276789Sdim  void UnlockAll();
37276789Sdim
38276789Sdim private:
39276789Sdim  static Node *find(Node *s, args_type args, u32 hash);
40276789Sdim  static Node *lock(atomic_uintptr_t *p);
41276789Sdim  static void unlock(atomic_uintptr_t *p, Node *s);
42276789Sdim
43276789Sdim  static const int kTabSize = 1 << kTabSizeLog;  // Hash table size.
44276789Sdim  static const int kPartBits = 8;
45276789Sdim  static const int kPartShift = sizeof(u32) * 8 - kPartBits - kReservedBits;
46276789Sdim  static const int kPartCount =
47276789Sdim      1 << kPartBits;  // Number of subparts in the table.
48276789Sdim  static const int kPartSize = kTabSize / kPartCount;
49276789Sdim  static const int kMaxId = 1 << kPartShift;
50276789Sdim
51276789Sdim  atomic_uintptr_t tab[kTabSize];   // Hash table of Node's.
52276789Sdim  atomic_uint32_t seq[kPartCount];  // Unique id generators.
53276789Sdim
54276789Sdim  StackDepotStats stats;
55276789Sdim
56276789Sdim  friend class StackDepotReverseMap;
57276789Sdim};
58276789Sdim
59276789Sdimtemplate <class Node, int kReservedBits, int kTabSizeLog>
60276789SdimNode *StackDepotBase<Node, kReservedBits, kTabSizeLog>::find(Node *s,
61276789Sdim                                                             args_type args,
62276789Sdim                                                             u32 hash) {
63276789Sdim  // Searches linked list s for the stack, returns its id.
64276789Sdim  for (; s; s = s->link) {
65276789Sdim    if (s->eq(hash, args)) {
66276789Sdim      return s;
67276789Sdim    }
68276789Sdim  }
69296417Sdim  return nullptr;
70276789Sdim}
71276789Sdim
72276789Sdimtemplate <class Node, int kReservedBits, int kTabSizeLog>
73276789SdimNode *StackDepotBase<Node, kReservedBits, kTabSizeLog>::lock(
74276789Sdim    atomic_uintptr_t *p) {
75276789Sdim  // Uses the pointer lsb as mutex.
76276789Sdim  for (int i = 0;; i++) {
77276789Sdim    uptr cmp = atomic_load(p, memory_order_relaxed);
78276789Sdim    if ((cmp & 1) == 0 &&
79276789Sdim        atomic_compare_exchange_weak(p, &cmp, cmp | 1, memory_order_acquire))
80276789Sdim      return (Node *)cmp;
81276789Sdim    if (i < 10)
82276789Sdim      proc_yield(10);
83276789Sdim    else
84276789Sdim      internal_sched_yield();
85276789Sdim  }
86276789Sdim}
87276789Sdim
88276789Sdimtemplate <class Node, int kReservedBits, int kTabSizeLog>
89276789Sdimvoid StackDepotBase<Node, kReservedBits, kTabSizeLog>::unlock(
90276789Sdim    atomic_uintptr_t *p, Node *s) {
91276789Sdim  DCHECK_EQ((uptr)s & 1, 0);
92276789Sdim  atomic_store(p, (uptr)s, memory_order_release);
93276789Sdim}
94276789Sdim
95276789Sdimtemplate <class Node, int kReservedBits, int kTabSizeLog>
96276789Sdimtypename StackDepotBase<Node, kReservedBits, kTabSizeLog>::handle_type
97276789SdimStackDepotBase<Node, kReservedBits, kTabSizeLog>::Put(args_type args,
98276789Sdim                                                      bool *inserted) {
99276789Sdim  if (inserted) *inserted = false;
100276789Sdim  if (!Node::is_valid(args)) return handle_type();
101276789Sdim  uptr h = Node::hash(args);
102276789Sdim  atomic_uintptr_t *p = &tab[h % kTabSize];
103276789Sdim  uptr v = atomic_load(p, memory_order_consume);
104276789Sdim  Node *s = (Node *)(v & ~1);
105276789Sdim  // First, try to find the existing stack.
106276789Sdim  Node *node = find(s, args, h);
107276789Sdim  if (node) return node->get_handle();
108276789Sdim  // If failed, lock, retry and insert new.
109276789Sdim  Node *s2 = lock(p);
110276789Sdim  if (s2 != s) {
111276789Sdim    node = find(s2, args, h);
112276789Sdim    if (node) {
113276789Sdim      unlock(p, s2);
114276789Sdim      return node->get_handle();
115276789Sdim    }
116276789Sdim  }
117276789Sdim  uptr part = (h % kTabSize) / kPartSize;
118276789Sdim  u32 id = atomic_fetch_add(&seq[part], 1, memory_order_relaxed) + 1;
119276789Sdim  stats.n_uniq_ids++;
120276789Sdim  CHECK_LT(id, kMaxId);
121276789Sdim  id |= part << kPartShift;
122276789Sdim  CHECK_NE(id, 0);
123276789Sdim  CHECK_EQ(id & (((u32)-1) >> kReservedBits), id);
124276789Sdim  uptr memsz = Node::storage_size(args);
125276789Sdim  s = (Node *)PersistentAlloc(memsz);
126276789Sdim  stats.allocated += memsz;
127276789Sdim  s->id = id;
128276789Sdim  s->store(args, h);
129276789Sdim  s->link = s2;
130276789Sdim  unlock(p, s);
131276789Sdim  if (inserted) *inserted = true;
132276789Sdim  return s->get_handle();
133276789Sdim}
134276789Sdim
135276789Sdimtemplate <class Node, int kReservedBits, int kTabSizeLog>
136276789Sdimtypename StackDepotBase<Node, kReservedBits, kTabSizeLog>::args_type
137276789SdimStackDepotBase<Node, kReservedBits, kTabSizeLog>::Get(u32 id) {
138276789Sdim  if (id == 0) {
139276789Sdim    return args_type();
140276789Sdim  }
141276789Sdim  CHECK_EQ(id & (((u32)-1) >> kReservedBits), id);
142276789Sdim  // High kPartBits contain part id, so we need to scan at most kPartSize lists.
143276789Sdim  uptr part = id >> kPartShift;
144276789Sdim  for (int i = 0; i != kPartSize; i++) {
145276789Sdim    uptr idx = part * kPartSize + i;
146276789Sdim    CHECK_LT(idx, kTabSize);
147276789Sdim    atomic_uintptr_t *p = &tab[idx];
148276789Sdim    uptr v = atomic_load(p, memory_order_consume);
149276789Sdim    Node *s = (Node *)(v & ~1);
150276789Sdim    for (; s; s = s->link) {
151276789Sdim      if (s->id == id) {
152276789Sdim        return s->load();
153276789Sdim      }
154276789Sdim    }
155276789Sdim  }
156276789Sdim  return args_type();
157276789Sdim}
158276789Sdim
159276789Sdimtemplate <class Node, int kReservedBits, int kTabSizeLog>
160276789Sdimvoid StackDepotBase<Node, kReservedBits, kTabSizeLog>::LockAll() {
161276789Sdim  for (int i = 0; i < kTabSize; ++i) {
162276789Sdim    lock(&tab[i]);
163276789Sdim  }
164276789Sdim}
165276789Sdim
166276789Sdimtemplate <class Node, int kReservedBits, int kTabSizeLog>
167276789Sdimvoid StackDepotBase<Node, kReservedBits, kTabSizeLog>::UnlockAll() {
168276789Sdim  for (int i = 0; i < kTabSize; ++i) {
169276789Sdim    atomic_uintptr_t *p = &tab[i];
170276789Sdim    uptr s = atomic_load(p, memory_order_relaxed);
171276789Sdim    unlock(p, (Node *)(s & ~1UL));
172276789Sdim  }
173276789Sdim}
174276789Sdim
175296417Sdim} // namespace __sanitizer
176296417Sdim
177296417Sdim#endif // SANITIZER_STACKDEPOTBASE_H
178