From 53ffd56d87a2590182e1b8c2535c5824ad77462c Mon Sep 17 00:00:00 2001
From: Denis Yaroshevskiy <dyaroshev@fb.com>
Date: Wed, 14 Aug 2024 06:42:43 -0700
Subject: [PATCH] folly::span_cast

Summary: a simple utility to convert span types to a different span.

Differential Revision: D61202580
---
 folly/BUCK                |   1 +
 folly/Memory.h            |  57 +++++++++++++++++++
 folly/Traits.h            |   1 +
 folly/test/MemoryTest.cpp | 116 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 175 insertions(+)

diff --git a/folly/BUCK b/folly/BUCK
index de765471a05..4eef0773a45 100644
--- a/folly/BUCK
+++ b/folly/BUCK
@@ -915,6 +915,7 @@ cpp_library(
         "//folly/lang:thunk",
         "//folly/memory:malloc",
         "//folly/portability:config",
+        "//folly/portability:constexpr",
         "//folly/portability:malloc",
     ],
 )
diff --git a/folly/Memory.h b/folly/Memory.h
index 91a7533afb5..260131b3ce4 100644
--- a/folly/Memory.h
+++ b/folly/Memory.h
@@ -38,8 +38,13 @@
 #include <folly/lang/Thunk.h>
 #include <folly/memory/Malloc.h>
 #include <folly/portability/Config.h>
+#include <folly/portability/Constexpr.h>
 #include <folly/portability/Malloc.h>
 
+#if __cpp_lib_span
+#include <span>
+#endif
+
 namespace folly {
 
 #if (defined(_POSIX_C_SOURCE) && _POSIX_C_SOURCE >= 200112L) || \
@@ -845,4 +850,56 @@ template <typename Value, typename T>
 struct AllocatorHasDefaultObjectDestroy<std::allocator<Value>, T>
     : std::true_type {};
 
+#if __cpp_lib_span
+
+namespace detail {
+
+template <typename U, typename T, std::size_t Extend>
+constexpr auto span_cast_impl(std::span<T, Extend> in, U* castData) {
+  // check alignment
+  if (!folly::is_constant_evaluated_or(true)) {
+    assert(reinterpret_cast<std::uintptr_t>(in.data()) % sizeof(U) == 0);
+  }
+
+  if constexpr (Extend == std::dynamic_extent) {
+    assert(in.size() * sizeof(T) % sizeof(U) == 0);
+    return std::span<U>(castData, in.size() * sizeof(T) / sizeof(U));
+  } else {
+    static_assert(in.size() * sizeof(T) % sizeof(U) == 0);
+    constexpr std::size_t kResSize = Extend * sizeof(T) / sizeof(U);
+    return std::span<U, kResSize>(castData, kResSize);
+  }
+}
+
+} // namespace detail
+
+/**
+ * converting a span to a different span.
+ * (you get a span to the same bytes but treated as different type)
+ *
+ * Example:
+ *
+ *    enum class SomeEnum : int { ... };
+ *
+ *    std::span<SomeEnum> s = ...
+ *    std::span<int> as_ints = folly::reinterpret_span_cast<int>(s);
+ */
+
+template <typename U, typename T, std::size_t Extend>
+constexpr auto static_span_cast(std::span<T, Extend> in) {
+  return detail::span_cast_impl(in, static_cast<U*>(in.data()));
+}
+
+template <typename U, typename T, std::size_t Extend>
+constexpr auto reinterpret_span_cast(std::span<T, Extend> in) {
+  return detail::span_cast_impl(in, reinterpret_cast<U*>(in.data()));
+}
+
+template <typename U, typename T, std::size_t Extend>
+constexpr auto const_span_cast(std::span<T, Extend> in) {
+  return detail::span_cast_impl(in, const_cast<U*>(in.data()));
+}
+
+#endif
+
 } // namespace folly
diff --git a/folly/Traits.h b/folly/Traits.h
index 6ae5c064c49..3474b827353 100644
--- a/folly/Traits.h
+++ b/folly/Traits.h
@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <cassert>
 #include <cstdint>
 #include <functional>
 #include <limits>
diff --git a/folly/test/MemoryTest.cpp b/folly/test/MemoryTest.cpp
index e23ed677604..7c660521a15 100644
--- a/folly/test/MemoryTest.cpp
+++ b/folly/test/MemoryTest.cpp
@@ -585,3 +585,119 @@ TEST(allocateOverAligned, defaultOverCustomAlloc) {
   folly::deallocateOverAligned(a, p, 1);
   EXPECT_EQ(folly::allocationBytesForOverAligned<decltype(a)>(1), 128);
 }
+
+#if __cpp_lib_span
+
+template <typename To, typename From, std::size_t Extend>
+using reinterpret_span_cast_result_type =
+    decltype(folly::reinterpret_span_cast<To>(
+        std::declval<std::span<From, Extend>>()));
+
+template <typename To, typename From, std::size_t Extend>
+using static_span_cast_result_type = decltype(folly::const_span_cast<To>(
+    std::declval<std::span<From, Extend>>()));
+
+template <typename To, typename From, std::size_t Extend>
+using const_span_cast_result_type = decltype(folly::const_span_cast<To>(
+    std::declval<std::span<From, Extend>>()));
+
+TEST(Traits, SpanCast) {
+  auto tstSpanCast = [](auto to, auto from) {
+    ASSERT_EQ(
+        static_cast<const void*>(from.data()),
+        static_cast<const void*>(to.data()));
+
+    ASSERT_EQ(
+        static_cast<const void*>(from.data() + from.size()),
+        static_cast<const void*>(to.data() + to.size()));
+  };
+
+  {
+    std::array<int, 4> a;
+    tstSpanCast(
+        folly::reinterpret_span_cast<const char>(std::span(a)), std::span(a));
+    tstSpanCast(
+        folly::reinterpret_span_cast<double>(std::span(a)), std::span(a));
+  }
+
+  {
+    std::vector<int> a(4u, 1);
+    tstSpanCast(
+        folly::reinterpret_span_cast<const char>(std::span(a)), std::span(a));
+    tstSpanCast(
+        folly::reinterpret_span_cast<double>(std::span(a)), std::span(a));
+  }
+
+  {
+    const std::vector<int> a(4u, 1);
+    tstSpanCast(folly::const_span_cast<int>(std::span(a)), std::span(a));
+
+    std::vector<int> b(4u, 1);
+    tstSpanCast(folly::static_span_cast<const int>(std::span(b)), std::span(b));
+    tstSpanCast(folly::const_span_cast<const int>(std::span(b)), std::span(b));
+    tstSpanCast(
+        folly::reinterpret_span_cast<const int>(std::span(b)), std::span(b));
+  }
+
+  // types
+  {
+    static_assert(
+        std::is_same_v<
+            std::span<char>,
+            reinterpret_span_cast_result_type<char, int, std::dynamic_extent>>);
+    static_assert(std::is_same_v<
+                  std::span<const char>,
+                  reinterpret_span_cast_result_type<
+                      const char,
+                      const int,
+                      std::dynamic_extent>>);
+    static_assert(std::is_same_v<
+                  std::span<const char, 12>,
+                  reinterpret_span_cast_result_type<const char, const int, 3>>);
+
+    static_assert(std::is_same_v<
+                  std::span<const char, 12>,
+                  const_span_cast_result_type<const char, char, 12>>);
+    static_assert(std::is_same_v<
+                  std::span<char, 12>,
+                  const_span_cast_result_type<char, const char, 12>>);
+    static_assert(std::is_same_v<
+                  std::span<const char>,
+                  const_span_cast_result_type<
+                      const char,
+                      char,
+                      std::dynamic_extent>>);
+    static_assert(std::is_same_v<
+                  std::span<char>,
+                  const_span_cast_result_type<
+                      char,
+                      const char,
+                      std::dynamic_extent>>);
+
+    static_assert(std::is_same_v<
+                  std::span<const char, 12>,
+                  static_span_cast_result_type<const char, char, 12>>);
+    static_assert(std::is_same_v<
+                  std::span<const char>,
+                  static_span_cast_result_type<
+                      const char,
+                      char,
+                      std::dynamic_extent>>);
+  }
+
+  // constexpr
+  {
+    [[maybe_unused]] constexpr auto _ = [] {
+      std::array<int, 4> a{0, 1, 2, 3};
+      std::span<int, 4> mutableAFixed(a);
+      std::span<int> mutableADynamic(a);
+      auto resFixed = folly::static_span_cast<const int>(mutableAFixed);
+      (void)resFixed;
+      auto resDynamic = folly::static_span_cast<const int>(mutableADynamic);
+      (void)resDynamic;
+      return 0;
+    }();
+  }
+}
+
+#endif