From 7b2eefb3fd2c94223e04c90d071bb0eba90d4e24 Mon Sep 17 00:00:00 2001
From: Yedidya Feldblum <yfeldblum@meta.com>
Date: Wed, 4 Dec 2024 13:31:17 -0800
Subject: [PATCH] atomic_grow_array::as_ptr_span

Summary:
A facility like `atomic_grow_array::as_view` but which returns a `folly::span` which, if `std::span` is available, is an alias to `std::span`. This allows using, for example, `span::subspan`, or calling any function which takes `span` for deduced template-type.

Note that this does expose the internal array and is the only member which does. `atomic_grow_array::view`, for example, does not expose the internal array.

Differential Revision: D66737577

fbshipit-source-id: ae92e3653ac575dc314a067e715a1fb90b8a27b3
---
 folly/concurrency/container/BUCK              |  1 +
 .../concurrency/container/atomic_grow_array.h | 25 +++++++++++++++++++
 .../container/test/atomic_grow_array_test.cpp | 12 +++++++++
 3 files changed, 38 insertions(+)

diff --git a/folly/concurrency/container/BUCK b/folly/concurrency/container/BUCK
index 223b8f7e1cb..477d4272eb5 100644
--- a/folly/concurrency/container/BUCK
+++ b/folly/concurrency/container/BUCK
@@ -12,6 +12,7 @@ cpp_library(
         "//folly:constexpr_math",
         "//folly:likely",
         "//folly:scope_guard",
+        "//folly/container:span",
         "//folly/lang:align",
         "//folly/lang:bits",
         "//folly/lang:new",
diff --git a/folly/concurrency/container/atomic_grow_array.h b/folly/concurrency/container/atomic_grow_array.h
index 88bc118b379..33db58d7423 100644
--- a/folly/concurrency/container/atomic_grow_array.h
+++ b/folly/concurrency/container/atomic_grow_array.h
@@ -25,6 +25,7 @@
 #include <folly/ConstexprMath.h>
 #include <folly/Likely.h>
 #include <folly/ScopeGuard.h>
+#include <folly/container/span.h>
 #include <folly/lang/Align.h>
 #include <folly/lang/Bits.h>
 #include <folly/lang/New.h>
@@ -80,6 +81,9 @@ class atomic_grow_array : private Policy {
   using size_type = std::size_t;
   using value_type = Item;
 
+  using pointer_span = span<value_type* const>;
+  using const_pointer_span = span<value_type const* const>;
+
   class iterator;
   class const_iterator;
 
@@ -206,6 +210,8 @@ class atomic_grow_array : private Policy {
     using size_type = typename atomic_grow_array::size_type;
     using reference = maybe_add_const_t<value_type>&;
     using const_reference = value_type const&;
+    using pointer = maybe_add_const_t<value_type>*;
+    using const_pointer = value_type const*;
     using iterator = conditional_t<Const, up::const_iterator, up::iterator>;
     using const_iterator = up::const_iterator;
 
@@ -227,6 +233,15 @@ class atomic_grow_array : private Policy {
     const_reference operator[](size_type index) const noexcept {
       return *array_->list[index];
     }
+
+    span<pointer const> as_ptr_span() noexcept {
+      using type = span<pointer const>;
+      return array_ ? type{array_->list, array_->size} : type{};
+    }
+    span<const_pointer const> as_ptr_span() const noexcept {
+      using type = span<const_pointer const>;
+      return array_ ? type{array_->list, array_->size} : type{};
+    }
   };
 
  public:
@@ -361,6 +376,7 @@ class atomic_grow_array : private Policy {
     using typename base::size_type;
     using typename base::value_type;
 
+    using base::as_ptr_span;
     using base::base;
     using base::begin;
     using base::cbegin;
@@ -390,6 +406,7 @@ class atomic_grow_array : private Policy {
     using typename base::size_type;
     using typename base::value_type;
 
+    using base::as_ptr_span;
     using base::base;
     using base::begin;
     using base::cbegin;
@@ -425,6 +442,14 @@ class atomic_grow_array : private Policy {
   view as_view() noexcept { return view{array_.load(mo_acquire)}; }
   const_view as_view() const noexcept { return view{array_.load(mo_acquire)}; }
 
+  /// as_ptr_span
+  ///
+  /// Convenience wrapper for view::as_ptr_span.
+  pointer_span as_ptr_span() noexcept { return as_view().as_ptr_span(); }
+  const_pointer_span as_ptr_span() const noexcept {
+    return as_view().as_ptr_span();
+  }
+
  private:
   static constexpr auto mo_acquire = std::memory_order_acquire;
   static constexpr auto mo_release = std::memory_order_release;
diff --git a/folly/concurrency/container/test/atomic_grow_array_test.cpp b/folly/concurrency/container/test/atomic_grow_array_test.cpp
index 479177ce694..863f64d3fb2 100644
--- a/folly/concurrency/container/test/atomic_grow_array_test.cpp
+++ b/folly/concurrency/container/test/atomic_grow_array_test.cpp
@@ -127,21 +127,33 @@ TEST_F(AtomicGrowArrayTest, empty) {
   auto const count = [](auto& v) {
     return std::accumulate(v.begin(), v.end(), 0);
   };
+  auto const countp = [](auto& v) {
+    auto const sum = [](auto a, auto* b) { return a + *b; };
+    return std::accumulate(v.begin(), v.end(), 0, sum);
+  };
   {
     auto view = array.as_view();
     EXPECT_EQ(expected, count(view));
+    auto span = array.as_ptr_span();
+    EXPECT_EQ(expected, countp(span));
   }
   {
     auto const view = array.as_view();
     EXPECT_EQ(expected, count(view));
+    auto const span = array.as_ptr_span();
+    EXPECT_EQ(expected, countp(span));
   }
   {
     auto view = std::as_const(array).as_view();
     EXPECT_EQ(expected, count(view));
+    auto span = std::as_const(array).as_ptr_span();
+    EXPECT_EQ(expected, countp(span));
   }
   {
     auto const view = std::as_const(array).as_view();
     EXPECT_EQ(expected, count(view));
+    auto const span = std::as_const(array).as_ptr_span();
+    EXPECT_EQ(expected, countp(span));
   }
 }