3 files changed, 48 insertions, 5 deletions

diff --git a/compiler-rt/lib/xray/tests/unit/segmented_array_test.cc b/compiler-rt/lib/xray/tests/unit/segmented_array_test.cc
index 035674ccfaf..80991b1b97a 100644
--- a/compiler-rt/lib/xray/tests/unit/segmented_array_test.cc
+++ b/compiler-rt/lib/xray/tests/unit/segmented_array_test.cc
@@ -1,9 +1,13 @@
+#include "test_helpers.h"
 #include "xray_segmented_array.h"
+#include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace __xray {
 namespace {
 
+using ::testing::SizeIs;
+
 struct TestData {
   s64 First;
   s64 Second;
@@ -12,6 +16,10 @@ struct TestData {
   TestData(s64 F, s64 S) : First(F), Second(S) {}
 };
 
+void PrintTo(const TestData &D, std::ostream *OS) {
+  *OS << "{ " << D.First << ", " << D.Second << " }";
+}
+
 TEST(SegmentedArrayTest, ConstructWithAllocators) {
   using AllocatorType = typename Array<TestData>::AllocatorType;
   AllocatorType A(1 << 4);
@@ -161,6 +169,23 @@ TEST(SegmentedArrayTest, IteratorTrimBehaviour) {
   EXPECT_EQ(Data.size(), SegmentX2);
 }
 
+TEST(SegmentedArrayTest, HandleExhaustedAllocator) {
+  using AllocatorType = typename Array<TestData>::AllocatorType;
+  constexpr auto Segment = Array<TestData>::SegmentSize;
+  constexpr auto MaxElements = Array<TestData>::ElementsPerSegment;
+  AllocatorType A(Segment);
+  Array<TestData> Data(A);
+  for (auto i = MaxElements; i > 0u; --i)
+    EXPECT_NE(Data.AppendEmplace(static_cast<s64>(i), static_cast<s64>(i)),
+              nullptr);
+  EXPECT_EQ(Data.AppendEmplace(0, 0), nullptr);
+  EXPECT_THAT(Data, SizeIs(MaxElements));
+
+  // Trimming more elements than there are in the container should be fine.
+  Data.trim(MaxElements + 1);
+  EXPECT_THAT(Data, SizeIs(0u));
+}
+
 struct ShadowStackEntry {
   uint64_t EntryTSC = 0;
   uint64_t *NodePtr = nullptr;
diff --git a/compiler-rt/lib/xray/tests/unit/test_helpers.h b/compiler-rt/lib/xray/tests/unit/test_helpers.h
index 6eca7795e31..6c23ce413b9 100644
--- a/compiler-rt/lib/xray/tests/unit/test_helpers.h
+++ b/compiler-rt/lib/xray/tests/unit/test_helpers.h
@@ -14,8 +14,9 @@
 #define COMPILER_RT_LIB_XRAY_TESTS_TEST_HELPERS_H_
 
 #include "xray_buffer_queue.h"
-#include "llvm/XRay/XRayRecord.h"
+#include "xray_segmented_array.h"
 #include "llvm/XRay/Trace.h"
+#include "llvm/XRay/XRayRecord.h"
 #include "gmock/gmock.h"
 
 // TODO: Move these to llvm/include/Testing/XRay/...
@@ -54,6 +55,19 @@ namespace __xray {
 
 std::string serialize(BufferQueue &Buffers, int32_t Version);
 
+template <class T> void PrintTo(const Array<T> &A, std::ostream *OS) {
+  *OS << "[";
+  bool first = true;
+  for (const auto &E : A) {
+    if (!first) {
+      *OS << ", ";
+    }
+    PrintTo(E, OS);
+    first = false;
+  }
+  *OS << "]";
+}
+
 } // namespace __xray
 
 #endif // COMPILER_RT_LIB_XRAY_TESTS_TEST_HELPERS_H_
diff --git a/compiler-rt/lib/xray/xray_segmented_array.h b/compiler-rt/lib/xray/xray_segmented_array.h
index c723c7de083..42f53be1e01 100644
--- a/compiler-rt/lib/xray/xray_segmented_array.h
+++ b/compiler-rt/lib/xray/xray_segmented_array.h
@@ -78,6 +78,8 @@ public:
 
   static SegmentBase SentinelSegment;
 
+  using size_type = size_t;
+
 private:
   AllocatorType *Alloc;
   SegmentBase *Head = &SentinelSegment;
@@ -334,9 +336,8 @@ public:
     if (Elements == 0)
       return;
 
-    DCHECK_LE(Elements, Size);
-    DCHECK_GT(Size, 0);
     auto OldSize = Size;
+    Elements = Elements >= Size ? Size : Elements;
     Size -= Elements;
 
     DCHECK_NE(Head, &SentinelSegment);
@@ -346,8 +347,11 @@ public:
                                 nearest_boundary(Size, ElementsPerSegment)) /
                                ElementsPerSegment;
          SegmentsToTrim > 0; --SegmentsToTrim) {
-      DCHECK_NE(Head, &SentinelSegment);
-      DCHECK_NE(Tail, &SentinelSegment);
+
+      // We want to short-circuit if the trace is already empty.
+      if (Head == &SentinelSegment && Head == Tail)
+        return;
+
       // Put the tail into the Freelist.
       auto *FreeSegment = Tail;
       Tail = Tail->Prev;