[ADT] Implement the Waymarking as an independent utility

This is the Waymarking algorithm implemented as an independent utility.
The utility is operating on a range of sequential elements.
First we "tag" the elements, by calling `fillWaymarks`.
Then we can "follow" the tags from every element inside the tagged
range, and reach the "head" (the first element), by calling
`followWaymarks`.

Differential Revision: https://reviews.llvm.org/D74415

include/llvm/ADT/Waymarking.h

@@ -0,0 +1,325 @@
+//===- Waymarking.h - Array waymarking algorithm ----------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Utility to backtrace an array's head, from a pointer into it. For the
+// backtrace to work, we use "Waymarks", which are special tags embedded into
+// the array's elements.
+//
+// A Tag of n-bits (in size) is composed as follows:
+//
+// bits: |   n-1   |             n-2 ... 0              |
+//       .---------.------------------------------------.
+//       |Stop Mask|(2^(n-1))-ary numeric system - digit|
+//       '---------'------------------------------------'
+//
+// Backtracing is done as follows:
+// Walk back (starting from a given pointer to an element into the array), until
+// a tag with a "Stop Mask" is reached. Then start calculating the "Offset" from
+// the array's head, by picking up digits along the way, until another stop is
+// reached. The "Offset" is then subtracted from the current pointer, and the
+// result is the array's head.
+// A special case - if we first encounter a Tag with a Stop and a zero digit,
+// then this is already the head.
+//
+// For example:
+// In case of 2 bits:
+//
+// Tags:
+// x0 - binary digit 0
+// x1 - binary digit 1
+// 1x - stop and calculate (s)
+//
+// Array:
+//         .---.---.---.---.---.---.---.---.---.---.---.---.---.---.---.---.
+// head -> |s0 |s1 | 0 |s1 | 0 | 0 |s1 | 1 | 1 |s1 | 0 | 1 | 0 |s1 | 0 | 1 |
+//         '---'---'---'---'---'---'---'---'---'---'---'---'---'---'---'---'
+//             |-1 |-2     |-4         |-7         |-10            |-14
+//          <_ |   |       |           |           |               |
+//          <_____ |       |           |           |               |
+//          <_____________ |           |           |               |
+//          <_________________________ |           |               |
+//          <_____________________________________ |               |
+//          <_____________________________________________________ |
+//
+//
+// In case of 3 bits:
+//
+// Tags:
+// x00 - quaternary digit 0
+// x01 - quaternary digit 1
+// x10 - quaternary digit 2
+// x11 - quaternary digit 3
+// 1xy - stop and calculate (s)
+//
+// Array:
+//         .---.---.---.---.---.---.---.---.---.---.---.---.---.---.---.---.
+// head -> |s0 |s1 |s2 |s3 | 0 |s1 | 2 |s1 | 0 |s2 | 2 |s2 | 0 |s3 | 2 |s3 |
+//         '---'---'---'---'---'---'---'---'---'---'---'---'---'---'---'---'
+//             |-1 |-2 |-3 |-4     |-6     |-8     |-10    |-12    |-14    |-16
+//          <_ |   |   |   |       |       |       |       |       |       |
+//          <_____ |   |   |       |       |       |       |       |       |
+//          <_________ |   |       |       |       |       |       |       |
+//          <_____________ |       |       |       |       |       |       |
+//          <_____________________ |       |       |       |       |       |
+//          <_____________________________ |       |       |       |       |
+//          <_____________________________________ |       |       |       |
+//          <_____________________________________________ |       |       |
+//          <_____________________________________________________ |       |
+//          <_____________________________________________________________ |
+//
+//
+// The API introduce 2 functions:
+// 1. fillWaymarks
+// 2. followWaymarks
+//
+// Example:
+//   int N = 10;
+//   int M = 5;
+//   int **A = new int *[N + M];   // Define the array.
+//   for (int I = 0; I < N + M; ++I)
+//     A[I] = new int(I);
+//
+//   fillWaymarks(A, A + N);       // Set the waymarks for the first N elements
+//                                 // of the array.
+//                                 // Note that it must be done AFTER we fill
+//                                 // the array's elements.
+//
+//   ...                           // Elements which are not in the range
+//                                 // [A, A+N) will not be marked, and we won't
+//                                 // be able to call followWaymarks on them.
+//
+//   ...                           // Elements which will be changed after the
+//                                 // call to fillWaymarks, will have to be
+//                                 // retagged.
+//
+//   fillWaymarks(A + N, A + N + M, N); // Set the waymarks of the remaining M
+//                                      // elements.
+//   ...
+//   int **It = A + N + 1;
+//   int **B = followWaymarks(It); // Find the head of the array containing It.
+//   assert(B == A);
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_ADT_WAYMARKING_H
+#define LLVM_ADT_WAYMARKING_H
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/PointerLikeTypeTraits.h"
+
+namespace llvm {
+
+namespace detail {
+
+template <unsigned NumBits> struct WaymarkingTraits {
+  enum : unsigned {
+    // The number of bits of a Waymarking Tag.
+    NUM_BITS = NumBits,
+
+    // A Tag is composed from a Mark and a Stop mask.
+    MARK_SIZE = NUM_BITS - 1,
+    STOP_MASK = (1 << MARK_SIZE),
+    MARK_MASK = (STOP_MASK - 1),
+    TAG_MASK = (MARK_MASK | STOP_MASK),
+
+    // The number of pre-calculated tags (for fast fill).
+    NUM_STATIC_TAGS = 32
+  };
+
+private:
+  // Add a new tag, calculated from Count and Stop, to the Vals pack, while
+  // continuing recursively to decrease Len down to 0.
+  template <unsigned Len, bool Stop, unsigned Count, uint8_t... Vals>
+  struct AddTag;
+
+  // Delegate to the specialized AddTag according to the need of a Stop mask.
+  template <unsigned Len, unsigned Count, uint8_t... Vals> struct GenTag {
+    typedef
+        typename AddTag<Len, (Count <= MARK_MASK), Count, Vals...>::Xdata Xdata;
+  };
+
+  // Start adding tags while calculating the next Count, which is actually the
+  // number of already calculated tags (equivalent to the position in the
+  // array).
+  template <unsigned Len, uint8_t... Vals> struct GenOffset {
+    typedef typename GenTag<Len, sizeof...(Vals), Vals...>::Xdata Xdata;
+  };
+
+  // Add the tag and remove it from Count.
+  template <unsigned Len, unsigned Count, uint8_t... Vals>
+  struct AddTag<Len, false, Count, Vals...> {
+    typedef typename GenTag<Len - 1, (Count >> MARK_SIZE), Vals...,
+                            Count & MARK_MASK>::Xdata Xdata;
+  };
+
+  // We have reached the end of this Count, so start with a new Count.
+  template <unsigned Len, unsigned Count, uint8_t... Vals>
+  struct AddTag<Len, true, Count, Vals...> {
+    typedef typename GenOffset<Len - 1, Vals...,
+                               (Count & MARK_MASK) | STOP_MASK>::Xdata Xdata;
+  };
+
+  template <unsigned Count, uint8_t... Vals> struct TagsData {
+    // The remaining number for calculating the next tag, following the last one
+    // in Values.
+    static const unsigned Remain = Count;
+
+    // The array of ordered pre-calculated Tags.
+    static const uint8_t Values[sizeof...(Vals)];
+  };
+
+  // Specialize the case when Len equals 0, as the recursion stop condition.
+  template <unsigned Count, uint8_t... Vals>
+  struct AddTag<0, false, Count, Vals...> {
+    typedef TagsData<Count, Vals...> Xdata;
+  };
+
+  template <unsigned Count, uint8_t... Vals>
+  struct AddTag<0, true, Count, Vals...> {
+    typedef TagsData<Count, Vals...> Xdata;
+  };
+
+public:
+  typedef typename GenOffset<NUM_STATIC_TAGS>::Xdata Tags;
+};
+
+template <unsigned NumBits>
+template <unsigned Count, uint8_t... Vals>
+const uint8_t WaymarkingTraits<NumBits>::TagsData<
+    Count, Vals...>::Values[sizeof...(Vals)] = {Vals...};
+
+} // end namespace detail
+
+/// This class is responsible for tagging (and retrieving the tag of) a given
+/// element of type T.
+template <class T, class WTraits = detail::WaymarkingTraits<
+                       PointerLikeTypeTraits<T>::NumLowBitsAvailable>>
+struct Waymarker {
+  using Traits = WTraits;
+  static void setWaymark(T &N, unsigned Tag) { N.setWaymark(Tag); }
+  static unsigned getWaymark(const T &N) { return N.getWaymark(); }
+};
+
+template <class T, class WTraits> struct Waymarker<T *, WTraits> {
+  using Traits = WTraits;
+  static void setWaymark(T *&N, unsigned Tag) {
+    reinterpret_cast<uintptr_t &>(N) |= static_cast<uintptr_t>(Tag);
+  }
+  static unsigned getWaymark(const T *N) {
+    return static_cast<unsigned>(reinterpret_cast<uintptr_t>(N)) &
+           Traits::TAG_MASK;
+  }
+};
+
+/// Sets up the waymarking algorithm's tags for a given range [Begin, End).
+///
+/// \param Begin The beginning of the range to mark with tags (inclusive).
+/// \param End The ending of the range to mark with tags (exclusive).
+/// \param Offset The position in the supposed tags array from which to start
+/// marking the given range.
+template <class TIter, class Marker = Waymarker<
+                           typename std::iterator_traits<TIter>::value_type>>
+void fillWaymarks(TIter Begin, TIter End, size_t Offset = 0) {
+  if (Begin == End)
+    return;
+
+  size_t Count = Marker::Traits::Tags::Remain;
+  if (Offset <= Marker::Traits::NUM_STATIC_TAGS) {
+    // Start by filling the pre-calculated tags, starting from the given offset.
+    while (Offset != Marker::Traits::NUM_STATIC_TAGS) {
+      Marker::setWaymark(*Begin, Marker::Traits::Tags::Values[Offset]);
+
+      ++Offset;
+      ++Begin;
+
+      if (Begin == End)
+        return;
+    }
+  } else {
+    // The given offset is larger than the number of pre-calculated tags, so we
+    // must do it the hard way.
+    // Calculate the next remaining Count, as if we have filled the tags up to
+    // the given offset.
+    size_t Off = Marker::Traits::NUM_STATIC_TAGS;
+    do {
+      ++Off;
+
+      unsigned Tag = Count & Marker::Traits::MARK_MASK;
+
+      // If the count can fit into the tag, then the counting must stop.
+      if (Count <= Marker::Traits::MARK_MASK) {
+        Tag |= Marker::Traits::STOP_MASK;
+        Count = Off;
+      } else
+        Count >>= Marker::Traits::MARK_SIZE;
+    } while (Off != Offset);
+  }
+
+  // By now, we have the matching remaining Count for the current offset.
+  do {
+    ++Offset;
+
+    unsigned Tag = Count & Marker::Traits::MARK_MASK;
+
+    // If the count can fit into the tag, then the counting must stop.
+    if (Count <= Marker::Traits::MARK_MASK) {
+      Tag |= Marker::Traits::STOP_MASK;
+      Count = Offset;
+    } else
+      Count >>= Marker::Traits::MARK_SIZE;
+
+    Marker::setWaymark(*Begin, Tag);
+    ++Begin;
+  } while (Begin != End);
+}
+
+/// Sets up the waymarking algorithm's tags for a given range.
+///
+/// \param Range The range to mark with tags.
+/// \param Offset The position in the supposed tags array from which to start
+/// marking the given range.
+template <typename R, class Marker = Waymarker<typename std::remove_reference<
+                          decltype(*std::begin(std::declval<R &>()))>::type>>
+void fillWaymarks(R &&Range, size_t Offset = 0) {
+  return fillWaymarks<decltype(std::begin(std::declval<R &>())), Marker>(
+      adl_begin(Range), adl_end(Range), Offset);
+}
+
+/// Retrieves the element marked with tag of only STOP_MASK, by following the
+/// waymarks. This is the first element in a range passed to a previous call to
+/// \c fillWaymarks with \c Offset 0.
+///
+/// For the trivial usage of calling \c fillWaymarks(Array), and \I is an
+/// iterator inside \c Array, this function retrieves the head of \c Array, by
+/// following the waymarks.
+///
+/// \param I The iterator into an array which was marked by the waymarking tags
+/// (by a previous call to \c fillWaymarks).
+template <class TIter, class Marker = Waymarker<
+                           typename std::iterator_traits<TIter>::value_type>>
+TIter followWaymarks(TIter I) {
+  unsigned Tag;
+  do
+    Tag = Marker::getWaymark(*I--);
+  while (!(Tag & Marker::Traits::STOP_MASK));
+
+  // Special case for the first Use.
+  if (Tag != Marker::Traits::STOP_MASK) {
+    ptrdiff_t Offset = Tag & Marker::Traits::MARK_MASK;
+    while (!((Tag = Marker::getWaymark(*I)) & Marker::Traits::STOP_MASK)) {
+      Offset = (Offset << Marker::Traits::MARK_SIZE) + Tag;
+      --I;
+    }
+    I -= Offset;
+  }
+  return ++I;
+}
+
+} // end namespace llvm
+
+#endif // LLVM_ADT_WAYMARKING_H

unittests/ADT/CMakeLists.txt

@@ -73,6 +73,7 @@ add_llvm_unittest(ADTTests
   TinyPtrVectorTest.cpp
   TripleTest.cpp
   TwineTest.cpp
+  WaymarkingTest.cpp
   )
 
 target_link_libraries(ADTTests PRIVATE LLVMTestingSupport)

unittests/ADT/WaymarkingTest.cpp

@@ -0,0 +1,142 @@
+//===- llvm/unittest/IR/WaymarkTest.cpp - Waymarking unit tests -----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/Waymarking.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+static const int N = 100;
+
+// Get the Waymarking Tag of the pointer.
+static int tag(int *P) {
+  return static_cast<int>(reinterpret_cast<uintptr_t>(P) &
+                          uintptr_t(alignof(int *) - 1));
+}
+
+// Get the actual pointer, by stripping the Waymarking Tag.
+static int *ref(int *P) {
+  return reinterpret_cast<int *>(reinterpret_cast<uintptr_t>(P) &
+                                 ~uintptr_t(alignof(int *) - 1));
+}
+
+static int **createArray(int Len) {
+  int **A = new int *[Len];
+  for (int I = 0; I < Len; ++I)
+    A[I] = new int(I);
+  return A;
+}
+
+static void freeArray(int **A, int Len) {
+  for (int I = 0; I < Len; ++I)
+    delete ref(A[I]);
+  delete[] A;
+}
+
+// Verify the values stored in the array are as expected, and did not change due
+// to fillWaymarks.
+static void verifyArrayValues(int **A, int Begin, int End) {
+  for (int I = Begin; I < End; ++I)
+    EXPECT_EQ(I, *ref(A[I]));
+}
+
+static void verifyArrayValues(int **A, int Len) {
+  verifyArrayValues(A, 0, Len);
+}
+
+// Verify that we can follow the waymarks to the array's head from each element
+// of the array.
+static void verifyFollowWaymarks(int **A, int Len) {
+  for (int I = 0; I < Len; ++I) {
+    int **P = followWaymarks(A + I);
+    EXPECT_EQ(A, P);
+  }
+}
+
+namespace {
+
+// Test filling and following the waymarks of a single array.
+TEST(WaymarkingTest, SingleHead) {
+  const int N2 = 2 * N;
+  int **volatile A = createArray(N2);
+
+  // Fill the first half of the array with waymarks.
+  fillWaymarks(A, A + N, 0);
+  verifyArrayValues(A, N2);
+
+  verifyFollowWaymarks(A, N);
+
+  // Fill the rest of the waymarks (continuing from where we stopped).
+  fillWaymarks(A + N, A + N2, N);
+  verifyArrayValues(A, N2);
+
+  verifyFollowWaymarks(A, N);
+
+  freeArray(A, N2);
+}
+
+// Test filling and following the waymarks of an array split into several
+// different sections of waymarks (treated just like separate arrays).
+TEST(WaymarkingTest, MultiHead) {
+  const int N2 = 2 * N;
+  const int N3 = 3 * N;
+  int **volatile A = createArray(N3);
+
+  // Separate the array into 3 sections of waymarks.
+  fillWaymarks(A, A + N, 0);
+  fillWaymarks(A + N, A + N2, 0);
+  fillWaymarks(A + N2, A + N3, 0);
+  verifyArrayValues(A, N3);
+
+  verifyFollowWaymarks(A, N);
+  verifyFollowWaymarks(A + N, N2 - N);
+  verifyFollowWaymarks(A + N2, N3 - N2);
+
+  freeArray(A, N3);
+}
+
+// Test reseting (value and tag of) elements inside an array of waymarks.
+TEST(WaymarkingTest, Reset) {
+  int **volatile A = createArray(N);
+
+  fillWaymarks(A, A + N, 0);
+  verifyArrayValues(A, N);
+
+  const int N2 = N / 2;
+  const int N3 = N / 3;
+  const int N4 = N / 4;
+
+  // Reset specific elements and check that the tag remains the same.
+  int T2 = tag(A[N2]);
+  delete ref(A[N2]);
+  A[N2] = new int(N2);
+  fillWaymarks(A + N2, A + N2 + 1, N2);
+  verifyArrayValues(A, N2, N2 + 1);
+  EXPECT_EQ(T2, tag(A[N2]));
+
+  int T3 = tag(A[N3]);
+  delete ref(A[N3]);
+  A[N3] = new int(N3);
+  fillWaymarks(A + N3, A + N3 + 1, N3);
+  verifyArrayValues(A, N3, N3 + 1);
+  EXPECT_EQ(T3, tag(A[N3]));
+
+  int T4 = tag(A[N4]);
+  delete ref(A[N4]);
+  A[N4] = new int(N4);
+  fillWaymarks(A + N4, A + N4 + 1, N4);
+  verifyArrayValues(A, N4, N4 + 1);
+  EXPECT_EQ(T4, tag(A[N4]));
+
+  verifyArrayValues(A, N);
+  verifyFollowWaymarks(A, N);
+
+  freeArray(A, N);
+}
+
+} // end anonymous namespace