1 #include "test_helpers.h"
2 #include "xray_segmented_array.h"
3 #include "gmock/gmock.h"
4 #include "gtest/gtest.h"
5 #include <algorithm>
6 #include <numeric>
7 #include <vector>
8 
9 namespace __xray {
10 namespace {
11 
12 using ::testing::SizeIs;
13 
14 struct TestData {
15   s64 First;
16   s64 Second;
17 
18   // Need a constructor for emplace operations.
19   TestData(s64 F, s64 S) : First(F), Second(S) {}
20 };
21 
22 void PrintTo(const TestData &D, std::ostream *OS) {
23   *OS << "{ " << D.First << ", " << D.Second << " }";
24 }
25 
26 TEST(SegmentedArrayTest, ConstructWithAllocators) {
27   using AllocatorType = typename Array<TestData>::AllocatorType;
28   AllocatorType A(1 << 4);
29   Array<TestData> Data(A);
30   (void)Data;
31 }
32 
33 TEST(SegmentedArrayTest, ConstructAndPopulate) {
34   using AllocatorType = typename Array<TestData>::AllocatorType;
35   AllocatorType A(1 << 4);
36   Array<TestData> data(A);
37   ASSERT_NE(data.Append(TestData{0, 0}), nullptr);
38   ASSERT_NE(data.Append(TestData{1, 1}), nullptr);
39   ASSERT_EQ(data.size(), 2u);
40 }
41 
42 TEST(SegmentedArrayTest, ConstructPopulateAndLookup) {
43   using AllocatorType = typename Array<TestData>::AllocatorType;
44   AllocatorType A(1 << 4);
45   Array<TestData> data(A);
46   ASSERT_NE(data.Append(TestData{0, 1}), nullptr);
47   ASSERT_EQ(data.size(), 1u);
48   ASSERT_EQ(data[0].First, 0);
49   ASSERT_EQ(data[0].Second, 1);
50 }
51 
52 TEST(SegmentedArrayTest, PopulateWithMoreElements) {
53   using AllocatorType = typename Array<TestData>::AllocatorType;
54   AllocatorType A(1 << 24);
55   Array<TestData> data(A);
56   static const auto kMaxElements = 100u;
57   for (auto I = 0u; I < kMaxElements; ++I) {
58     ASSERT_NE(data.Append(TestData{I, I + 1}), nullptr);
59   }
60   ASSERT_EQ(data.size(), kMaxElements);
61   for (auto I = 0u; I < kMaxElements; ++I) {
62     ASSERT_EQ(data[I].First, I);
63     ASSERT_EQ(data[I].Second, I + 1);
64   }
65 }
66 
67 TEST(SegmentedArrayTest, AppendEmplace) {
68   using AllocatorType = typename Array<TestData>::AllocatorType;
69   AllocatorType A(1 << 4);
70   Array<TestData> data(A);
71   ASSERT_NE(data.AppendEmplace(1, 1), nullptr);
72   ASSERT_EQ(data[0].First, 1);
73   ASSERT_EQ(data[0].Second, 1);
74 }
75 
76 TEST(SegmentedArrayTest, AppendAndTrim) {
77   using AllocatorType = typename Array<TestData>::AllocatorType;
78   AllocatorType A(1 << 4);
79   Array<TestData> data(A);
80   ASSERT_NE(data.AppendEmplace(1, 1), nullptr);
81   ASSERT_EQ(data.size(), 1u);
82   data.trim(1);
83   ASSERT_EQ(data.size(), 0u);
84   ASSERT_TRUE(data.empty());
85 }
86 
87 TEST(SegmentedArrayTest, IteratorAdvance) {
88   using AllocatorType = typename Array<TestData>::AllocatorType;
89   AllocatorType A(1 << 4);
90   Array<TestData> data(A);
91   ASSERT_TRUE(data.empty());
92   ASSERT_EQ(data.begin(), data.end());
93   auto I0 = data.begin();
94   ASSERT_EQ(I0++, data.begin());
95   ASSERT_NE(I0, data.begin());
96   for (const auto &D : data) {
97     (void)D;
98     FAIL();
99   }
100   ASSERT_NE(data.AppendEmplace(1, 1), nullptr);
101   ASSERT_EQ(data.size(), 1u);
102   ASSERT_NE(data.begin(), data.end());
103   auto &D0 = *data.begin();
104   ASSERT_EQ(D0.First, 1);
105   ASSERT_EQ(D0.Second, 1);
106 }
107 
108 TEST(SegmentedArrayTest, IteratorRetreat) {
109   using AllocatorType = typename Array<TestData>::AllocatorType;
110   AllocatorType A(1 << 4);
111   Array<TestData> data(A);
112   ASSERT_TRUE(data.empty());
113   ASSERT_EQ(data.begin(), data.end());
114   ASSERT_NE(data.AppendEmplace(1, 1), nullptr);
115   ASSERT_EQ(data.size(), 1u);
116   ASSERT_NE(data.begin(), data.end());
117   auto &D0 = *data.begin();
118   ASSERT_EQ(D0.First, 1);
119   ASSERT_EQ(D0.Second, 1);
120 
121   auto I0 = data.end();
122   ASSERT_EQ(I0--, data.end());
123   ASSERT_NE(I0, data.end());
124   ASSERT_EQ(I0, data.begin());
125   ASSERT_EQ(I0->First, 1);
126   ASSERT_EQ(I0->Second, 1);
127 }
128 
129 TEST(SegmentedArrayTest, IteratorTrimBehaviour) {
130   using AllocatorType = typename Array<TestData>::AllocatorType;
131   AllocatorType A(1 << 20);
132   Array<TestData> Data(A);
133   ASSERT_TRUE(Data.empty());
134   auto I0Begin = Data.begin(), I0End = Data.end();
135   // Add enough elements in Data to have more than one chunk.
136   constexpr auto Segment = Array<TestData>::SegmentSize;
137   constexpr auto SegmentX2 = Segment * 2;
138   for (auto i = SegmentX2; i > 0u; --i) {
139     Data.AppendEmplace(static_cast<s64>(i), static_cast<s64>(i));
140   }
141   ASSERT_EQ(Data.size(), SegmentX2);
142   {
143     auto &Back = Data.back();
144     ASSERT_EQ(Back.First, 1);
145     ASSERT_EQ(Back.Second, 1);
146   }
147 
148   // Trim one chunk's elements worth.
149   Data.trim(Segment);
150   ASSERT_EQ(Data.size(), Segment);
151 
152   // Check that we are still able to access 'back' properly.
153   {
154     auto &Back = Data.back();
155     ASSERT_EQ(Back.First, static_cast<s64>(Segment + 1));
156     ASSERT_EQ(Back.Second, static_cast<s64>(Segment + 1));
157   }
158 
159   // Then trim until it's empty.
160   Data.trim(Segment);
161   ASSERT_TRUE(Data.empty());
162 
163   // Here our iterators should be the same.
164   auto I1Begin = Data.begin(), I1End = Data.end();
165   EXPECT_EQ(I0Begin, I1Begin);
166   EXPECT_EQ(I0End, I1End);
167 
168   // Then we ensure that adding elements back works just fine.
169   for (auto i = SegmentX2; i > 0u; --i) {
170     Data.AppendEmplace(static_cast<s64>(i), static_cast<s64>(i));
171   }
172   EXPECT_EQ(Data.size(), SegmentX2);
173 }
174 
175 TEST(SegmentedArrayTest, HandleExhaustedAllocator) {
176   using AllocatorType = typename Array<TestData>::AllocatorType;
177   constexpr auto Segment = Array<TestData>::SegmentSize;
178   constexpr auto MaxElements = Array<TestData>::ElementsPerSegment;
179   AllocatorType A(Segment);
180   Array<TestData> Data(A);
181   for (auto i = MaxElements; i > 0u; --i)
182     EXPECT_NE(Data.AppendEmplace(static_cast<s64>(i), static_cast<s64>(i)),
183               nullptr);
184   EXPECT_EQ(Data.AppendEmplace(0, 0), nullptr);
185   EXPECT_THAT(Data, SizeIs(MaxElements));
186 
187   // Trimming more elements than there are in the container should be fine.
188   Data.trim(MaxElements + 1);
189   EXPECT_THAT(Data, SizeIs(0u));
190 }
191 
192 struct ShadowStackEntry {
193   uint64_t EntryTSC = 0;
194   uint64_t *NodePtr = nullptr;
195   ShadowStackEntry(uint64_t T, uint64_t *N) : EntryTSC(T), NodePtr(N) {}
196 };
197 
198 TEST(SegmentedArrayTest, SimulateStackBehaviour) {
199   using AllocatorType = typename Array<ShadowStackEntry>::AllocatorType;
200   AllocatorType A(1 << 10);
201   Array<ShadowStackEntry> Data(A);
202   static uint64_t Dummy = 0;
203   constexpr uint64_t Max = 9;
204 
205   for (uint64_t i = 0; i < Max; ++i) {
206     auto P = Data.Append({i, &Dummy});
207     ASSERT_NE(P, nullptr);
208     ASSERT_EQ(P->NodePtr, &Dummy);
209     auto &Back = Data.back();
210     ASSERT_EQ(Back.NodePtr, &Dummy);
211     ASSERT_EQ(Back.EntryTSC, i);
212   }
213 
214   // Simulate a stack by checking the data from the end as we're trimming.
215   auto Counter = Max;
216   ASSERT_EQ(Data.size(), size_t(Max));
217   while (!Data.empty()) {
218     const auto &Top = Data.back();
219     uint64_t *TopNode = Top.NodePtr;
220     EXPECT_EQ(TopNode, &Dummy) << "Counter = " << Counter;
221     Data.trim(1);
222     --Counter;
223     ASSERT_EQ(Data.size(), size_t(Counter));
224   }
225 }
226 
227 TEST(SegmentedArrayTest, PlacementNewOnAlignedStorage) {
228   using AllocatorType = typename Array<ShadowStackEntry>::AllocatorType;
229   typename std::aligned_storage<sizeof(AllocatorType),
230                                 alignof(AllocatorType)>::type AllocatorStorage;
231   new (&AllocatorStorage) AllocatorType(1 << 10);
232   auto *A = reinterpret_cast<AllocatorType *>(&AllocatorStorage);
233   typename std::aligned_storage<sizeof(Array<ShadowStackEntry>),
234                                 alignof(Array<ShadowStackEntry>)>::type
235       ArrayStorage;
236   new (&ArrayStorage) Array<ShadowStackEntry>(*A);
237   auto *Data = reinterpret_cast<Array<ShadowStackEntry> *>(&ArrayStorage);
238 
239   static uint64_t Dummy = 0;
240   constexpr uint64_t Max = 9;
241 
242   for (uint64_t i = 0; i < Max; ++i) {
243     auto P = Data->Append({i, &Dummy});
244     ASSERT_NE(P, nullptr);
245     ASSERT_EQ(P->NodePtr, &Dummy);
246     auto &Back = Data->back();
247     ASSERT_EQ(Back.NodePtr, &Dummy);
248     ASSERT_EQ(Back.EntryTSC, i);
249   }
250 
251   // Simulate a stack by checking the data from the end as we're trimming.
252   auto Counter = Max;
253   ASSERT_EQ(Data->size(), size_t(Max));
254   while (!Data->empty()) {
255     const auto &Top = Data->back();
256     uint64_t *TopNode = Top.NodePtr;
257     EXPECT_EQ(TopNode, &Dummy) << "Counter = " << Counter;
258     Data->trim(1);
259     --Counter;
260     ASSERT_EQ(Data->size(), size_t(Counter));
261   }
262 
263   // Once the stack is exhausted, we re-use the storage.
264   for (uint64_t i = 0; i < Max; ++i) {
265     auto P = Data->Append({i, &Dummy});
266     ASSERT_NE(P, nullptr);
267     ASSERT_EQ(P->NodePtr, &Dummy);
268     auto &Back = Data->back();
269     ASSERT_EQ(Back.NodePtr, &Dummy);
270     ASSERT_EQ(Back.EntryTSC, i);
271   }
272 
273   // We re-initialize the storage, by calling the destructor and
274   // placement-new'ing again.
275   Data->~Array();
276   A->~AllocatorType();
277   new (A) AllocatorType(1 << 10);
278   new (Data) Array<ShadowStackEntry>(*A);
279 
280   // Then re-do the test.
281   for (uint64_t i = 0; i < Max; ++i) {
282     auto P = Data->Append({i, &Dummy});
283     ASSERT_NE(P, nullptr);
284     ASSERT_EQ(P->NodePtr, &Dummy);
285     auto &Back = Data->back();
286     ASSERT_EQ(Back.NodePtr, &Dummy);
287     ASSERT_EQ(Back.EntryTSC, i);
288   }
289 
290   // Simulate a stack by checking the data from the end as we're trimming.
291   Counter = Max;
292   ASSERT_EQ(Data->size(), size_t(Max));
293   while (!Data->empty()) {
294     const auto &Top = Data->back();
295     uint64_t *TopNode = Top.NodePtr;
296     EXPECT_EQ(TopNode, &Dummy) << "Counter = " << Counter;
297     Data->trim(1);
298     --Counter;
299     ASSERT_EQ(Data->size(), size_t(Counter));
300   }
301 
302   // Once the stack is exhausted, we re-use the storage.
303   for (uint64_t i = 0; i < Max; ++i) {
304     auto P = Data->Append({i, &Dummy});
305     ASSERT_NE(P, nullptr);
306     ASSERT_EQ(P->NodePtr, &Dummy);
307     auto &Back = Data->back();
308     ASSERT_EQ(Back.NodePtr, &Dummy);
309     ASSERT_EQ(Back.EntryTSC, i);
310   }
311 }
312 
313 TEST(SegmentedArrayTest, ArrayOfPointersIteratorAccess) {
314   using PtrArray = Array<int *>;
315   PtrArray::AllocatorType Alloc(16384);
316   Array<int *> A(Alloc);
317   static constexpr size_t Count = 100;
318   std::vector<int> Integers(Count);
319   std::iota(Integers.begin(), Integers.end(), 0);
320   for (auto &I : Integers)
321     ASSERT_NE(A.Append(&I), nullptr);
322   int V = 0;
323   ASSERT_EQ(A.size(), Count);
324   for (auto P : A) {
325     ASSERT_NE(P, nullptr);
326     ASSERT_EQ(*P, V++);
327   }
328 }
329 
330 TEST(SegmentedArrayTest, ArrayOfPointersIteratorAccessExhaustion) {
331   using PtrArray = Array<int *>;
332   PtrArray::AllocatorType Alloc(4096);
333   Array<int *> A(Alloc);
334   static constexpr size_t Count = 1000;
335   std::vector<int> Integers(Count);
336   std::iota(Integers.begin(), Integers.end(), 0);
337   for (auto &I : Integers)
338     if (A.Append(&I) == nullptr)
339       break;
340   int V = 0;
341   ASSERT_LT(A.size(), Count);
342   for (auto P : A) {
343     ASSERT_NE(P, nullptr);
344     ASSERT_EQ(*P, V++);
345   }
346 }
347 
348 } // namespace
349 } // namespace __xray
350