[STLExtras] Add take_begin and take_end range adaptors.

llvm/include/llvm/ADT/STLExtras.h

@@ -325,6 +325,20 @@ template <typename T> auto drop_end(T &&RangeOrContainer, size_t N = 1) {
                     std::prev(adl_end(RangeOrContainer), N));
 }
 
+/// Return a range covering \p RangeOrContainer with the first N elements
+/// included.
+template <typename T> auto take_begin(T &&RangeOrContainer, size_t N = 1) {
+  return make_range(adl_begin(RangeOrContainer),
+                    std::next(adl_begin(RangeOrContainer), N));
+}
+
+/// Return a range covering \p RangeOrContainer with the last N elements
+/// included.
+template <typename T> auto take_end(T &&RangeOrContainer, size_t N = 1) {
+  return make_range(std::prev(adl_end(RangeOrContainer), N),
+                    adl_end(RangeOrContainer));
+}
+
 // mapped_iterator - This is a simple iterator adapter that causes a function to
 // be applied whenever operator* is invoked on the iterator.
 

llvm/unittests/ADT/STLExtrasTest.cpp

@@ -800,6 +800,35 @@ TEST(STLExtrasTest, DropEndDefaultTest) {
   EXPECT_THAT(drop_end(vec), ElementsAre(0, 1, 2, 3));
 }
 
+TEST(STLExtrasTest, TakeBeginTest) {
+  SmallVector<int, 5> vec{0, 1, 2, 3, 4};
+
+  for (int n = 0; n < 5; ++n) {
+    EXPECT_THAT(take_begin(vec, n), ElementsAreArray(ArrayRef(vec.data(), n)));
+  }
+}
+
+TEST(STLExtrasTest, TakeBeginDefaultTest) {
+  SmallVector<int, 5> vec{0, 1, 2, 3, 4};
+
+  EXPECT_THAT(take_begin(vec), ElementsAre(0));
+}
+
+TEST(STLExtrasTest, TakeEndTest) {
+  SmallVector<int, 5> vec{0, 1, 2, 3, 4};
+
+  for (int n = 0; n < 5; ++n) {
+    EXPECT_THAT(take_end(vec, n),
+                ElementsAreArray(ArrayRef(&vec[vec.size() - n], n)));
+  }
+}
+
+TEST(STLExtrasTest, TakeEndDefaultTest) {
+  SmallVector<int, 5> vec{0, 1, 2, 3, 4};
+
+  EXPECT_THAT(take_end(vec), ElementsAre(4));
+}
+
 TEST(STLExtrasTest, MapRangeTest) {
   SmallVector<int, 5> Vec{0, 1, 2};
   EXPECT_THAT(map_range(Vec, [](int V) { return V + 1; }),