[XRay][compiler-rt] xray::Array Freelist and Iterator Updates

[dberris]

Summary

We found a bug while working on a benchmark for the profiling mode which
manifests as a segmentation fault in the profiling handler's
implementation. This change adds unit tests which replicate the
issues in isolation.

We've tracked this down as a bug in the implementation of the Freelist
in the xray::Array type. This happens when we trim the array by a
number of elements, where we've been incorrectly assigning pointers for
the links in the freelist of chunk nodes. We've taken the chance to add
more debug-only assertions to the code path and allow us to verify these
assumptions in debug builds.

In the process, we also took the opportunity to use iterators to
implement both front() and back() which exposes a bug in the
iterator decrement operation. In particular, when we decrement past a
chunk size boundary, we end up moving too far back and reaching the
SentinelChunk prematurely.

This change unblocks us to allow for contributing the non-crashing
version of the benchmarks in the test-suite as well.

[dberris]

• fixup: Add a test to an additional edge-case; propagate extents into iterator

[dberris]

• Rebase to master.
  • fixup: Add a test to an additional edge-case; propagate extents into iterator
  • fixup: play with alignment, flags, and testing properly
  • fixup: fix bug in trim implementation in isolation
  • fixup: all changes to fix benchmark run

[dberris]

This change is now bigger than the original posted, and I'm going to spend a bit more time cleaning it up. In particular I intend to:

• Explore more edge cases in the unit tests for the segmented array.

• Document and test the assumptions for the allocator implementation.

• Reduce the memory requirements by changing the sizing algorithm used for the chunks in the segmented array.

• Reduce the size of the segmented array by attempting to remove the 'FreeElements' member.

Let's determine whether I can split it up further to make it easier to review when all the changes are done (if it's even possible to break it up).

[dberris]

In D48653#1150774, @dberris wrote:

This change is now bigger than the original posted, and I'm going to spend a bit more time cleaning it up. In particular I intend to:

• Explore more edge cases in the unit tests for the segmented array.

This didn't seem to be blocking, so I ended up not doing it.

• Document and test the assumptions for the allocator implementation.

It also looked like this could be done after the fact.

• Reduce the memory requirements by changing the sizing algorithm used for the chunks in the segmented array.

I tried going down this route and it was clear that we needed to get the alignment correct, so being conservative here for correctness seems fine for now.

• Reduce the size of the segmented array by attempting to remove the 'FreeElements' member.

I've done this, which simplifies the code a lot.

Let's determine whether I can split it up further to make it easier to review when all the changes are done (if it's even possible to break it up).

I thought about this more and it seems easier to get this in as a batch instead of trying to do this in smaller batches.

[dberris]

• fixup: clean up some more, now ready for a look.

[dberris]

Update description, rebase on top of bug-fix change in profiling mode.

[dberris]

Reverting change, restoring "good" state.

[kpw]

Seems like half of this change is DCHECKS. Much better coverage! Was that just a result of your strategy for diagnosing what actually caused the problems and leaving those in?

compiler-rt/lib/xray/tests/unit/function_call_trie_test.cc

183

No EXPECT statement for the time tracking?

compiler-rt/lib/xray/tests/unit/segmented_array_test.cc

129–139

You can replace "ChunkX2 / 2" with the equivalent "Chunk" here and below.

150–186

Why Append here and AppendEmplace above?

compiler-rt/lib/xray/xray_allocator.h

153

So trim of segmented array was free-ing the "block" with the wrong allocator when trimming the first element?

compiler-rt/lib/xray/xray_function_call_trie.h

371

Should the NewRoots point back to the old roots as a parent? Don't we want to keep this isolated?

compiler-rt/lib/xray/xray_segmented_array.h

44–45

I assume this is all alignment related?

172

This size check is applicable when decrementing end()?

205–206

Lol. Checks out.

294–296

This is common enough you might want to put it into an always inlined private member function. Maybe currentElementData()

308

I would comment operator-- with the fact that it must not mutate the iterator. It's counterintuitive that it's non-destructive while ++ advances state.

336

This feels like it probably simplifies to an equivalent expression.

[dberris]

In D48653#1156953, @kpw wrote:

Seems like half of this change is DCHECKS. Much better coverage! Was that just a result of your strategy for diagnosing what actually caused the problems and leaving those in?

Yep! Turns out that using assertions is a great way to... well... assert things. :)

compiler-rt/lib/xray/tests/unit/segmented_array_test.cc

150–186

No good reason. :) Fixed.

compiler-rt/lib/xray/xray_allocator.h

153

Yeah, this one was hard to track. We were running into undefined behaviour in a lot of places given alignment issues with reinterpret_cast's.

compiler-rt/lib/xray/xray_function_call_trie.h

371

No, but what we're doing here is a parallel traversal of the tries -- we add the root from the original trie alongside the new root in the new new/destination trie. That way we can recreate the structure from the source trie into the new trie.

compiler-rt/lib/xray/xray_segmented_array.h

44–45

Yes, we'd like to have a power-of-two size for the structures, to allow easier pointer alignment calculations.

172

Yes -- we know that when PreviousOffset == Size and neither are 0, that the offset is now pointing to a valid object.

205–206

Yeah, ensuring correct alignment is hard. :P

308

operator-- actually does mutate the iterator. What we're doing here is being explicit about storing the state of the iterator into It first, mutating It, and then de-referencing it.

336

Actually turns out the math was slightly wrong, updated and fixed.

[dberris]

Address comments, rebase, before landing.

compiler-rt/lib/xray/tests/CMakeLists.txt

if(APPLE)
     set(tgt_name "RTXRay.test.osx")
   else()
     set(tgt_name "RTXRay.test.${arch}")
   endif()
   set(${lib} "${tgt_name}" PARENT_SCOPE)
 endfunction()
 
 set(XRAY_TEST_ARCH ${XRAY_SUPPORTED_ARCH})
-set(XRAY_UNITTEST_LINK_FLAGS ${CMAKE_THREAD_LIBS_INIT})
+set(XRAY_UNITTEST_LINK_FLAGS
+  ${CMAKE_THREAD_LIBS_INIT}
+  -fxray-instrument
+  -lstdc++  # TODO: Detect the correct standard library used!
+  )
 if (NOT APPLE)
   append_list_if(COMPILER_RT_HAS_LIBM -lm XRAY_UNITTEST_LINK_FLAGS)
   append_list_if(COMPILER_RT_HAS_LIBRT -lrt XRAY_UNITTEST_LINK_FLAGS)
   append_list_if(COMPILER_RT_HAS_LIBDL -ldl XRAY_UNITTEST_LINK_FLAGS)
   append_list_if(COMPILER_RT_HAS_LIBPTHREAD -pthread XRAY_UNITTEST_LINK_FLAGS)
 endif()
 
 macro(add_xray_unittest testname)
foreach(arch ${XRAY_TEST_ARCH})
         # Note that any change in the implementations will cause all the unit
         # tests to be re-built. This is by design, but may be cumbersome during
         # the build/test cycle.
         COMPILE_DEPS ${TEST_SOURCES} ${COMPILER_RT_GTEST_SOURCE}
         ${XRAY_HEADERS} ${XRAY_IMPL_FILES}
         RUNTIME "${XRAY_RUNTIME_LIBS}"
         DEPS gtest xray
         CFLAGS ${XRAY_UNITTEST_CFLAGS}
-        LINK_FLAGS ${TARGET_LINK_FLAGS} ${XRAY_UNITTEST_LINK_FLAGS} -lstdc++)
+        LINK_FLAGS ${TARGET_LINK_FLAGS} ${XRAY_UNITTEST_LINK_FLAGS})
       set_target_properties(XRayUnitTests
         PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
     endforeach()
   endif()
 endmacro()
 
 if(COMPILER_RT_CAN_EXECUTE_TESTS)
   if (APPLE)

compiler-rt/lib/xray/tests/unit/function_call_trie_test.cc

TEST(FunctionCallTrieTest, MissingFunctionEntry) {
   auto A = FunctionCallTrie::InitAllocators();
   FunctionCallTrie Trie(A);
   Trie.exitFunction(1, 1);
   const auto &R = Trie.getRoots();
 
   ASSERT_TRUE(R.empty());
 }
 
+TEST(FunctionCallTrieTest, NoMatchingEntersForExit) {
+  auto A = FunctionCallTrie::InitAllocators();
+  FunctionCallTrie Trie(A);
+  Trie.enterFunction(2, 1);
+  Trie.enterFunction(3, 3);
+  Trie.exitFunction(1, 5);
+  const auto &R = Trie.getRoots();
+
+  ASSERT_TRUE(R.empty());
+}
+
 TEST(FunctionCallTrieTest, MissingFunctionExit) {
   auto A = FunctionCallTrie::InitAllocators();
   FunctionCallTrie Trie(A);
   Trie.enterFunction(1, 1);
   const auto &R = Trie.getRoots();
 
   ASSERT_TRUE(R.empty());
 }
TEST(FunctionCallTrieTest, MissingIntermediaryExit) {
   EXPECT_EQ(F3.CallCount, 1);
   EXPECT_EQ(F2.CallCount, 1);
   EXPECT_EQ(F1.CallCount, 1);
   EXPECT_EQ(F3.CumulativeLocalTime, 100);
   EXPECT_EQ(F2.CumulativeLocalTime, 300);
   EXPECT_EQ(F1.CumulativeLocalTime, 100);
 }
 
+TEST(FunctionCallTrieTest, DeepCallStack) {
+  // Simulate a relatively deep call stack (32 levels) and ensure that we can
+  // properly pop all the way up the stack.
+  profilingFlags()->setDefaults();
+  auto A = FunctionCallTrie::InitAllocators();
+  FunctionCallTrie Trie(A);
+  for (int i = 0; i < 32; ++i)
+    Trie.enterFunction(i + 1, i);
+  Trie.exitFunction(1, 33);
+
+  // Here, validate that we have a 32-level deep function call path from the
+  // root (1) down to the leaf (33).
+  const auto &R = Trie.getRoots();
+  ASSERT_EQ(R.size(), 1u);
+  auto F = R[0];
+  for (int i = 0; i < 32; ++i) {
+    EXPECT_EQ(F->FId, i + 1);
+    if (F->Callees.empty() && i != 31)
+      FAIL() << "Empty callees for FId " << F->FId;
+    if (i != 31)
+      F = F->Callees[0].NodePtr;
+  }
+}
+
 // TODO: Test that we can handle cross-CPU migrations, where TSCs are not
 // guaranteed to be synchronised.
 TEST(FunctionCallTrieTest, DeepCopy) {
   profilingFlags()->setDefaults();
   auto A = FunctionCallTrie::InitAllocators();
   FunctionCallTrie Trie(A);
 
   Trie.enterFunction(1, 0);

compiler-rt/lib/xray/tests/unit/segmented_array_test.cc

TEST(SegmentedArrayTest, IteratorRetreat) {
   ASSERT_EQ(I0--, data.end());
   ASSERT_NE(I0, data.end());
   ASSERT_EQ(I0, data.begin());
   ASSERT_EQ(I0->First, 1);
   ASSERT_EQ(I0->Second, 1);
 }
 
 TEST(SegmentedArrayTest, IteratorTrimBehaviour) {
-  Array<TestData> data;
-  ASSERT_TRUE(data.empty());
-  auto I0Begin = data.begin(), I0End = data.end();
-  // Add enough elements in data to have more than one chunk.
-  constexpr auto ChunkX2 = Array<TestData>::ChunkSize * 2;
+  using AllocatorType = typename Array<TestData>::AllocatorType;
+  AllocatorType A(1 << 10, 0);
+  Array<TestData> Data(A);
+  ASSERT_TRUE(Data.empty());
+  auto I0Begin = Data.begin(), I0End = Data.end();
+  // Add enough elements in Data to have more than one chunk.
+  constexpr auto Chunk = Array<TestData>::ChunkSize;
+  constexpr auto ChunkX2 = Chunk * 2;
   for (auto i = ChunkX2; i > 0u; --i) {
-    data.Append({static_cast<s64>(i), static_cast<s64>(i)});
+    Data.AppendEmplace(static_cast<s64>(i), static_cast<s64>(i));
   }
-  ASSERT_EQ(data.size(), ChunkX2);
+  ASSERT_EQ(Data.size(), ChunkX2);
+  {
+    auto &Back = Data.back();
+    ASSERT_EQ(Back.First, 1);
+    ASSERT_EQ(Back.Second, 1);
+  }
+
   // Trim one chunk's elements worth.
-  data.trim(ChunkX2 / 2);
-  ASSERT_EQ(data.size(), ChunkX2 / 2);
+  Data.trim(ChunkX2 / 2);
+  ASSERT_EQ(Data.size(), ChunkX2 / 2);
+
+  // Check that we are still able to access 'back' properly.
+  {
+    auto &Back = Data.back();
+    ASSERT_EQ(Back.First, static_cast<s64>(Chunk + 1));
+    ASSERT_EQ(Back.Second, static_cast<s64>(Chunk + 1));
+  }
+
   // Then trim until it's empty.
-  data.trim(ChunkX2 / 2);
-  ASSERT_TRUE(data.empty());
+  Data.trim(ChunkX2 / 2);
+  ASSERT_TRUE(Data.empty());
 
   // Here our iterators should be the same.
-  auto I1Begin = data.begin(), I1End = data.end();
+  auto I1Begin = Data.begin(), I1End = Data.end();
   EXPECT_EQ(I0Begin, I1Begin);
   EXPECT_EQ(I0End, I1End);
 
   // Then we ensure that adding elements back works just fine.
   for (auto i = ChunkX2; i > 0u; --i) {
-    data.Append({static_cast<s64>(i), static_cast<s64>(i)});
+    Data.Append({static_cast<s64>(i), static_cast<s64>(i)});
+  }
+  EXPECT_EQ(Data.size(), ChunkX2);
+}
+
+struct ShadowStackEntry {
+  uint64_t EntryTSC = 0;
+  uint64_t *NodePtr = nullptr;
+  ShadowStackEntry(uint64_t T, uint64_t *N) : EntryTSC(T), NodePtr(N) {}
+};
+
+TEST(SegmentedArrayTest, SimulateStackBehaviour) {
+  using AllocatorType = typename Array<ShadowStackEntry>::AllocatorType;
+  AllocatorType A(1 << 10, 0);
+  Array<ShadowStackEntry> Data(A);
+  static uint64_t Dummy = 0;
+  constexpr uint64_t Max = 9;
+
+  for (uint64_t i = 0; i < Max; ++i) {
+    auto P = Data.Append({i, &Dummy});
+    ASSERT_NE(P, nullptr);
+    ASSERT_EQ(P->NodePtr, &Dummy);
+    auto &Back = Data.back();
+    ASSERT_EQ(Back.NodePtr, &Dummy);
+    ASSERT_EQ(Back.EntryTSC, i);
+  }
+
+  // Simulate a stack by checking the data from the end as we're trimming.
+  auto Counter = Max;
+  ASSERT_EQ(Data.size(), size_t(Max));
+  while (!Data.empty()) {
+    const auto &Top = Data.back();
+    uint64_t *TopNode = Top.NodePtr;
+    EXPECT_EQ(TopNode, &Dummy) << "Counter = " << Counter;
+    Data.trim(1);
+    --Counter;
+    ASSERT_EQ(Data.size(), size_t(Counter));
   }
-  EXPECT_EQ(data.size(), ChunkX2);
 }
 
 } // namespace
 } // namespace __xray

compiler-rt/lib/xray/xray_allocator.h

 // the profiling runtime.
 //
 //===----------------------------------------------------------------------===//
 #ifndef XRAY_ALLOCATOR_H
 #define XRAY_ALLOCATOR_H
 
 #include "sanitizer_common/sanitizer_allocator_internal.h"
 #include "sanitizer_common/sanitizer_common.h"
+#include "sanitizer_common/sanitizer_internal_defs.h"
 #include "sanitizer_common/sanitizer_mutex.h"
+#include "xray_utils.h"
 #include <cstddef>
 #include <cstdint>
 
-#include "sanitizer_common/sanitizer_internal_defs.h"
-
 namespace __xray {
 
 /// The Allocator type hands out fixed-sized chunks of memory that are
 /// cache-line aligned and sized. This is useful for placement of
 /// performance-sensitive data in memory that's frequently accessed. The
 /// allocator also self-limits the peak memory usage to a dynamically defined
 /// maximum.
 ///
 /// N is the lower-bound size of the block of memory to return from the
 /// allocation function. N is used to compute the size of a block, which is
 /// cache-line-size multiples worth of memory. We compute the size of a block by
 /// determining how many cache lines worth of memory is required to subsume N.
 template <size_t N> struct Allocator {
   // The Allocator returns memory as Block instances.
   struct Block {
     /// Compute the minimum cache-line size multiple that is >= N.
-    static constexpr auto Size =
-        kCacheLineSize * ((N / kCacheLineSize) + (N % kCacheLineSize ? 1 : 0));
+    static constexpr auto Size = nearest_boundary(N, kCacheLineSize);
     void *Data = nullptr;
   };
 
 private:
   // A BlockLink will contain a fixed number of blocks, each with an identifier
   // to specify whether it's been handed out or not. We keep track of BlockLink
   // iterators, which are basically a pointer to the link and an offset into
   // the fixed set of blocks associated with a link. The iterators are
   // bidirectional.
   //
   // We're calling it a "link" in the context of seeing these as a chain of
   // block pointer containers (i.e. links in a chain).
   struct BlockLink {
     static_assert(kCacheLineSize % sizeof(void *) == 0,
                   "Cache line size is not divisible by size of void*; none of "
                   "the assumptions of the BlockLink will hold.");
 
     // We compute the number of pointers to areas in memory where we consider as
     // individual blocks we've allocated. To ensure that instances of the
-    // BlockLink object are cache-line sized, we deduct one additional
-    // pointers worth representing the pointer to the previous link.
+    // BlockLink object are cache-line sized, we deduct two pointers worth
+    // representing the pointer to the previous link and the backing store for
+    // the whole block.
     //
     // This structure corresponds to the following layout:
     //
-    //   Blocks [ 0, 1, 2, .., BlockPtrCount - 1]
+    //   Blocks [ 0, 1, 2, .., BlockPtrCount - 2]
     //
     static constexpr auto BlockPtrCount =
-        (kCacheLineSize / sizeof(Block *)) - 1;
+        (kCacheLineSize / sizeof(Block *)) - 2;
 
     BlockLink() {
       // Zero out Blocks.
       // FIXME: Use a braced member initializer when we drop support for GCC
       // 4.8.
       internal_memset(Blocks, 0, sizeof(Blocks));
     }
 
     // FIXME: Align this to cache-line address boundaries?
     Block Blocks[BlockPtrCount];
     BlockLink *Prev = nullptr;
+    void *BackingStore = nullptr;
   };
 
   static_assert(sizeof(BlockLink) == kCacheLineSize,
                 "BlockLink instances must be cache-line-sized.");
 
   static BlockLink NullLink;
 
   // FIXME: Implement a freelist, in case we actually do intend to return memory
   // to the allocator, as opposed to just de-allocating everything in one go?
 
   size_t MaxMemory;
   SpinMutex Mutex{};
   BlockLink *Tail = &NullLink;
   size_t Counter = 0;
 
-  BlockLink *NewChainLink() {
+  BlockLink *NewChainLink(uint64_t Alignment) {
     auto NewChain = reinterpret_cast<BlockLink *>(
         InternalAlloc(sizeof(BlockLink), nullptr, kCacheLineSize));
     auto BackingStore = reinterpret_cast<char *>(InternalAlloc(
-        BlockLink::BlockPtrCount * Block::Size, nullptr, kCacheLineSize));
+        (BlockLink::BlockPtrCount + 1) * Block::Size, nullptr, Alignment));
     size_t Offset = 0;
     DCHECK_NE(NewChain, nullptr);
     DCHECK_NE(BackingStore, nullptr);
+    NewChain->BackingStore = BackingStore;
+
+    // Here we ensure that the alignment of the pointers we're handing out
+    // adhere to the alignment requirements of the call to Allocate().
     for (auto &B : NewChain->Blocks) {
-      B.Data = BackingStore + Offset;
-      Offset += Block::Size;
+      auto AlignmentAdjustment =
+          nearest_boundary(reinterpret_cast<uintptr_t>(BackingStore + Offset),
+                           Alignment) -
+          reinterpret_cast<uintptr_t>(BackingStore + Offset);
+      B.Data = BackingStore + AlignmentAdjustment + Offset;
+      DCHECK_EQ(reinterpret_cast<uintptr_t>(B.Data) % Alignment, 0);
+      Offset += AlignmentAdjustment + Block::Size;
     }
     NewChain->Prev = Tail;
     return NewChain;
   }
 
 public:
   Allocator(size_t M, size_t PreAllocate) : MaxMemory(M) {
     // FIXME: Implement PreAllocate support!
   }
 
-  Block Allocate() {
+  Block Allocate(uint64_t Alignment = 8) {
     SpinMutexLock Lock(&Mutex);
     // Check whether we're over quota.
     if (Counter * Block::Size >= MaxMemory)
       return {};
 
     size_t ChainOffset = Counter % BlockLink::BlockPtrCount;
 
     Block B{};
     BlockLink *Link = Tail;
     if (UNLIKELY(Counter == 0 || ChainOffset == 0))
-      Tail = Link = NewChainLink();
+      Tail = Link = NewChainLink(Alignment);
 
     B = Link->Blocks[ChainOffset];
     ++Counter;
     return B;
   }
 
   ~Allocator() NOEXCEPT {
     // We need to deallocate all the blocks, including the chain links.
     for (auto *C = Tail; C != &NullLink;) {
       // We know that the data block is a large contiguous page, we deallocate
       // that at once.
-      InternalFree(C->Blocks[0].Data);
+      InternalFree(C->BackingStore);
       auto Prev = C->Prev;
       InternalFree(C);
       C = Prev;
     }
   }
 }; // namespace __xray
 
 // Storage for the NullLink sentinel.
 template <size_t N> typename Allocator<N>::BlockLink Allocator<N>::NullLink;
 
 } // namespace __xray
 
 #endif // XRAY_ALLOCATOR_H

compiler-rt/lib/xray/xray_function_call_trie.h

 // This file defines the interface for a function call trie.
 //
 //===----------------------------------------------------------------------===//
 #ifndef XRAY_FUNCTION_CALL_TRIE_H
 #define XRAY_FUNCTION_CALL_TRIE_H
 
 #include "xray_profiling_flags.h"
 #include "xray_segmented_array.h"
-#include <utility>
-#include <memory>  // For placement new.
+#include <memory> // For placement new.
+#include <utility>
 
 namespace __xray {
 
 /// A FunctionCallTrie represents the stack traces of XRay instrumented
 /// functions that we've encountered, where a node corresponds to a function and
 /// the path from the root to the node its stack trace. Each node in the trie
 /// will contain some useful values, including:
 ///
public:
   struct NodeIdPair {
     Node *NodePtr;
     int32_t FId;
 
     // Constructor for inplace-construction.
     NodeIdPair(Node *N, int32_t F) : NodePtr(N), FId(F) {}
   };
 
+  static_assert(sizeof(NodeIdPair) == 16, "");
+
   using NodeIdPairArray = Array<NodeIdPair>;
   using NodeIdPairAllocatorType = NodeIdPairArray::AllocatorType;
 
   // A Node in the FunctionCallTrie gives us a list of callees, the cumulative
   // number of times this node actually appeared, the cumulative amount of time
   // for this particular node including its children call times, and just the
   // local time spent on this node. Each Node will have the ID of the XRay
   // instrumented function that it is associated to.
Node(Node *P, NodeIdPairAllocatorType &A, int64_t CC, int64_t CLT,
         : Parent(P), Callees(A), CallCount(CC), CumulativeLocalTime(CLT),
           FId(F) {}
 
     // TODO: Include the compact histogram.
   };
 
 private:
   struct ShadowStackEntry {
-    int32_t FId; // We're copying the function ID into the stack to avoid having
-                 // to reach into the node just to get the function ID.
     uint64_t EntryTSC;
     Node *NodePtr;
+    int32_t FuncId;
 
     // We add a constructor here to allow us to inplace-construct through
     // Array<...>'s AppendEmplace.
-    ShadowStackEntry(int32_t F, uint64_t T, Node *N)
-        : FId(F), EntryTSC(T), NodePtr(N) {}
+    ShadowStackEntry(uint64_t T, Node *N, int32_t F)
+        : EntryTSC{T}, NodePtr{N}, FuncId{F} {}
   };
 
   using NodeArray = Array<Node>;
   using RootArray = Array<Node *>;
   using ShadowStackArray = Array<ShadowStackEntry>;
 
 public:
   // We collate the allocators we need into a single struct, as a convenience to
~Allocators() {
         NodeIdPairAllocator->~NodeIdPairAllocatorType();
         InternalFree(NodeIdPairAllocator);
       }
     }
   };
 
   // TODO: Support configuration of options through the arguments.
   static Allocators InitAllocators() {
+    return InitAllocatorsCustom(profilingFlags()->per_thread_allocator_max);
+  }
+
+  static Allocators InitAllocatorsCustom(uptr Max) {
     Allocators A;
     auto NodeAllocator = reinterpret_cast<Allocators::NodeAllocatorType *>(
         InternalAlloc(sizeof(Allocators::NodeAllocatorType)));
-    new (NodeAllocator) Allocators::NodeAllocatorType(
-        profilingFlags()->per_thread_allocator_max, 0);
+    new (NodeAllocator) Allocators::NodeAllocatorType(Max, 0);
     A.NodeAllocator = NodeAllocator;
 
     auto RootAllocator = reinterpret_cast<Allocators::RootAllocatorType *>(
         InternalAlloc(sizeof(Allocators::RootAllocatorType)));
-    new (RootAllocator) Allocators::RootAllocatorType(
-        profilingFlags()->per_thread_allocator_max, 0);
+    new (RootAllocator) Allocators::RootAllocatorType(Max, 0);
     A.RootAllocator = RootAllocator;
 
     auto ShadowStackAllocator =
         reinterpret_cast<Allocators::ShadowStackAllocatorType *>(
             InternalAlloc(sizeof(Allocators::ShadowStackAllocatorType)));
-    new (ShadowStackAllocator) Allocators::ShadowStackAllocatorType(
-        profilingFlags()->per_thread_allocator_max, 0);
+    new (ShadowStackAllocator) Allocators::ShadowStackAllocatorType(Max, 0);
     A.ShadowStackAllocator = ShadowStackAllocator;
 
     auto NodeIdPairAllocator = reinterpret_cast<NodeIdPairAllocatorType *>(
         InternalAlloc(sizeof(NodeIdPairAllocatorType)));
-    new (NodeIdPairAllocator)
-        NodeIdPairAllocatorType(profilingFlags()->per_thread_allocator_max, 0);
+    new (NodeIdPairAllocator) NodeIdPairAllocatorType(Max, 0);
     A.NodeIdPairAllocator = NodeIdPairAllocator;
     return A;
   }
 
 private:
   NodeArray Nodes;
   RootArray Roots;
   ShadowStackArray ShadowStack;
   NodeIdPairAllocatorType *NodeIdPairAllocator = nullptr;
 
   const Allocators &GetGlobalAllocators() {
     static const Allocators A = [] { return InitAllocators(); }();
     return A;
   }
 
 public:
   explicit FunctionCallTrie(const Allocators &A)
       : Nodes(*A.NodeAllocator), Roots(*A.RootAllocator),
         ShadowStack(*A.ShadowStackAllocator),
         NodeIdPairAllocator(A.NodeIdPairAllocator) {}
 
   FunctionCallTrie() : FunctionCallTrie(GetGlobalAllocators()) {}
 
-  void enterFunction(int32_t FId, uint64_t TSC) {
+  void enterFunction(const int32_t FId, uint64_t TSC) {
+    DCHECK_NE(FId, 0);
     // This function primarily deals with ensuring that the ShadowStack is
     // consistent and ready for when an exit event is encountered.
     if (UNLIKELY(ShadowStack.empty())) {
       auto NewRoot =
           Nodes.AppendEmplace(nullptr, *NodeIdPairAllocator, 0, 0, FId);
       if (UNLIKELY(NewRoot == nullptr))
         return;
       Roots.Append(NewRoot);
-      ShadowStack.AppendEmplace(FId, TSC, NewRoot);
+      ShadowStack.AppendEmplace(TSC, NewRoot, FId);
       return;
     }
 
     auto &Top = ShadowStack.back();
     auto TopNode = Top.NodePtr;
+    DCHECK_NE(TopNode, nullptr);
 
     // If we've seen this callee before, then we just access that node and place
     // that on the top of the stack.
     auto Callee = TopNode->Callees.find_element(
         [FId](const NodeIdPair &NR) { return NR.FId == FId; });
     if (Callee != nullptr) {
       CHECK_NE(Callee->NodePtr, nullptr);
-      ShadowStack.AppendEmplace(FId, TSC, Callee->NodePtr);
+      ShadowStack.AppendEmplace(TSC, Callee->NodePtr, FId);
       return;
     }
 
     // This means we've never seen this stack before, create a new node here.
     auto NewNode =
         Nodes.AppendEmplace(TopNode, *NodeIdPairAllocator, 0, 0, FId);
     if (UNLIKELY(NewNode == nullptr))
       return;
+    DCHECK_NE(NewNode, nullptr);
     TopNode->Callees.AppendEmplace(NewNode, FId);
-    ShadowStack.AppendEmplace(FId, TSC, NewNode);
+    ShadowStack.AppendEmplace(TSC, NewNode, FId);
+    DCHECK_NE(ShadowStack.back().NodePtr, nullptr);
     return;
   }
 
   void exitFunction(int32_t FId, uint64_t TSC) {
     // When we exit a function, we look up the ShadowStack to see whether we've
     // entered this function before. We do as little processing here as we can,
     // since most of the hard work would have already been done at function
     // entry.
-    if (UNLIKELY(ShadowStack.empty()))
-      return;
-
     uint64_t CumulativeTreeTime = 0;
     while (!ShadowStack.empty()) {
-      auto &Top = ShadowStack.back();
+      const auto &Top = ShadowStack.back();
       auto TopNode = Top.NodePtr;
-      auto TopFId = TopNode->FId;
+      DCHECK_NE(TopNode, nullptr);
       auto LocalTime = TSC - Top.EntryTSC;
       TopNode->CallCount++;
       TopNode->CumulativeLocalTime += LocalTime - CumulativeTreeTime;
       CumulativeTreeTime += LocalTime;
       ShadowStack.trim(1);
 
       // TODO: Update the histogram for the node.
-      if (TopFId == FId)
+      if (Top.FuncId == FId)
         break;
     }
   }
 
   const RootArray &getRoots() const { return Roots; }
 
   // The deepCopyInto operation will update the provided FunctionCallTrie by
   // re-creating the contents of this particular FunctionCallTrie in the other
   // FunctionCallTrie. It will do this using a Depth First Traversal from the
   // roots, and while doing so recreating the traversal in the provided
   // FunctionCallTrie.
   //
   // This operation will *not* destroy the state in `O`, and thus may cause some
   // duplicate entries in `O` if it is not empty.
   //
   // This function is *not* thread-safe, and may require external
   // synchronisation of both "this" and |O|.
   //
   // This function must *not* be called with a non-empty FunctionCallTrie |O|.
   void deepCopyInto(FunctionCallTrie &O) const {
     DCHECK(O.getRoots().empty());
-    for (const auto Root : getRoots()) {
-      // Add a node in O for this root.
-      auto NewRoot = O.Nodes.AppendEmplace(
-          nullptr, *O.NodeIdPairAllocator, Root->CallCount,
-          Root->CumulativeLocalTime, Root->FId);
-      O.Roots.Append(NewRoot);
 
-      // We then push the root into a stack, to use as the parent marker for new
-      // nodes we push in as we're traversing depth-first down the call tree.
-      struct NodeAndParent {
-        FunctionCallTrie::Node *Node;
-        FunctionCallTrie::Node *NewNode;
-      };
-      using Stack = Array<NodeAndParent>;
 
-      typename Stack::AllocatorType StackAllocator(
-          profilingFlags()->stack_allocator_max, 0);
-      Stack DFSStack(StackAllocator);
 
+    for (const auto Root : getRoots()) {
+      // Add a node in O for this root.
+      auto NewRoot = O.Nodes.AppendEmplace(
+          nullptr, *O.NodeIdPairAllocator, Root->CallCount,
+          Root->CumulativeLocalTime, Root->FId);
+      O.Roots.Append(NewRoot);
+
       // TODO: Figure out what to do if we fail to allocate any more stack
       // space. Maybe warn or report once?
       DFSStack.Append(NodeAndParent{Root, NewRoot});
       while (!DFSStack.empty()) {
         NodeAndParent NP = DFSStack.back();
         DCHECK_NE(NP.Node, nullptr);
         DCHECK_NE(NP.NewNode, nullptr);
         DFSStack.trim(1);
         for (const auto Callee : NP.Node->Callees) {
           auto NewNode = O.Nodes.AppendEmplace(
               NP.NewNode, *O.NodeIdPairAllocator, Callee.NodePtr->CallCount,

compiler-rt/lib/xray/xray_profile_collector.cc

 //
 // This implements the interface for the profileCollectorService.
 //
 //===----------------------------------------------------------------------===//
 #include "xray_profile_collector.h"
 #include "sanitizer_common/sanitizer_common.h"
 #include "sanitizer_common/sanitizer_vector.h"
 #include "xray_profiling_flags.h"
-#include <pthread.h>
 #include <memory>
+#include <pthread.h>
 #include <utility>
 
 namespace __xray {
 namespace profileCollectorService {
 
 namespace {
 
 SpinMutex GlobalMutex;

 
 void post(const FunctionCallTrie &T, tid_t TId) {
   static pthread_once_t Once = PTHREAD_ONCE_INIT;
   pthread_once(&Once, +[] {
     SpinMutexLock Lock(&GlobalMutex);
     GlobalAllocators = reinterpret_cast<FunctionCallTrie::Allocators *>(
         InternalAlloc(sizeof(FunctionCallTrie::Allocators)));
     new (GlobalAllocators) FunctionCallTrie::Allocators();
-    *GlobalAllocators = FunctionCallTrie::InitAllocators();
+    *GlobalAllocators = FunctionCallTrie::InitAllocatorsCustom(
+        profilingFlags()->global_allocator_max);
   });
   DCHECK_NE(GlobalAllocators, nullptr);
 
   ThreadTrie *Item = nullptr;
   {
     SpinMutexLock Lock(&GlobalMutex);
     if (GlobalAllocators == nullptr)
       return;
ThreadTrie *Item = nullptr;
     //    trivially destructible, where FunctionCallTrie isn't.
     //
     // 2) Using a managed allocator means we need to manage that separately,
     //    which complicates the nature of this code. To get around that, we're
     //    using the internal allocator instead, which has its own global state
     //    and is decoupled from the lifetime management required by the managed
     //    allocator we have in XRay.
     //
-    Item->Trie = reinterpret_cast<FunctionCallTrie *>(
-        InternalAlloc(sizeof(FunctionCallTrie)));
+    Item->Trie = reinterpret_cast<FunctionCallTrie *>(InternalAlloc(
+        sizeof(FunctionCallTrie), nullptr, alignof(FunctionCallTrie)));
     DCHECK_NE(Item->Trie, nullptr);
     new (Item->Trie) FunctionCallTrie(*GlobalAllocators);
   }
-  DCHECK_NE(Item, nullptr);
 
   T.deepCopyInto(*Item->Trie);
 }
 
 // A PathArray represents the function id's representing a stack trace. In this
 // context a path is almost always represented from the leaf function in a call
 // stack to a root of the call trie.
 using PathArray = Array<int32_t>;

compiler-rt/lib/xray/xray_profiling.cc

if (Verbosity())
       Report("Cannot initialize already initialised profiling "
              "implementation.\n");
     return static_cast<XRayLogInitStatus>(CurrentStatus);
   }
 
   {
     SpinMutexLock Lock(&ProfilerOptionsMutex);
     FlagParser ConfigParser;
-    auto *F = profilingFlags();
-    F->setDefaults();
-    registerProfilerFlags(&ConfigParser, F);
+    ProfilerFlags Flags;
+    Flags.setDefaults();
+    registerProfilerFlags(&ConfigParser, &Flags);
     ConfigParser.ParseString(profilingCompilerDefinedFlags());
     const char *Env = GetEnv("XRAY_PROFILING_OPTIONS");
     if (Env == nullptr)
       Env = "";
     ConfigParser.ParseString(Env);
 
     // Then parse the configuration string provided.
     ConfigParser.ParseString(static_cast<const char *>(Options));
     if (Verbosity())
       ReportUnrecognizedFlags();
+    *profilingFlags() = Flags;
   }
 
   // We need to reset the profile data collection implementation now.
   profileCollectorService::reset();
 
   // We need to set up the at-thread-exit handler.
   static pthread_once_t Once = PTHREAD_ONCE_INIT;
   pthread_once(&Once, +[] {

compiler-rt/lib/xray/xray_profiling_flags.inc

 //
 //===----------------------------------------------------------------------===//
 #ifndef XRAY_FLAG
 #error "Define XRAY_FLAG prior to including this file!"
 #endif
 
 XRAY_FLAG(uptr, per_thread_allocator_max, 2 << 20,
           "Maximum size of any single per-thread allocator.")
-XRAY_FLAG(uptr, global_allocator_max, 2 << 24,
+XRAY_FLAG(uptr, global_allocator_max, 2uL << 48,
           "Maximum size of the global allocator for profile storage.")
 XRAY_FLAG(uptr, stack_allocator_max, 2 << 24,
           "Maximum size of the traversal stack allocator.")
 XRAY_FLAG(int, grace_period_ms, 100,
           "Profile collection will wait this much time in milliseconds before "
           "resetting the global state. This gives a chance to threads to "
           "notice that the profiler has been finalized and clean up.")
 XRAY_FLAG(bool, no_flush, false,
           "Set to true if we want the profiling implementation to not write "
           "out files.")

compiler-rt/lib/xray/xray_segmented_array.h

 // backing the segments.
 //
 //===----------------------------------------------------------------------===//
 #ifndef XRAY_SEGMENTED_ARRAY_H
 #define XRAY_SEGMENTED_ARRAY_H
 
 #include "sanitizer_common/sanitizer_allocator.h"
 #include "xray_allocator.h"
+#include "xray_utils.h"
+#include <cassert>
 #include <type_traits>
 #include <utility>
 
 namespace __xray {
 
-namespace {
-
-constexpr size_t gcd(size_t a, size_t b) {
-  return (b == 0) ? a : gcd(b, a % b);
-}
-
-constexpr size_t lcm(size_t a, size_t b) { return a * b / gcd(a, b); }
-
-} // namespace
-
 /// The Array type provides an interface similar to std::vector<...> but does
 /// not shrink in size. Once constructed, elements can be appended but cannot be
 /// removed. The implementation is heavily dependent on the contract provided by
 /// the Allocator type, in that all memory will be released when the Allocator
 /// is destroyed. When an Array is destroyed, it will destroy elements in the
 /// backing store but will not free the memory. The parameter N defines how many
 /// elements of T there should be in a single block.
 ///
 /// We compute the least common multiple of the size of T and the cache line
 /// size, to allow us to maximise the number of T objects we can place in
 /// cache-line multiple sized blocks. To get back the number of T's, we divide
 /// this least common multiple by the size of T.
-template <class T, size_t N = lcm(sizeof(T), kCacheLineSize) / sizeof(T)>
+template <class T,
+          size_t N = lcm(next_pow2(sizeof(T)), kCacheLineSize) / sizeof(T)>
 struct Array {
   static constexpr size_t ChunkSize = N;
-  static constexpr size_t AllocatorChunkSize = sizeof(T) * ChunkSize;
+  static constexpr size_t ElementStorageSize = next_pow2(sizeof(T));
+  static constexpr size_t AllocatorChunkSize = ElementStorageSize * ChunkSize;
   using AllocatorType = Allocator<AllocatorChunkSize>;
   static_assert(std::is_trivially_destructible<T>::value,
                 "T must be trivially destructible.");
 
 private:
   // TODO: Consider co-locating the chunk information with the data in the
   // Block, as in an intrusive list -- i.e. putting the next and previous
   // pointer values inside the Block storage.
   struct Chunk {
     typename AllocatorType::Block Block;
     static constexpr size_t Size = N;
-    Chunk *Prev = nullptr;
-    Chunk *Next = nullptr;
+    Chunk *Prev = this;
+    Chunk *Next = this;
   };
 
   static Chunk SentinelChunk;
 
   AllocatorType *Alloc;
   Chunk *Head = &SentinelChunk;
   Chunk *Tail = &SentinelChunk;
   size_t Size = 0;
   size_t FreeElements = 0;
 
   // Here we keep track of chunks in the freelist, to allow us to re-use chunks
   // when elements are trimmed off the end.
   Chunk *Freelist = &SentinelChunk;
 
   Chunk *NewChunk() {
     // We need to handle the case in which enough elements have been trimmed to
     // allow us to re-use chunks we've allocated before. For this we look into
     // the Freelist, to see whether we need to actually allocate new blocks or
     // just re-use blocks we've already seen before.
     if (Freelist != &SentinelChunk) {
       auto *FreeChunk = Freelist;
       Freelist = FreeChunk->Next;
       FreeChunk->Next = &SentinelChunk;
+      Freelist->Prev = &SentinelChunk;
       return FreeChunk;
     }
 
-    auto Block = Alloc->Allocate();
+    auto Block = Alloc->Allocate(ElementStorageSize);
     if (Block.Data == nullptr)
       return nullptr;
+
     // TODO: Maybe use a separate managed allocator for Chunk instances?
-    auto C = reinterpret_cast<Chunk *>(InternalAlloc(sizeof(Chunk)));
+    auto C = reinterpret_cast<Chunk *>(
+        InternalAlloc(sizeof(Chunk), nullptr, kCacheLineSize));
     if (C == nullptr)
       return nullptr;
     C->Block = Block;
+    C->Prev = &SentinelChunk;
+    C->Next = &SentinelChunk;
     return C;
   }
 
   static AllocatorType &GetGlobalAllocator() {
     static AllocatorType *const GlobalAllocator = [] {
       AllocatorType *A = reinterpret_cast<AllocatorType *>(
           InternalAlloc(sizeof(AllocatorType)));
       new (A) AllocatorType(2 << 10, 0);
       return A;
     }();
 
     return *GlobalAllocator;
   }
 
   Chunk *InitHeadAndTail() {
     DCHECK_EQ(Head, &SentinelChunk);
     DCHECK_EQ(Tail, &SentinelChunk);
     auto Chunk = NewChunk();
     if (Chunk == nullptr)
       return nullptr;
-    Chunk->Prev = &SentinelChunk;
-    Chunk->Next = &SentinelChunk;
-    Head = Chunk;
-    Tail = Chunk;
-    return Chunk;
+    DCHECK_EQ(Chunk->Next, &SentinelChunk);
+    DCHECK_EQ(Chunk->Prev, &SentinelChunk);
+    return Head = Tail = Chunk;
   }
 
   Chunk *AppendNewChunk() {
     auto Chunk = NewChunk();
     if (Chunk == nullptr)
       return nullptr;
+    DCHECK_NE(Tail, &SentinelChunk);
+    DCHECK_EQ(Tail->Next, &SentinelChunk);
+    DCHECK_EQ(Chunk->Prev, &SentinelChunk);
+    DCHECK_EQ(Chunk->Next, &SentinelChunk);
     Tail->Next = Chunk;
     Chunk->Prev = Tail;
-    Chunk->Next = &SentinelChunk;
     Tail = Chunk;
-    return Chunk;
+    return Tail;
   }
 
   // This Iterator models a BidirectionalIterator.
   template <class U> class Iterator {
-    Chunk *C = nullptr;
+    Chunk *C = &SentinelChunk;
     size_t Offset = 0;
+    size_t Size = 0;
 
   public:
-    Iterator(Chunk *IC, size_t Off) : C(IC), Offset(Off) {}
+    Iterator(Chunk *IC, size_t Off, size_t S) : C(IC), Offset(Off), Size(S) {}
+    Iterator(const Iterator &) noexcept = default;
+    Iterator() noexcept = default;
+    Iterator(Iterator &&) noexcept = default;
+    Iterator &operator=(const Iterator &) = default;
+    Iterator &operator=(Iterator &&) = default;
+    ~Iterator() = default;
 
     Iterator &operator++() {
-      if (++Offset % N)
+      if (++Offset % N || Offset == Size)
         return *this;
 
-      DCHECK_NE(C, &SentinelChunk);
-
       // At this point, we know that Offset % N == 0, so we must advance the
       // chunk pointer.
       DCHECK_EQ(Offset % N, 0);
+      DCHECK_NE(Offset, Size);
+      DCHECK_NE(C, &SentinelChunk);
+      DCHECK_NE(C->Next, &SentinelChunk);
       C = C->Next;
+      DCHECK_NE(C, &SentinelChunk);
       return *this;
     }
 
     Iterator &operator--() {
       DCHECK_NE(C, &SentinelChunk);
       DCHECK_GT(Offset, 0);
 
-      // We check whether the offset was on a boundary before decrement, to see
-      // whether we need to retreat to the previous chunk.
-      if ((Offset-- % N) == 0)
+      auto PreviousOffset = Offset--;
+      if (PreviousOffset != Size && PreviousOffset % N == 0) {
+        DCHECK_NE(C->Prev, &SentinelChunk);
         C = C->Prev;
+      }
+
       return *this;
     }
 
     Iterator operator++(int) {
       Iterator Copy(*this);
       ++(*this);
       return Copy;
     }
public:
 
     template <class V, class W>
     friend bool operator!=(const Iterator<V> &L, const Iterator<W> &R) {
       return !(L == R);
     }
 
     U &operator*() const {
       DCHECK_NE(C, &SentinelChunk);
-      return reinterpret_cast<U *>(C->Block.Data)[Offset % N];
+      auto RelOff = Offset % N;
+      return *reinterpret_cast<U *>(reinterpret_cast<char *>(C->Block.Data) +
+                                    (RelOff * ElementStorageSize));
     }
 
     U *operator->() const {
       DCHECK_NE(C, &SentinelChunk);
-      return reinterpret_cast<U *>(C->Block.Data) + (Offset % N);
+      auto RelOff = Offset % N;
+      return reinterpret_cast<U *>(reinterpret_cast<char *>(C->Block.Data) +
+                                   (RelOff * ElementStorageSize));
     }
   };
 
 public:
   explicit Array(AllocatorType &A) : Alloc(&A) {}
   Array() : Array(GetGlobalAllocator()) {}
 
   Array(const Array &) = delete;
if (UNLIKELY(Head == &SentinelChunk))
       if (InitHeadAndTail() == nullptr)
         return nullptr;
 
     auto Offset = Size % N;
     if (UNLIKELY(Size != 0 && Offset == 0))
       if (AppendNewChunk() == nullptr)
         return nullptr;
 
-    auto Position = reinterpret_cast<T *>(Tail->Block.Data) + Offset;
+    auto Position =
+        reinterpret_cast<T *>(reinterpret_cast<char *>(Tail->Block.Data) +
+                              (Offset * ElementStorageSize));
     *Position = E;
     ++Size;
-    FreeElements -= FreeElements ? 1 : 0;
+    if (FreeElements)
+      --FreeElements;
     return Position;
   }
 
   template <class... Args> T *AppendEmplace(Args &&... args) {
     if (UNLIKELY(Head == &SentinelChunk))
       if (InitHeadAndTail() == nullptr)
         return nullptr;
 
     auto Offset = Size % N;
-    if (UNLIKELY(Size != 0 && Offset == 0))
-      if (AppendNewChunk() == nullptr)
+    auto *LatestChunk = Tail;
+    if (UNLIKELY(Size != 0 && Offset == 0)) {
+      LatestChunk = AppendNewChunk();
+      if (LatestChunk == nullptr)
         return nullptr;
+    }
 
-    auto Position = reinterpret_cast<T *>(Tail->Block.Data) + Offset;
+    DCHECK_NE(Tail, &SentinelChunk);
+    auto Position = reinterpret_cast<char *>(LatestChunk->Block.Data) +
+                    (Offset * ElementStorageSize);
+    DCHECK_EQ(reinterpret_cast<uintptr_t>(Position) % ElementStorageSize, 0);
     // In-place construct at Position.
-    new (Position) T(std::forward<Args>(args)...);
+    new (Position) T{std::forward<Args>(args)...};
+    if (FreeElements)
+      --FreeElements;
     ++Size;
-    FreeElements -= FreeElements ? 1 : 0;
-    return Position;
+    return reinterpret_cast<T *>(Position);
   }
 
   T &operator[](size_t Offset) const {
     DCHECK_LE(Offset, Size);
     // We need to traverse the array enough times to find the element at Offset.
     auto C = Head;
     while (Offset >= N) {
       C = C->Next;
       Offset -= N;
       DCHECK_NE(C, &SentinelChunk);
     }
-    auto Position = reinterpret_cast<T *>(C->Block.Data) + Offset;
-    return *Position;
+    return *reinterpret_cast<T *>(reinterpret_cast<char *>(C->Block.Data) +
+                                  (Offset * ElementStorageSize));
   }
 
   T &front() const {
     DCHECK_NE(Head, &SentinelChunk);
     DCHECK_NE(Size, 0u);
-    return *reinterpret_cast<T *>(Head->Block.Data);
+    return *begin();
   }
 
   T &back() const {
     DCHECK_NE(Tail, &SentinelChunk);
-    auto Offset = (Size - 1) % N;
-    return *(reinterpret_cast<T *>(Tail->Block.Data) + Offset);
+    DCHECK_NE(Size, 0u);
+    auto It = end();
+    --It;
+    return *It;
   }
 
   template <class Predicate> T *find_element(Predicate P) const {
     if (empty())
       return nullptr;
 
     auto E = end();
     for (auto I = begin(); I != E; ++I)
       if (P(*I))
         return &(*I);
 
     return nullptr;
   }
 
   /// Remove N Elements from the end. This leaves the blocks behind, and not
   /// require allocation of new blocks for new elements added after trimming.
   void trim(size_t Elements) {
     DCHECK_LE(Elements, Size);
+    DCHECK_GT(Size, 0);
+    auto OldSize = Size;
     Size -= Elements;
-    FreeElements += Elements;
 
-    // Here we need to check whether we've cleared enough elements to warrant
-    // putting blocks on to the freelist. We determine whether we need to
-    // right-size the internal list, by keeping track of the number of "free"
-    // elements still in the array.
-    auto ChunksToTrim = FreeElements / N;
-    for (size_t i = 0; i < ChunksToTrim; ++i, FreeElements -= N) {
+    DCHECK_NE(Head, &SentinelChunk);
+    DCHECK_NE(Tail, &SentinelChunk);
+
+    for (auto ChunksToTrim =
+             nearest_boundary(OldSize - nearest_boundary(Size, N), N) / N;
+         ChunksToTrim > 0; --ChunksToTrim) {
+      DCHECK_NE(Head, &SentinelChunk);
+      DCHECK_NE(Tail, &SentinelChunk);
       // Put the tail into the Freelist.
       auto *FreeChunk = Tail;
       Tail = Tail->Prev;
       if (Tail == &SentinelChunk)
         Head = Tail;
       else
         Tail->Next = &SentinelChunk;
+
+      DCHECK_EQ(Tail->Next, &SentinelChunk);
       FreeChunk->Next = Freelist;
-      FreeChunk->Prev = Freelist->Prev;
+      FreeChunk->Prev = &SentinelChunk;
+      if (Freelist != &SentinelChunk)
+        Freelist->Prev = FreeChunk;
       Freelist = FreeChunk;
     }
+    FreeElements = nearest_boundary(Size, N) - Size;
   }
 
   // Provide iterators.
-  Iterator<T> begin() const { return Iterator<T>(Head, 0); }
-  Iterator<T> end() const { return Iterator<T>(Tail, Size); }
-  Iterator<const T> cbegin() const { return Iterator<const T>(Head, 0); }
-  Iterator<const T> cend() const { return Iterator<const T>(Tail, Size); }
+  Iterator<T> begin() const { return Iterator<T>(Head, 0, Size); }
+  Iterator<T> end() const { return Iterator<T>(Tail, Size, Size); }
+  Iterator<const T> cbegin() const { return Iterator<const T>(Head, 0, Size); }
+  Iterator<const T> cend() const { return Iterator<const T>(Tail, Size, Size); }
 };
 
 // We need to have this storage definition out-of-line so that the compiler can
 // ensure that storage for the SentinelChunk is defined and has a single
 // address.
 template <class T, size_t N>
 typename Array<T, N>::Chunk Array<T, N>::SentinelChunk;
 
 } // namespace __xray
 
 #endif // XRAY_SEGMENTED_ARRAY_H

compiler-rt/lib/xray/xray_utils.h

 
 // EINTR-safe read routine, providing a file descriptor and a character range.
 std::pair<ssize_t, bool> retryingReadSome(int Fd, char *Begin, char *End);
 
 // EINTR-safe open routine, uses flag-provided values for initialising a log
 // file.
 int getLogFD();
 
+constexpr size_t gcd(size_t a, size_t b) {
+  return (b == 0) ? a : gcd(b, a % b);
+}
+
+constexpr size_t lcm(size_t a, size_t b) { return a * b / gcd(a, b); }
+
+template <class T> constexpr T nearest_boundary(T number, T multiple) {
+  return multiple * ((number / multiple) + (number % multiple ? 1 : 0));
+}
+
+constexpr size_t next_pow2_helper(size_t num, size_t acc) {
+  return (1u << acc) >= num ? (1u << acc) : next_pow2_helper(num, acc + 1);
+}
+
+constexpr size_t next_pow2(size_t number) {
+  return next_pow2_helper(number, 1);
+}
+
 } // namespace __xray
 
 #endif // XRAY_UTILS_H