From 61819f03205556ca7d83cb13ec402849a45f4894 Mon Sep 17 00:00:00 2001
From: Boris Zbarsky <bzbarsky@apple.com>
Date: Thu, 4 May 2023 10:12:49 -0400
Subject: [PATCH] Add element accessors on Span.

This matches std::span and avoids people having to write span.data()[index] and
whatnot.
---
 src/lib/support/Span.h             | 24 +++++++++++++--
 src/lib/support/tests/TestSpan.cpp | 48 ++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 2 deletions(-)

diff --git a/src/lib/support/Span.h b/src/lib/support/Span.h
index 4d8f521eb35c0c..4ba15145f07558 100644
--- a/src/lib/support/Span.h
+++ b/src/lib/support/Span.h
@@ -38,7 +38,8 @@ template <class T>
 class Span
 {
 public:
-    using pointer = T *;
+    using pointer   = T *;
+    using reference = T &;
 
     constexpr Span() : mDataBuf(nullptr), mDataLen(0) {}
     constexpr Span(pointer databuf, size_t datalen) : mDataBuf(databuf), mDataLen(datalen) {}
@@ -75,6 +76,15 @@ class Span
     constexpr pointer begin() const { return data(); }
     constexpr pointer end() const { return data() + size(); }
 
+    // Element accessors, matching the std::span API.
+    constexpr reference operator[](size_t index) const
+    {
+        VerifyOrDie(index < size());
+        return data()[index];
+    }
+    constexpr reference front() const { return (*this)[0]; }
+    constexpr reference back() const { return (*this)[size() - 1]; }
+
     template <class U, typename = std::enable_if_t<std::is_same<std::remove_const_t<T>, std::remove_const_t<U>>::value>>
     bool data_equal(const Span<U> & other) const
     {
@@ -142,7 +152,8 @@ template <class T, size_t N>
 class FixedSpan
 {
 public:
-    using pointer = T *;
+    using pointer   = T *;
+    using reference = T &;
 
     constexpr FixedSpan() : mDataBuf(nullptr) {}
 
@@ -188,6 +199,15 @@ class FixedSpan
     constexpr pointer begin() const { return mDataBuf; }
     constexpr pointer end() const { return mDataBuf + N; }
 
+    // Element accessors, matching the std::span API.
+    constexpr reference operator[](size_t index) const
+    {
+        VerifyOrDie(index < size());
+        return data()[index];
+    }
+    constexpr reference front() const { return (*this)[0]; }
+    constexpr reference back() const { return (*this)[size() - 1]; }
+
     // Allow data_equal for spans that are over the same type up to const-ness.
     template <class U, typename = std::enable_if_t<std::is_same<std::remove_const_t<T>, std::remove_const_t<U>>::value>>
     bool data_equal(const FixedSpan<U, N> & other) const
diff --git a/src/lib/support/tests/TestSpan.cpp b/src/lib/support/tests/TestSpan.cpp
index 06a9a57070bc02..a7bf8d618afd55 100644
--- a/src/lib/support/tests/TestSpan.cpp
+++ b/src/lib/support/tests/TestSpan.cpp
@@ -56,6 +56,30 @@ static void TestByteSpan(nlTestSuite * inSuite, void * inContext)
     NL_TEST_ASSERT(inSuite, s2.data_equal(s2));
     NL_TEST_ASSERT(inSuite, !s2.data_equal(s1));
     NL_TEST_ASSERT(inSuite, IsSpanUsable(s2) == true);
+    NL_TEST_ASSERT(inSuite, s2.front() == 1);
+    NL_TEST_ASSERT(inSuite, s2.back() == 3);
+    NL_TEST_ASSERT(inSuite, s2[0] == 1);
+    NL_TEST_ASSERT(inSuite, s2[1] == 2);
+    NL_TEST_ASSERT(inSuite, s2[2] == 3);
+
+    // Test that for a constexpr span we can use its data as compile-time
+    // constants.
+    {
+        static constexpr uint8_t constArr[] = { 2, 3, 4 };
+        constexpr ByteSpan constSpan(constArr);
+
+        static_assert(constSpan.front() == 2, "Unexpected front()");
+        uint8_t t1[constSpan.front()];
+        NL_TEST_ASSERT(inSuite, ArraySize(t1) == 2);
+
+        static_assert(constSpan.back() == 4, "Unexpected back()");
+        uint8_t t2[constSpan.back()];
+        NL_TEST_ASSERT(inSuite, ArraySize(t2) == 4);
+
+        static_assert(constSpan[1] == 3, "Unexpected [1]");
+        uint8_t t3[constSpan[1]];
+        NL_TEST_ASSERT(inSuite, ArraySize(t3) == 3);
+    }
 
     ByteSpan s3 = s2;
     NL_TEST_ASSERT(inSuite, s3.data() == arr);
@@ -175,6 +199,30 @@ static void TestFixedByteSpan(nlTestSuite * inSuite, void * inContext)
     NL_TEST_ASSERT(inSuite, s2.data()[2] == 3);
     NL_TEST_ASSERT(inSuite, !s2.empty());
     NL_TEST_ASSERT(inSuite, s2.data_equal(s2));
+    NL_TEST_ASSERT(inSuite, s2.front() == 1);
+    NL_TEST_ASSERT(inSuite, s2.back() == 3);
+    NL_TEST_ASSERT(inSuite, s2[0] == 1);
+    NL_TEST_ASSERT(inSuite, s2[1] == 2);
+    NL_TEST_ASSERT(inSuite, s2[2] == 3);
+
+    // Test that for a constexpr span we can use its data as compile-time
+    // constants.
+    {
+        static constexpr uint8_t constArr[] = { 2, 3, 4 };
+        constexpr FixedByteSpan<ArraySize(constArr)> constSpan(constArr);
+
+        static_assert(constSpan.front() == 2, "Unexpected front()");
+        uint8_t t1[constSpan.front()];
+        NL_TEST_ASSERT(inSuite, ArraySize(t1) == 2);
+
+        static_assert(constSpan.back() == 4, "Unexpected back()");
+        uint8_t t2[constSpan.back()];
+        NL_TEST_ASSERT(inSuite, ArraySize(t2) == 4);
+
+        static_assert(constSpan[1] == 3, "Unexpected [1]");
+        uint8_t t3[constSpan[1]];
+        NL_TEST_ASSERT(inSuite, ArraySize(t3) == 3);
+    }
 
     FixedByteSpan<3> s3 = s2;
     NL_TEST_ASSERT(inSuite, s3.data() == arr);