Skip to content
rachel/workerd
File

Blob: src/workerd/util/small-set-test.c++

8.2 KB
1// Copyright (c) 2017-2022 Cloudflare, Inc.
2// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3//     https://opensource.org/licenses/Apache-2.0
4 
5#include "small-set.h"
6#include "weak-refs.h"
7 
8#include <kj/refcount.h>
9#include <kj/test.h>
10 
11namespace workerd {
12namespace {
13 
14// A simple refcounted type for testing SmallSet.
15// Supports WeakRef creation for testing forEach.
16struct TestItem final: public kj::Refcounted {
17  int value;
18  kj::Rc<WeakRef<TestItem>> selfRef;
19 
20  explicit TestItem(int v)
21      : value(v),
22        selfRef(kj::rc<WeakRef<TestItem>>(kj::Badge<TestItem>{}, *this)) {}
23 
24  ~TestItem() noexcept(false) {
25    selfRef->invalidate();
26  }
27 
28  kj::Rc<WeakRef<TestItem>> getWeakRef() {
29    return selfRef.addRef();
30  }
31};
32 
33// Helper to check if set contains an item with a given value.
34bool containsValue(SmallSet<kj::Rc<TestItem>>& set, int value) {
35  return set.containsIf([value](const kj::Rc<TestItem>& item) { return item->value == value; });
36}
37 
38// Helper to remove an item with a given value.
39bool removeValue(SmallSet<kj::Rc<TestItem>>& set, int value) {
40  return set.removeIf([value](const kj::Rc<TestItem>& item) { return item->value == value; });
41}
42 
43// Helper for WeakRef-based tests (used with forEach).
44bool removeWeakRefValue(SmallSet<kj::Rc<WeakRef<TestItem>>>& set, int value) {
45  return set.removeIf([value](const kj::Rc<WeakRef<TestItem>>& item) {
46    KJ_IF_SOME(ref, item->tryGet()) {
47      return ref.value == value;
48    }
49    return false;
50  });
51}
52 
53KJ_TEST("SmallSet: empty set") {
54  SmallSet<kj::Rc<TestItem>> set;
55  KJ_EXPECT(set.empty());
56  KJ_EXPECT(set.size() == 0);
57  KJ_EXPECT(!containsValue(set, 42));
58}
59 
60KJ_TEST("SmallSet: add and remove single item") {
61  SmallSet<kj::Rc<TestItem>> set;
62 
63  set.add(kj::rc<TestItem>(1));
64  KJ_EXPECT(!set.empty());
65  KJ_EXPECT(set.size() == 1);
66  KJ_EXPECT(containsValue(set, 1));
67 
68  KJ_EXPECT(removeValue(set, 1));
69  KJ_EXPECT(set.empty());
70  KJ_EXPECT(set.size() == 0);
71  KJ_EXPECT(!containsValue(set, 1));
72 
73  // Removing again should return false
74  KJ_EXPECT(!removeValue(set, 1));
75}
76 
77KJ_TEST("SmallSet: add and remove two items") {
78  SmallSet<kj::Rc<TestItem>> set;
79 
80  set.add(kj::rc<TestItem>(1));
81  set.add(kj::rc<TestItem>(2));
82  KJ_EXPECT(set.size() == 2);
83  KJ_EXPECT(containsValue(set, 1));
84  KJ_EXPECT(containsValue(set, 2));
85 
86  KJ_EXPECT(removeValue(set, 1));
87  KJ_EXPECT(set.size() == 1);
88  KJ_EXPECT(!containsValue(set, 1));
89  KJ_EXPECT(containsValue(set, 2));
90 
91  KJ_EXPECT(removeValue(set, 2));
92  KJ_EXPECT(set.empty());
93}
94 
95KJ_TEST("SmallSet: add and remove multiple items") {
96  SmallSet<kj::Rc<TestItem>> set;
97 
98  set.add(kj::rc<TestItem>(1));
99  set.add(kj::rc<TestItem>(2));
100  set.add(kj::rc<TestItem>(3));
101  set.add(kj::rc<TestItem>(4));
102  KJ_EXPECT(set.size() == 4);
103 
104  KJ_EXPECT(containsValue(set, 1));
105  KJ_EXPECT(containsValue(set, 2));
106  KJ_EXPECT(containsValue(set, 3));
107  KJ_EXPECT(containsValue(set, 4));
108 
109  KJ_EXPECT(removeValue(set, 2));
110  KJ_EXPECT(set.size() == 3);
111  KJ_EXPECT(!containsValue(set, 2));
112 
113  KJ_EXPECT(removeValue(set, 3));
114  KJ_EXPECT(set.size() == 2);
115  KJ_EXPECT(containsValue(set, 1));
116  KJ_EXPECT(containsValue(set, 4));
117 
118  KJ_EXPECT(removeValue(set, 1));
119  KJ_EXPECT(set.size() == 1);
120  KJ_EXPECT(containsValue(set, 4));
121 
122  KJ_EXPECT(removeValue(set, 4));
123  KJ_EXPECT(set.empty());
124}
125 
126KJ_TEST("SmallSet: state transitions") {
127  SmallSet<kj::Rc<TestItem>> set;
128 
129  // None -> Single
130  set.add(kj::rc<TestItem>(1));
131  KJ_EXPECT(set.size() == 1);
132 
133  // Single -> Double
134  set.add(kj::rc<TestItem>(2));
135  KJ_EXPECT(set.size() == 2);
136 
137  // Double -> Multiple
138  set.add(kj::rc<TestItem>(3));
139  KJ_EXPECT(set.size() == 3);
140 
141  // Multiple stays Multiple
142  set.add(kj::rc<TestItem>(4));
143  KJ_EXPECT(set.size() == 4);
144 
145  // Multiple -> Multiple (one less)
146  KJ_EXPECT(removeValue(set, 4));
147  KJ_EXPECT(set.size() == 3);
148 
149  // Multiple -> Double
150  KJ_EXPECT(removeValue(set, 3));
151  KJ_EXPECT(set.size() == 2);
152 
153  // Double -> Single
154  KJ_EXPECT(removeValue(set, 2));
155  KJ_EXPECT(set.size() == 1);
156 
157  // Single -> None
158  KJ_EXPECT(removeValue(set, 1));
159  KJ_EXPECT(set.size() == 0);
160}
161 
162KJ_TEST("SmallSet: iteration") {
163  SmallSet<kj::Rc<TestItem>> set;
164 
165  // Empty iteration
166  int count = 0;
167  for (auto& item: set) {
168    (void)item;
169    count++;
170  }
171  KJ_EXPECT(count == 0);
172 
173  // Single item
174  set.add(kj::rc<TestItem>(1));
175  count = 0;
176  for (auto& item: set) {
177    KJ_EXPECT(item->value == 1);
178    count++;
179  }
180  KJ_EXPECT(count == 1);
181 
182  // Two items
183  set.add(kj::rc<TestItem>(2));
184  count = 0;
185  for (auto& item: set) {
186    KJ_EXPECT(item->value == 1 || item->value == 2);
187    count++;
188  }
189  KJ_EXPECT(count == 2);
190 
191  // Multiple items
192  set.add(kj::rc<TestItem>(3));
193  count = 0;
194  for (auto& item: set) {
195    KJ_EXPECT(item->value >= 1 && item->value <= 3);
196    count++;
197  }
198  KJ_EXPECT(count == 3);
199}
200 
201KJ_TEST("SmallSet: clear") {
202  SmallSet<kj::Rc<TestItem>> set;
203 
204  set.add(kj::rc<TestItem>(1));
205  set.add(kj::rc<TestItem>(2));
206  set.add(kj::rc<TestItem>(3));
207  KJ_EXPECT(set.size() == 3);
208 
209  set.clear();
210  KJ_EXPECT(set.empty());
211  KJ_EXPECT(set.size() == 0);
212  KJ_EXPECT(!containsValue(set, 1));
213  KJ_EXPECT(!containsValue(set, 2));
214  KJ_EXPECT(!containsValue(set, 3));
215}
216 
217KJ_TEST("SmallSet: forEach for safe iteration during modification") {
218  // This simulates the queue.h use case where items may be removed during iteration.
219  // forEach requires WeakRef items to handle invalidation during iteration.
220  SmallSet<kj::Rc<WeakRef<TestItem>>> set;
221 
222  auto item1 = kj::rc<TestItem>(1);
223  auto item2 = kj::rc<TestItem>(2);
224  auto item3 = kj::rc<TestItem>(3);
225 
226  set.add(item1->getWeakRef());
227  set.add(item2->getWeakRef());
228  set.add(item3->getWeakRef());
229  KJ_EXPECT(set.size() == 3);
230 
231  // forEach takes a snapshot internally, so items remain valid even if removed during iteration.
232  kj::Vector<int> foundValues;
233  set.forEach([&](TestItem& item) {
234    foundValues.add(item.value);
235    // Remove the item from the set during iteration - this is safe because
236    // forEach iterates over an internal snapshot, not the set itself.
237    removeWeakRefValue(set, item.value);
238  });
239 
240  KJ_EXPECT(foundValues.size() == 3);
241  KJ_EXPECT(set.empty());
242}
243 
244KJ_TEST("SmallSet: forEach from single state") {
245  SmallSet<kj::Rc<WeakRef<TestItem>>> set;
246 
247  auto item = kj::rc<TestItem>(42);
248  set.add(item->getWeakRef());
249 
250  size_t count = 0;
251  set.forEach([&](TestItem& ref) {
252    KJ_EXPECT(ref.value == 42);
253    count++;
254  });
255  KJ_EXPECT(count == 1);
256}
257 
258KJ_TEST("SmallSet: forEach from double state") {
259  SmallSet<kj::Rc<WeakRef<TestItem>>> set;
260 
261  auto item1 = kj::rc<TestItem>(1);
262  auto item2 = kj::rc<TestItem>(2);
263 
264  set.add(item1->getWeakRef());
265  set.add(item2->getWeakRef());
266 
267  kj::Vector<int> foundValues;
268  set.forEach([&](TestItem& ref) { foundValues.add(ref.value); });
269  KJ_EXPECT(foundValues.size() == 2);
270  // Order doesn't matter for set semantics, just check both values are present
271  bool found1 = false, found2 = false;
272  for (auto v: foundValues) {
273    if (v == 1) found1 = true;
274    if (v == 2) found2 = true;
275  }
276  KJ_EXPECT(found1);
277  KJ_EXPECT(found2);
278}
279 
280KJ_TEST("SmallSet: forEach from empty state") {
281  SmallSet<kj::Rc<WeakRef<TestItem>>> set;
282  size_t count = 0;
283  set.forEach([&](TestItem& ref) { count++; });
284  KJ_EXPECT(count == 0);
285}
286 
287KJ_TEST("SmallSet: forEach skips invalidated WeakRefs") {
288  // Test that forEach properly skips items whose WeakRefs have been invalidated
289  SmallSet<kj::Rc<WeakRef<TestItem>>> set;
290 
291  auto item1 = kj::rc<TestItem>(1);
292  auto item2 = kj::rc<TestItem>(2);
293 
294  set.add(item1->getWeakRef());
295  set.add(item2->getWeakRef());
296  KJ_EXPECT(set.size() == 2);
297 
298  // Invalidate item1's weak ref by destroying item1
299  item1 = nullptr;
300 
301  // forEach should only visit item2
302  kj::Vector<int> foundValues;
303  set.forEach([&](TestItem& ref) { foundValues.add(ref.value); });
304 
305  KJ_EXPECT(foundValues.size() == 1);
306  KJ_EXPECT(foundValues[0] == 2);
307}
308 
309KJ_TEST("SmallSet: reference counting works correctly") {
310  // Verify that items are properly reference counted
311  kj::Rc<TestItem> item = kj::rc<TestItem>(42);
312  KJ_EXPECT(!item->isShared());  // Only one reference
313 
314  {
315    SmallSet<kj::Rc<TestItem>> set;
316    set.add(item.addRef());
317    KJ_EXPECT(item->isShared());  // Now shared between item and set
318  }
319  // Set destroyed, only our original reference remains
320  KJ_EXPECT(!item->isShared());
321}
322 
323}  // namespace
324}  // namespace workerd