Make buffer manager context getter return a mutable reference.

PiperOrigin-RevId: 720347065
GesuBackups · Jan 28, 2025 · 799ada6 · 799ada6
1 parent ed6f659
commit 799ada6
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/tensorflow/lite/experimental/litert/core/model/buffer_manager.h b/tensorflow/lite/experimental/litert/core/model/buffer_manager.h
@@ -17,6 +17,7 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <functional>
 #include <optional>
 #include <utility>
 #include <variant>
@@ -30,6 +31,8 @@ namespace litert::internal {
 
 // Extra info about how the buffer is handled during load or serialization.
 struct BufferContext {
+  using Ref = std::reference_wrapper<BufferContext>;
+
   // Whether the buffer should be appended to the flatbuffer during
   // serialization.
   bool should_append = false;
@@ -71,11 +74,11 @@ class BufferManager {
   }
 
   // Get the context of the buffer at the given id.
-  Expected<BufferContext> GetContext(BufferId id) {
+  Expected<BufferContext::Ref> GetContext(BufferId id) {
     if (id >= buffers_.size()) {
       return Error(kLiteRtStatusErrorIndexOOB);
     }
-    return buffers_[id].second;
+    return std::ref(buffers_[id].second);
   }
 
   // Number of buffers. Ids will be 0 <-> num - 1.

diff --git a/tensorflow/lite/experimental/litert/core/model/buffer_manager_test.cc b/tensorflow/lite/experimental/litert/core/model/buffer_manager_test.cc
@@ -63,7 +63,7 @@ TEST(BufferManagerTest, RegisterWithContext) {
 
   EXPECT_EQ(manager.NumBuffers(), 2);
   EXPECT_EQ(manager.GetBuffer(id)->StrView(), kData);
-  EXPECT_EQ(manager.GetContext(id)->should_append, true);
+  EXPECT_EQ(manager.GetContext(id)->get().should_append, true);
 }
 
 }  // namespace