[STLExtras] Add take_begin and take_end range adaptors.

[vitalybuka]

These are similar to drop_begin and drop_end but they return the
first/last N elements instead of dropping them.

[llvmbot]

@llvm/pr-subscribers-llvm-adt

Author: Vitaly Buka (vitalybuka)

Changes

These are similar to drop_begin and drop_end but they return the
first/last N elements instead of dropping them.

---

Full diff: https://github.com/llvm/llvm-project/pull/164542.diff

2 Files Affected:

• (modified) llvm/include/llvm/ADT/STLExtras.h (+14)
• (modified) llvm/unittests/ADT/STLExtrasTest.cpp (+29)

diff --git a/llvm/include/llvm/ADT/STLExtras.h b/llvm/include/llvm/ADT/STLExtras.h
index a9841c6651b72..1cbc1d051aa18 100644
--- a/llvm/include/llvm/ADT/STLExtras.h
+++ b/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.
 
diff --git a/llvm/unittests/ADT/STLExtrasTest.cpp b/llvm/unittests/ADT/STLExtrasTest.cpp
index 966b1f01e8a31..ad71f8068366c 100644
--- a/llvm/unittests/ADT/STLExtrasTest.cpp
+++ b/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; }),

[kuhar]

Having these makes sense to me, I was just looking for them the other week

[kuhar]

Can you clarify what happens when the range has fewer than N elements?

[kuhar]

This is fine for random access iterators, but for forward or input iterators we should not be traversing the container twice. Instead, it would be better to have an iterator type that maintains a counter.
(Note that drop_begin doesn't suffer from the same issue because we don't come back to the front of the range.)

[kuhar]

Also here, specify what happens when the range has fewer elements than N

[kuhar]

nit: query end only once

Otherwise this implementation seems fine -- I don't think we can do much better

[kuhar]

nit: I would combine this with the previous function

[kuhar]

also here

[kuhar]

Can you add a test case with an input iterator?

[kuhar]

Can you add a testcase with a bidirectional iterator?

[kuhar]

can we make this constexpr?

[kuhar]

Can we make this constexpr?

[vitalybuka]

@kuhar Thanks for these comments, they make sense. But it's done in a way consistent to drop_begin/drop_end.

So I think we should either land it more or less as is, and fix both together, or start with fixing drop_begin/drop_end in a separate patch, and then update this one.

WDYT?

[kuhar]

I'm fine with the latter options -- cleaning up drop_begin/end should be very quick since there's no issue with input iterators. That code probably predates std::next/prev being constexpr. I don't think this cleanup should block this PR though.