[llvm] [LLVM][ADT] Add `consume_front` member to ArrayRef (PR #146741)

[Rahul Joshi via llvm-commits]

https://github.com/jurahul updated https://github.com/llvm/llvm-project/pull/146741

>From 04af3182621f95c0be974d00dda054e49f939e7c Mon Sep 17 00:00:00 2001
From: Rahul Joshi <rjoshi at nvidia.com>
Date: Wed, 2 Jul 2025 08:56:57 -0700
Subject: [PATCH 1/2] [LLVM][ADT] Add `consume_front` member to ArrayRef

Add `consume_front` that returns the front element and drops it
from the current ArrayRef.
---
 llvm/include/llvm/ADT/ArrayRef.h    | 14 ++++++++++++++
 llvm/unittests/ADT/ArrayRefTest.cpp |  9 +++++++++
 2 files changed, 23 insertions(+)

diff --git a/llvm/include/llvm/ADT/ArrayRef.h b/llvm/include/llvm/ADT/ArrayRef.h
index 7fe43c99e34a7..93481f285e126 100644
--- a/llvm/include/llvm/ADT/ArrayRef.h
+++ b/llvm/include/llvm/ADT/ArrayRef.h
@@ -158,6 +158,13 @@ namespace llvm {
       return Data[Length-1];
     }
 
+    /// consume_front() - Returns the first element and drops it from ArrayRef.
+    const T &consume_front() {
+      const T &Ret = front();
+      *this = drop_front();
+      return Ret;
+    }
+
     // copy - Allocate copy in Allocator and return ArrayRef<T> to it.
     template <typename Allocator> MutableArrayRef<T> copy(Allocator &A) {
       T *Buff = A.template Allocate<T>(Length);
@@ -352,6 +359,13 @@ namespace llvm {
       return data()[this->size()-1];
     }
 
+    /// consume_front() - Returns the first element and drops it from ArrayRef.
+    T &consume_front() {
+      const T &Ret = front();
+      *this = drop_front();
+      return Ret;
+    }
+
     /// slice(n, m) - Chop off the first N elements of the array, and keep M
     /// elements in the array.
     MutableArrayRef<T> slice(size_t N, size_t M) const {
diff --git a/llvm/unittests/ADT/ArrayRefTest.cpp b/llvm/unittests/ADT/ArrayRefTest.cpp
index 985db1625454c..1f40b1e401c0b 100644
--- a/llvm/unittests/ADT/ArrayRefTest.cpp
+++ b/llvm/unittests/ADT/ArrayRefTest.cpp
@@ -104,6 +104,15 @@ TEST(ArrayRefTest, DropFront) {
   EXPECT_TRUE(AR1.drop_front(2).equals(AR2));
 }
 
+TEST(ArrayRefTest, ConsumeFront) {
+  static const int TheNumbers[] = {4, 8, 15, 16, 23, 42};
+  ArrayRef<int> AR1(TheNumbers);
+  ArrayRef<int> AR2(&TheNumbers[2], AR1.size() - 2);
+  EXPECT_EQ(AR1.consume_front(), 4);
+  EXPECT_EQ(AR1.consume_front(), 8);
+  EXPECT_TRUE(AR1.equals(AR2));
+}
+
 TEST(ArrayRefTest, DropWhile) {
   static const int TheNumbers[] = {1, 3, 5, 8, 10, 11};
   ArrayRef<int> AR1(TheNumbers);

>From de6e4888af3f5629921cea0060f774f396a87641 Mon Sep 17 00:00:00 2001
From: Rahul Joshi <rjoshi at nvidia.com>
Date: Wed, 2 Jul 2025 12:07:42 -0700
Subject: [PATCH 2/2] Review feedback

---
 llvm/include/llvm/ADT/ArrayRef.h    | 16 ++++++++++++-
 llvm/unittests/ADT/ArrayRefTest.cpp | 37 +++++++++++++++++++++++++++--
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/llvm/include/llvm/ADT/ArrayRef.h b/llvm/include/llvm/ADT/ArrayRef.h
index 93481f285e126..8091be2036642 100644
--- a/llvm/include/llvm/ADT/ArrayRef.h
+++ b/llvm/include/llvm/ADT/ArrayRef.h
@@ -165,6 +165,13 @@ namespace llvm {
       return Ret;
     }
 
+    /// consume_back() - Returns the last element and drops it from ArrayRef.
+    const T &consume_back() {
+      const T &Ret = back();
+      *this = drop_back();
+      return Ret;
+    }
+
     // copy - Allocate copy in Allocator and return ArrayRef<T> to it.
     template <typename Allocator> MutableArrayRef<T> copy(Allocator &A) {
       T *Buff = A.template Allocate<T>(Length);
@@ -361,11 +368,18 @@ namespace llvm {
 
     /// consume_front() - Returns the first element and drops it from ArrayRef.
     T &consume_front() {
-      const T &Ret = front();
+      T &Ret = front();
       *this = drop_front();
       return Ret;
     }
 
+    /// consume_back() - Returns the last element and drops it from ArrayRef.
+    T &consume_back() {
+      T &Ret = back();
+      *this = drop_back();
+      return Ret;
+    }
+
     /// slice(n, m) - Chop off the first N elements of the array, and keep M
     /// elements in the array.
     MutableArrayRef<T> slice(size_t N, size_t M) const {
diff --git a/llvm/unittests/ADT/ArrayRefTest.cpp b/llvm/unittests/ADT/ArrayRefTest.cpp
index 1f40b1e401c0b..b5fc0a100571e 100644
--- a/llvm/unittests/ADT/ArrayRefTest.cpp
+++ b/llvm/unittests/ADT/ArrayRefTest.cpp
@@ -108,11 +108,44 @@ TEST(ArrayRefTest, ConsumeFront) {
   static const int TheNumbers[] = {4, 8, 15, 16, 23, 42};
   ArrayRef<int> AR1(TheNumbers);
   ArrayRef<int> AR2(&TheNumbers[2], AR1.size() - 2);
-  EXPECT_EQ(AR1.consume_front(), 4);
-  EXPECT_EQ(AR1.consume_front(), 8);
+  EXPECT_EQ(&AR1.consume_front(), &TheNumbers[0]);
+  EXPECT_EQ(&AR1.consume_front(), &TheNumbers[1]);
   EXPECT_TRUE(AR1.equals(AR2));
 }
 
+TEST(ArrayRefTest, ConsumeBack) {
+  static const int TheNumbers[] = {4, 8, 15, 16, 23, 42};
+  ArrayRef<int> AR1(TheNumbers);
+  ArrayRef<int> AR2(TheNumbers, AR1.size() - 2);
+  EXPECT_EQ(&AR1.consume_back(), &TheNumbers[5]);
+  EXPECT_EQ(&AR1.consume_back(), &TheNumbers[4]);
+  EXPECT_TRUE(AR1.equals(AR2));
+}
+
+TEST(ArrayRefTest, MutableArryaRefConsumeFront) {
+  int TheNumbers[] = {4, 8, 15, 16, 23, 42};
+  MutableArrayRef<int> AR1(TheNumbers);
+  MutableArrayRef<int> AR2(&TheNumbers[2], AR1.size() - 2);
+  EXPECT_EQ(&AR1.consume_front(), &TheNumbers[0]);
+  EXPECT_EQ(&AR1.consume_front(), &TheNumbers[1]);
+  EXPECT_TRUE(AR1.equals(AR2));
+
+  AR1.consume_front() = 33;
+  EXPECT_EQ(TheNumbers[2], 33);
+}
+
+TEST(ArrayRefTest, MutableArryaRefConsumeBack) {
+  int TheNumbers[] = {4, 8, 15, 16, 23, 42};
+  MutableArrayRef<int> AR1(TheNumbers);
+  MutableArrayRef<int> AR2(TheNumbers, AR1.size() - 2);
+  EXPECT_EQ(&AR1.consume_back(), &TheNumbers[5]);
+  EXPECT_EQ(&AR1.consume_back(), &TheNumbers[4]);
+  EXPECT_TRUE(AR1.equals(AR2));
+
+  AR1.consume_back() = 33;
+  EXPECT_EQ(TheNumbers[3], 33);
+}
+
 TEST(ArrayRefTest, DropWhile) {
   static const int TheNumbers[] = {1, 3, 5, 8, 10, 11};
   ArrayRef<int> AR1(TheNumbers);