[LLVM][Support] Add getTrailingObjects() for single trailing type (#138970)

Add a specialization of getTrailingObjects() for a single trailing type.
This is a common case and with the specialization you don't need to
specify the single trailing type redundantly. Also add an overload for
getTrailingObjects which takes size and returns an
ArryaRef/MutableArrayRef as that's a common use case as well.

Author: jurahul

llvm/include/llvm/Support/TrailingObjects.h

@@ -46,6 +46,7 @@
 #ifndef LLVM_SUPPORT_TRAILINGOBJECTS_H
 #define LLVM_SUPPORT_TRAILINGOBJECTS_H
 
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/MathExtras.h"
@@ -301,6 +302,41 @@ class TrailingObjects : private trailing_objects_internal::TrailingObjectsImpl<
         static_cast<BaseTy *>(this), TrailingObjectsBase::OverloadToken<T>());
   }
 
+  // getTrailingObjects() specialization for a single trailing type.
+  using FirstTrailingType =
+      typename std::tuple_element_t<0, std::tuple<TrailingTys...>>;
+
+  const FirstTrailingType *getTrailingObjects() const {
+    static_assert(sizeof...(TrailingTys) == 1,
+                  "Can use non-templated getTrailingObjects() only when there "
+                  "is a single trailing type");
+    return getTrailingObjects<FirstTrailingType>();
+  }
+
+  FirstTrailingType *getTrailingObjects() {
+    static_assert(sizeof...(TrailingTys) == 1,
+                  "Can use non-templated getTrailingObjects() only when there "
+                  "is a single trailing type");
+    return getTrailingObjects<FirstTrailingType>();
+  }
+
+  // Functions that return the trailing objects as ArrayRefs.
+  template <typename T> MutableArrayRef<T> getTrailingObjects(size_t N) {
+    return MutableArrayRef(getTrailingObjects<T>(), N);
+  }
+
+  template <typename T> ArrayRef<T> getTrailingObjects(size_t N) const {
+    return ArrayRef(getTrailingObjects<T>(), N);
+  }
+
+  MutableArrayRef<FirstTrailingType> getTrailingObjects(size_t N) {
+    return MutableArrayRef(getTrailingObjects(), N);
+  }
+
+  ArrayRef<FirstTrailingType> getTrailingObjects(size_t N) const {
+    return ArrayRef(getTrailingObjects(), N);
+  }
+
   /// Returns the size of the trailing data, if an object were
   /// allocated with the given counts (The counts are in the same order
   /// as the template arguments). This does not include the size of the

llvm/unittests/Support/TrailingObjectsTest.cpp

@@ -26,7 +26,9 @@ class Class1 final : protected TrailingObjects<Class1, short> {
   size_t numTrailingObjects(OverloadToken<short>) const { return NumShorts; }
 
   Class1(ArrayRef<int> ShortArray) : NumShorts(ShortArray.size()) {
-    llvm::copy(ShortArray, getTrailingObjects<short>());
+    // This tests the non-templated getTrailingObjects() that returns a pointer
+    // when using a single trailing type.
+    llvm::copy(ShortArray, getTrailingObjects());
   }
 
 public:
@@ -36,7 +38,8 @@ class Class1 final : protected TrailingObjects<Class1, short> {
   }
   void operator delete(void *Ptr) { ::operator delete(Ptr); }
 
-  short get(unsigned Num) const { return getTrailingObjects<short>()[Num]; }
+  // This indexes into the ArrayRef<> returned by `getTrailingObjects`.
+  short get(unsigned Num) const { return getTrailingObjects(NumShorts)[Num]; }
 
   unsigned numShorts() const { return NumShorts; }
 
@@ -128,6 +131,9 @@ TEST(TrailingObjects, OneArg) {
   EXPECT_EQ(C->getTrailingObjects<short>(), reinterpret_cast<short *>(C + 1));
   EXPECT_EQ(C->get(0), 1);
   EXPECT_EQ(C->get(2), 3);
+
+  EXPECT_EQ(C->getTrailingObjects(), C->getTrailingObjects<short>());
+
   delete C;
 }