Change LockedBufferSpan to make it possible to create multiple spans from the same locked object.

PiperOrigin-RevId: 971629251

Author: qukhan

litert/tensor/buffer.h

@@ -54,37 +54,38 @@ class LockedBufferSpan {
   using iterator = T*;
   using const_iterator = const T*;
 
-  // `const std::byte*` if `T` is `const`, otherwise `std::byte`.
-  using MaybeConstByte =
-      std::conditional_t<std::is_const_v<T>, const std::byte, std::byte>;
-
-  LockedBufferSpan(MaybeConstByte* data,
-                   std::function<void(MaybeConstByte*)> unlock, size_t bytes)
-      : data_(data, std::move(unlock)), bytes_(bytes) {}
-
-  LockedBufferSpan(
-      std::unique_ptr<MaybeConstByte, std::function<void(MaybeConstByte*)>>
-          data,
-      size_t bytes)
-      : data_(std::move(data)), bytes_(bytes) {}
+  template <class Unlock>
+  LockedBufferSpan(T* data, Unlock&& unlock, size_t count)
+      : data_(data, std::forward<Unlock>(unlock)), bytes_(count * sizeof(T)) {}
+
+  template <class Y, class Deleter>
+  LockedBufferSpan(std::unique_ptr<Y, Deleter> data, size_t count)
+      : data_(data.release(), data.get_deleter()), bytes_(count * sizeof(T)) {}
 
   static LockedBufferSpan Empty() {
-    return LockedBufferSpan(nullptr, [](MaybeConstByte*) {}, 0);
+    return LockedBufferSpan(nullptr, [](T*) {}, 0);
   }
 
   // Casts the span to a specific type.
-  //
-  // Warning: This transfers the lock management to the returned
-  // `LockedBufferSpan`.
   template <class U>
-  [[nodiscard]] LockedBufferSpan<U> As() && {
+  [[nodiscard]] LockedBufferSpan<U> As() const {
     static_assert(
         std::is_const_v<U> || !std::is_const_v<T>,
         "Cannot cast from a constant buffer span to a non constant one.");
-    return LockedBufferSpan<U>(std::move(data_), bytes_);
+    return LockedBufferSpan<U>(data_, reinterpret_cast<U*>(data_.get()),
+                               bytes_ / sizeof(U));
+  }
+
+  [[nodiscard]] LockedBufferSpan SubSpan(size_t offset,
+                                         size_t count = SIZE_MAX) const {
+    if (offset >= size()) {
+      return Empty();
+    }
+    const size_t sub_count = std::min(count, size() - offset);
+    return LockedBufferSpan(data_, data_.get() + offset, sub_count);
   }
 
-  T* data() const& { return reinterpret_cast<T*>(data_.get()); }
+  T* data() const& { return data_.get(); }
   size_t size() const { return bytes_ / sizeof(T); }
   T* begin() & { return data(); }
   T* end() & { return data() + size(); }
@@ -102,8 +103,15 @@ class LockedBufferSpan {
   const T* cend() const&& = delete;
 
  private:
-  std::unique_ptr<MaybeConstByte, std::function<void(MaybeConstByte*)>> data_;
-  size_t bytes_;
+  template <class>
+  friend class LockedBufferSpan;
+
+  template <class U>
+  LockedBufferSpan(const std::shared_ptr<U>& data, T* ptr, size_t count)
+      : data_(data, ptr), bytes_(count * sizeof(T)) {}
+
+  std::shared_ptr<T> data_;
+  size_t bytes_ = 0;
 };
 
 // The main interface for buffers.

litert/tensor/buffer_test.cc

@@ -40,6 +40,7 @@ using ::litert::tensor::IsOk;
 using ::testing::Address;
 using ::testing::Contains;
 using ::testing::Each;
+using ::testing::ElementsAre;
 using ::testing::ElementsAreArray;
 using ::testing::Eq;
 using ::testing::Not;
@@ -82,6 +83,47 @@ TEST(LockedBufferSpanTest, CanBeCastAndUsedAsAContainer) {
               SizeIs(std::size(Lockable::kData)));
 }
 
+TEST(LockedBufferSpanTest, SharedOwnershipAcrossCopies) {
+  Lockable l;
+  ASSERT_EQ(l.i, 0);
+  {
+    LockedBufferSpan<std::byte> span1 = l.LockMutable();
+    ASSERT_EQ(l.i, 1);
+    {
+      LockedBufferSpan<std::byte> span2 = span1;
+      LockedBufferSpan<int> span3 = span1.As<int>();
+      ASSERT_EQ(l.i, 1);
+    }
+    ASSERT_EQ(l.i, 1);
+  }
+  ASSERT_EQ(l.i, 0);
+}
+
+TEST(LockedBufferSpanTest, SubSpan) {
+  const int data[] = {10, 20, 30, 40, 50};
+  LockedBufferSpan<const int> span(data, [](const int*) {}, std::size(data));
+
+  LockedBufferSpan<const int> sub1 = span.SubSpan(1, 3);
+  EXPECT_THAT(sub1, ElementsAre(20, 30, 40));
+  EXPECT_EQ(sub1.size(), 3);
+
+  LockedBufferSpan<const int> sub2 = span.SubSpan(2);
+  EXPECT_THAT(sub2, ElementsAre(30, 40, 50));
+  EXPECT_EQ(sub2.size(), 3);
+
+  LockedBufferSpan<const int> sub_oob = span.SubSpan(10);
+  EXPECT_EQ(sub_oob.size(), 0);
+}
+
+TEST(LockedBufferSpanTest, ConstructFromUniquePtr) {
+  auto ptr = std::make_unique<int[]>(3);
+  ptr[0] = 100;
+  ptr[1] = 200;
+  ptr[2] = 300;
+  LockedBufferSpan<int> span(std::move(ptr), 3);
+  EXPECT_THAT(span, ElementsAre(100, 200, 300));
+}
+
 TEST(SpanCpuBufferTest, BuildFromRawData) {
   const int32_t backing_array[] = {1, 2, 3, 4, 5};
   SpanCpuBuffer bv(