diff --git a/llvm/include/llvm/ADT/STLExtras.h b/llvm/include/llvm/ADT/STLExtras.h index af0e4a36be1b1..3c18aae1b519c 100644 --- a/llvm/include/llvm/ADT/STLExtras.h +++ b/llvm/include/llvm/ADT/STLExtras.h @@ -1415,6 +1415,18 @@ template auto make_second_range(ContainerTy &&c) { }); } +/// Return a range that conditionally reverses \p C. The collection is iterated +/// in reverse if \p ShouldReverse is true (otherwise, it is iterated forwards). +template +[[nodiscard]] auto reverse_conditionally(ContainerTy &&C, bool ShouldReverse) { + using IterTy = detail::IterOfRange; + using ReferenceTy = typename std::iterator_traits::reference; + return map_range(zip_equal(reverse(C), C), + [ShouldReverse](auto I) -> ReferenceTy { + return ShouldReverse ? std::get<0>(I) : std::get<1>(I); + }); +} + //===----------------------------------------------------------------------===// // Extra additions to //===----------------------------------------------------------------------===// diff --git a/llvm/unittests/ADT/STLExtrasTest.cpp b/llvm/unittests/ADT/STLExtrasTest.cpp index 85567775e4ebd..e356f6b540568 100644 --- a/llvm/unittests/ADT/STLExtrasTest.cpp +++ b/llvm/unittests/ADT/STLExtrasTest.cpp @@ -1693,6 +1693,30 @@ TEST(STLExtrasTest, ProductOf) { EXPECT_EQ(product_of(V3), 2.0f); } +TEST(STLExtrasTest, ReverseConditionally) { + std::vector foo = {'a', 'b', 'c'}; + + // Test backwards. + std::vector ReverseResults; + for (char Value : llvm::reverse_conditionally(foo, /*ShouldReverse=*/true)) { + ReverseResults.emplace_back(Value); + } + EXPECT_THAT(ReverseResults, ElementsAre('c', 'b', 'a')); + + // Test forwards. + std::vector ForwardResults; + for (char Value : llvm::reverse_conditionally(foo, /*ShouldReverse=*/false)) { + ForwardResults.emplace_back(Value); + } + EXPECT_THAT(ForwardResults, ElementsAre('a', 'b', 'c')); + + // Test modifying collection. + for (char &Value : llvm::reverse_conditionally(foo, /*ShouldReverse=*/true)) { + ++Value; + } + EXPECT_THAT(foo, ElementsAre('b', 'c', 'd')); +} + struct Foo; struct Bar {};