1 files changed, 163 insertions, 145 deletions

diff --git a/lib/xray/xray_segmented_array.h b/lib/xray/xray_segmented_array.h
index f001b230c..018c2aa9e 100644
--- a/lib/xray/xray_segmented_array.h
+++ b/lib/xray/xray_segmented_array.h
@@ -9,7 +9,7 @@
 //
 // This file is a part of XRay, a dynamic runtime instrumentation system.
 //
-// Defines the implementation of a segmented array, with fixed-size chunks
+// Defines the implementation of a segmented array, with fixed-size segments
 // backing the segments.
 //
 //===----------------------------------------------------------------------===//
@@ -24,113 +24,128 @@
 #include <utility>
 
 namespace __xray {
-struct Chunk {
-  void *Data;
-  Chunk *Prev;
-  Chunk *Next;
-};
-
-using ChunkAllocator = Allocator<next_pow2(sizeof(Chunk)) * 8>;
 
 /// 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(next_pow2(sizeof(T)), kCacheLineSize) / sizeof(T)>
-struct Array {
-  static constexpr size_t ChunkSize = N;
-  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.");
+/// backing store but will not free the memory.
+template <class T> class Array {
+  struct SegmentBase {
+    SegmentBase *Prev;
+    SegmentBase *Next;
+  };
 
-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.
-  static Chunk SentinelChunk;
+  // We want each segment of the array to be cache-line aligned, and elements of
+  // the array be offset from the beginning of the segment.
+  struct Segment : SegmentBase {
+    char Data[];
+  };
+
+public:
+  // Each segment of the array will be laid out with the following assumptions:
+  //
+  //   - Each segment will be on a cache-line address boundary (kCacheLineSize
+  //     aligned).
+  //
+  //   - The elements will be accessed through an aligned pointer, dependent on
+  //     the alignment of T.
+  //
+  //   - Each element is at least two-pointers worth from the beginning of the
+  //     Segment, aligned properly, and the rest of the elements are accessed
+  //     through appropriate alignment.
+  //
+  // We then compute the size of the segment to follow this logic:
+  //
+  //   - Compute the number of elements that can fit within
+  //     kCacheLineSize-multiple segments, minus the size of two pointers.
+  //
+  //   - Request cacheline-multiple sized elements from the allocator.
+  static constexpr size_t AlignedElementStorageSize =
+      sizeof(typename std::aligned_storage<sizeof(T), alignof(T)>::type);
+
+  static constexpr size_t SegmentSize =
+      nearest_boundary(sizeof(Segment) + next_pow2(sizeof(T)), kCacheLineSize);
+
+  using AllocatorType = Allocator<SegmentSize>;
+
+  static constexpr size_t ElementsPerSegment =
+      (SegmentSize - sizeof(Segment)) / next_pow2(sizeof(T));
+
+  static_assert(ElementsPerSegment > 0,
+                "Must have at least 1 element per segment.");
+
+  static SegmentBase SentinelSegment;
 
+private:
   AllocatorType *Alloc;
-  ChunkAllocator *ChunkAlloc;
-  Chunk *Head = &SentinelChunk;
-  Chunk *Tail = &SentinelChunk;
+  SegmentBase *Head = &SentinelSegment;
+  SegmentBase *Tail = &SentinelSegment;
   size_t Size = 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;
+  // Here we keep track of segments in the freelist, to allow us to re-use
+  // segments when elements are trimmed off the end.
+  SegmentBase *Freelist = &SentinelSegment;
 
-  Chunk *NewChunk() {
+  Segment *NewSegment() {
     // 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
+    // allow us to re-use segments 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;
+    if (Freelist != &SentinelSegment) {
+      auto *FreeSegment = Freelist;
+      Freelist = FreeSegment->Next;
+      FreeSegment->Next = &SentinelSegment;
+      Freelist->Prev = &SentinelSegment;
+      return static_cast<Segment *>(FreeSegment);
     }
 
-    auto Block = Alloc->Allocate();
-    if (Block.Data == nullptr)
-      return nullptr;
-
-    auto ChunkElement = ChunkAlloc->Allocate();
-    if (ChunkElement.Data == nullptr)
+    auto SegmentBlock = Alloc->Allocate();
+    if (SegmentBlock.Data == nullptr)
       return nullptr;
 
-    // Placement-new the Chunk element at the appropriate location.
-    auto C = reinterpret_cast<Chunk *>(ChunkElement.Data);
-    Chunk LocalC{Block.Data, &SentinelChunk, &SentinelChunk};
-    internal_memcpy(C, &LocalC, sizeof(LocalC));
-    return C;
+    // Placement-new the Segment element at the beginning of the SegmentBlock.
+    auto S = reinterpret_cast<Segment *>(SegmentBlock.Data);
+    new (S) SegmentBase{&SentinelSegment, &SentinelSegment};
+    return S;
   }
 
-  Chunk *InitHeadAndTail() {
-    DCHECK_EQ(Head, &SentinelChunk);
-    DCHECK_EQ(Tail, &SentinelChunk);
-    auto Chunk = NewChunk();
-    if (Chunk == nullptr)
+  Segment *InitHeadAndTail() {
+    DCHECK_EQ(Head, &SentinelSegment);
+    DCHECK_EQ(Tail, &SentinelSegment);
+    auto Segment = NewSegment();
+    if (Segment == nullptr)
       return nullptr;
-    DCHECK_EQ(Chunk->Next, &SentinelChunk);
-    DCHECK_EQ(Chunk->Prev, &SentinelChunk);
-    return Head = Tail = Chunk;
+    DCHECK_EQ(Segment->Next, &SentinelSegment);
+    DCHECK_EQ(Segment->Prev, &SentinelSegment);
+    Head = Tail = static_cast<SegmentBase *>(Segment);
+    return Segment;
   }
 
-  Chunk *AppendNewChunk() {
-    auto Chunk = NewChunk();
-    if (Chunk == nullptr)
+  Segment *AppendNewSegment() {
+    auto S = NewSegment();
+    if (S == 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;
-    Tail = Chunk;
-    return Tail;
+    DCHECK_NE(Tail, &SentinelSegment);
+    DCHECK_EQ(Tail->Next, &SentinelSegment);
+    DCHECK_EQ(S->Prev, &SentinelSegment);
+    DCHECK_EQ(S->Next, &SentinelSegment);
+    Tail->Next = S;
+    S->Prev = Tail;
+    Tail = S;
+    return static_cast<Segment *>(Tail);
   }
 
   // This Iterator models a BidirectionalIterator.
   template <class U> class Iterator {
-    Chunk *C = &SentinelChunk;
+    SegmentBase *S = &SentinelSegment;
     size_t Offset = 0;
     size_t Size = 0;
 
   public:
-    Iterator(Chunk *IC, size_t Off, size_t S) : C(IC), Offset(Off), Size(S) {}
+    Iterator(SegmentBase *IS, size_t Off, size_t S)
+        : S(IS), Offset(Off), Size(S) {}
     Iterator(const Iterator &) noexcept = default;
     Iterator() noexcept = default;
     Iterator(Iterator &&) noexcept = default;
@@ -139,28 +154,28 @@ private:
     ~Iterator() = default;
 
     Iterator &operator++() {
-      if (++Offset % N || Offset == Size)
+      if (++Offset % ElementsPerSegment || Offset == Size)
         return *this;
 
       // At this point, we know that Offset % N == 0, so we must advance the
-      // chunk pointer.
-      DCHECK_EQ(Offset % N, 0);
+      // segment pointer.
+      DCHECK_EQ(Offset % ElementsPerSegment, 0);
       DCHECK_NE(Offset, Size);
-      DCHECK_NE(C, &SentinelChunk);
-      DCHECK_NE(C->Next, &SentinelChunk);
-      C = C->Next;
-      DCHECK_NE(C, &SentinelChunk);
+      DCHECK_NE(S, &SentinelSegment);
+      DCHECK_NE(S->Next, &SentinelSegment);
+      S = S->Next;
+      DCHECK_NE(S, &SentinelSegment);
       return *this;
     }
 
     Iterator &operator--() {
-      DCHECK_NE(C, &SentinelChunk);
+      DCHECK_NE(S, &SentinelSegment);
       DCHECK_GT(Offset, 0);
 
       auto PreviousOffset = Offset--;
-      if (PreviousOffset != Size && PreviousOffset % N == 0) {
-        DCHECK_NE(C->Prev, &SentinelChunk);
-        C = C->Prev;
+      if (PreviousOffset != Size && PreviousOffset % ElementsPerSegment == 0) {
+        DCHECK_NE(S->Prev, &SentinelSegment);
+        S = S->Prev;
       }
 
       return *this;
@@ -180,7 +195,7 @@ private:
 
     template <class V, class W>
     friend bool operator==(const Iterator<V> &L, const Iterator<W> &R) {
-      return L.C == R.C && L.Offset == R.Offset;
+      return L.S == R.S && L.Offset == R.Offset;
     }
 
     template <class V, class W>
@@ -189,31 +204,29 @@ private:
     }
 
     U &operator*() const {
-      DCHECK_NE(C, &SentinelChunk);
-      auto RelOff = Offset % N;
-      return *reinterpret_cast<U *>(reinterpret_cast<char *>(C->Data) +
-                                    (RelOff * ElementStorageSize));
+      DCHECK_NE(S, &SentinelSegment);
+      auto RelOff = Offset % ElementsPerSegment;
+
+      // We need to compute the character-aligned pointer, offset from the
+      // segment's Data location to get the element in the position of Offset.
+      auto Base = static_cast<Segment *>(S)->Data;
+      auto AlignedOffset = Base + (RelOff * AlignedElementStorageSize);
+      return *reinterpret_cast<U *>(AlignedOffset);
     }
 
-    U *operator->() const {
-      DCHECK_NE(C, &SentinelChunk);
-      auto RelOff = Offset % N;
-      return reinterpret_cast<U *>(reinterpret_cast<char *>(C->Data) +
-                                   (RelOff * ElementStorageSize));
-    }
+    U *operator->() const { return &(**this); }
   };
 
 public:
-  explicit Array(AllocatorType &A, ChunkAllocator &CA)
-      : Alloc(&A), ChunkAlloc(&CA) {}
+  explicit Array(AllocatorType &A) : Alloc(&A) {}
 
   Array(const Array &) = delete;
   Array(Array &&O) NOEXCEPT : Alloc(O.Alloc),
                               Head(O.Head),
                               Tail(O.Tail),
                               Size(O.Size) {
-    O.Head = &SentinelChunk;
-    O.Tail = &SentinelChunk;
+    O.Head = &SentinelSegment;
+    O.Tail = &SentinelSegment;
     O.Size = 0;
   }
 
@@ -227,39 +240,41 @@ public:
   size_t size() const { return Size; }
 
   T *Append(const T &E) {
-    if (UNLIKELY(Head == &SentinelChunk))
+    if (UNLIKELY(Head == &SentinelSegment))
       if (InitHeadAndTail() == nullptr)
         return nullptr;
 
-    auto Offset = Size % N;
+    auto Offset = Size % ElementsPerSegment;
     if (UNLIKELY(Size != 0 && Offset == 0))
-      if (AppendNewChunk() == nullptr)
+      if (AppendNewSegment() == nullptr)
         return nullptr;
 
-    auto Position = reinterpret_cast<T *>(reinterpret_cast<char *>(Tail->Data) +
-                                          (Offset * ElementStorageSize));
+    auto Base = static_cast<Segment *>(Tail)->Data;
+    auto AlignedOffset = Base + (Offset * AlignedElementStorageSize);
+    auto Position = reinterpret_cast<T *>(AlignedOffset);
     *Position = E;
     ++Size;
     return Position;
   }
 
   template <class... Args> T *AppendEmplace(Args &&... args) {
-    if (UNLIKELY(Head == &SentinelChunk))
+    if (UNLIKELY(Head == &SentinelSegment))
       if (InitHeadAndTail() == nullptr)
         return nullptr;
 
-    auto Offset = Size % N;
-    auto *LatestChunk = Tail;
+    auto Offset = Size % ElementsPerSegment;
+    auto *LatestSegment = Tail;
     if (UNLIKELY(Size != 0 && Offset == 0)) {
-      LatestChunk = AppendNewChunk();
-      if (LatestChunk == nullptr)
+      LatestSegment = AppendNewSegment();
+      if (LatestSegment == nullptr)
         return nullptr;
     }
 
-    DCHECK_NE(Tail, &SentinelChunk);
-    auto Position = reinterpret_cast<char *>(LatestChunk->Data) +
-                    (Offset * ElementStorageSize);
-    DCHECK_EQ(reinterpret_cast<uintptr_t>(Position) % ElementStorageSize, 0);
+    DCHECK_NE(Tail, &SentinelSegment);
+    auto Base = static_cast<Segment *>(LatestSegment)->Data;
+    auto AlignedOffset = Base + (Offset * AlignedElementStorageSize);
+    auto Position = reinterpret_cast<T *>(AlignedOffset);
+
     // In-place construct at Position.
     new (Position) T{std::forward<Args>(args)...};
     ++Size;
@@ -269,24 +284,26 @@ public:
   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 S = Head;
+    while (Offset >= ElementsPerSegment) {
+      S = S->Next;
+      Offset -= ElementsPerSegment;
+      DCHECK_NE(S, &SentinelSegment);
     }
-    return *reinterpret_cast<T *>(reinterpret_cast<char *>(C->Data) +
-                                  (Offset * ElementStorageSize));
+    auto Base = static_cast<Segment *>(S)->Data;
+    auto AlignedOffset = Base + (Offset * AlignedElementStorageSize);
+    auto Position = reinterpret_cast<T *>(AlignedOffset);
+    return *reinterpret_cast<T *>(Position);
   }
 
   T &front() const {
-    DCHECK_NE(Head, &SentinelChunk);
+    DCHECK_NE(Head, &SentinelSegment);
     DCHECK_NE(Size, 0u);
     return *begin();
   }
 
   T &back() const {
-    DCHECK_NE(Tail, &SentinelChunk);
+    DCHECK_NE(Tail, &SentinelSegment);
     DCHECK_NE(Size, 0u);
     auto It = end();
     --It;
@@ -313,28 +330,29 @@ public:
     auto OldSize = Size;
     Size -= Elements;
 
-    DCHECK_NE(Head, &SentinelChunk);
-    DCHECK_NE(Tail, &SentinelChunk);
+    DCHECK_NE(Head, &SentinelSegment);
+    DCHECK_NE(Tail, &SentinelSegment);
 
-    for (auto ChunksToTrim =
-             (nearest_boundary(OldSize, N) - nearest_boundary(Size, N)) / N;
-         ChunksToTrim > 0; --ChunksToTrim) {
-      DCHECK_NE(Head, &SentinelChunk);
-      DCHECK_NE(Tail, &SentinelChunk);
+    for (auto SegmentsToTrim = (nearest_boundary(OldSize, ElementsPerSegment) -
+                                nearest_boundary(Size, ElementsPerSegment)) /
+                               ElementsPerSegment;
+         SegmentsToTrim > 0; --SegmentsToTrim) {
+      DCHECK_NE(Head, &SentinelSegment);
+      DCHECK_NE(Tail, &SentinelSegment);
       // Put the tail into the Freelist.
-      auto *FreeChunk = Tail;
+      auto *FreeSegment = Tail;
       Tail = Tail->Prev;
-      if (Tail == &SentinelChunk)
+      if (Tail == &SentinelSegment)
         Head = Tail;
       else
-        Tail->Next = &SentinelChunk;
-
-      DCHECK_EQ(Tail->Next, &SentinelChunk);
-      FreeChunk->Next = Freelist;
-      FreeChunk->Prev = &SentinelChunk;
-      if (Freelist != &SentinelChunk)
-        Freelist->Prev = FreeChunk;
-      Freelist = FreeChunk;
+        Tail->Next = &SentinelSegment;
+
+      DCHECK_EQ(Tail->Next, &SentinelSegment);
+      FreeSegment->Next = Freelist;
+      FreeSegment->Prev = &SentinelSegment;
+      if (Freelist != &SentinelSegment)
+        Freelist->Prev = FreeSegment;
+      Freelist = FreeSegment;
     }
   }
 
@@ -346,11 +364,11 @@ public:
 };
 
 // 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
+// ensure that storage for the SentinelSegment is defined and has a single
 // address.
-template <class T, size_t N>
-Chunk Array<T, N>::SentinelChunk{nullptr, &Array<T, N>::SentinelChunk,
-                                 &Array<T, N>::SentinelChunk};
+template <class T>
+typename Array<T>::SegmentBase Array<T>::SentinelSegment{
+    &Array<T>::SentinelSegment, &Array<T>::SentinelSegment};
 
 } // namespace __xray