Generalize the FlatBufferFile methods that take ownership of a string so that they can take any string or STL container of byte sized elements.

PiperOrigin-RevId: 272538399

iree/base/flatbuffer_util.cc

diff --git a/iree/base/flatbuffer_util.cc b/iree/base/flatbuffer_util.cc
index d4722ff..7ce44f4 100644
--- a/iree/base/flatbuffer_util.cc
+++ b/iree/base/flatbuffer_util.cc
@@ -51,7 +51,7 @@
 
   // Pass along the buffer provided so we keep it alive until the
   // FlatBufferFileBase is destructed.
-  auto backing_buffer_baton = MoveToLambda(backing_buffer);
+  auto backing_buffer_baton = IreeMoveToLambda(backing_buffer);
   deleter_ = [backing_buffer_baton]() { (void)backing_buffer_baton.value; };
 
   return OkStatus();
@@ -124,28 +124,6 @@
       identifier, buffer_data, []() {}, root_type_size, verifier_fn);
 }
 
-Status FlatBufferFileBase::FromString(Identifier identifier,
-                                      std::string buffer_data,
-                                      size_t root_type_size,
-                                      VerifierFn verifier_fn) {
-  IREE_TRACE_SCOPE0("FlatBufferFileBase::FromString");
-
-  // Reference right into the string buffer.
-  auto buffer_data_data = absl::MakeConstSpan(
-      reinterpret_cast<const uint8_t*>(buffer_data.data()), buffer_data.size());
-
-  // Use a baton to keep the string alive until the FlatBufferFileBase is
-  // destroyed.
-  auto buffer_data_baton = MoveToLambda(buffer_data);
-  return FromBuffer(
-      identifier, buffer_data_data,
-      [buffer_data_baton]() {
-        // Keeping the string alive.
-        (void)buffer_data_baton.value;
-      },
-      root_type_size, verifier_fn);
-}
-
 Status FlatBufferFileBase::LoadFile(Identifier identifier, std::string path,
                                     size_t root_type_size,
                                     VerifierFn verifier_fn) {
@@ -154,7 +132,7 @@
   ASSIGN_OR_RETURN(auto file_mapping, FileMapping::OpenRead(path));
   auto buffer_data = file_mapping->data();
 
-  auto handle_baton = MoveToLambda(file_mapping);
+  auto handle_baton = IreeMoveToLambda(file_mapping);
   return FromBuffer(
       identifier, buffer_data,
       [handle_baton]() {

iree/base/flatbuffer_util.h

diff --git a/iree/base/flatbuffer_util.h b/iree/base/flatbuffer_util.h
index 93ce144..0e9344f 100644
--- a/iree/base/flatbuffer_util.h
+++ b/iree/base/flatbuffer_util.h
@@ -27,6 +27,7 @@
 #include "absl/types/optional.h"
 #include "absl/types/span.h"
 #include "flatbuffers/flatbuffers.h"
+#include "iree/base/memory.h"
 #include "iree/base/status.h"
 
 namespace iree {
@@ -66,11 +67,14 @@
                     absl::Span<const uint8_t> buffer_data,
                     std::function<void()> deleter, size_t root_type_size,
                     VerifierFn verifier_fn);
+  // Initializes from an STL byte based container (string and vector of
+  // char/byte should be compatible).
+  template <typename Container>
+  Status FromContainer(Identifier identifier, Container container,
+                       size_t root_type_size, VerifierFn verifier_fn);
   Status WrapBuffer(Identifier identifier,
                     absl::Span<const uint8_t> buffer_data,
                     size_t root_type_size, VerifierFn verifier_fn);
-  Status FromString(Identifier identifier, std::string buffer_data,
-                    size_t root_type_size, VerifierFn verifier_fn);
   Status LoadFile(Identifier identifier, std::string path,
                   size_t root_type_size, VerifierFn verifier_fn);
 
@@ -132,10 +136,24 @@
   static StatusOr<std::unique_ptr<FlatBufferFile<T>>> WrapBuffer(
       Identifier identifier, absl::Span<const uint8_t> buffer_data);
 
+  // Loads the FlatBufferFile from a serialized byte-based STL container.
+  template <typename Container>
+  static StatusOr<std::unique_ptr<FlatBufferFile<T>>> FromContainer(
+      Identifier identifier, Container buffer_data);
+
   // Loads a FlatBufferFile from a serialized string.
   // The FlatBufferFile takes ownership of the string.
   static StatusOr<std::unique_ptr<FlatBufferFile<T>>> FromString(
-      Identifier identifier, std::string buffer_data);
+      Identifier identifier, std::string buffer_data) {
+    return FromContainer(identifier, std::move(buffer_data));
+  }
+
+  // Loads a FlatBufferFile from a serialized byte vector.
+  // The FlatBufferFile takes ownership of the vector.
+  static StatusOr<std::unique_ptr<FlatBufferFile<T>>> FromVector(
+      Identifier identifier, std::vector<uint8_t> buffer_data) {
+    return FromContainer(identifier, std::move(buffer_data));
+  }
 
   // Loads a FlatBufferFile from a serialized file on the file system.
   // This will attempt to mmap the file and is the preferred way of loading as
@@ -163,6 +181,29 @@
   }
 };
 
+template <typename Container>
+Status FlatBufferFileBase::FromContainer(Identifier identifier,
+                                         Container container,
+                                         size_t root_type_size,
+                                         VerifierFn verifier_fn) {
+  static_assert(sizeof(*container.data()) == 1,
+                "Expected container of byte sized elements");
+  auto buffer_data = absl::MakeConstSpan(
+      // Double static_cast through void is safer than reinterpret_cast.
+      static_cast<const uint8_t*>(static_cast<const void*>(container.data())),
+      container.size());
+  // Use a baton to keep the container alive until the FlatBufferFileBase is
+  // destroyed.
+  auto buffer_data_baton = IreeMoveToLambda(container);
+  return FromBuffer(
+      identifier, buffer_data,
+      [buffer_data_baton]() {
+        // Keeping the container alive.
+        (void)buffer_data_baton.value;
+      },
+      root_type_size, verifier_fn);
+}
+
 // static
 template <typename T>
 StatusOr<std::unique_ptr<FlatBufferFile<T>>> FlatBufferFile<T>::Create(
@@ -231,12 +272,13 @@
 
 // static
 template <typename T>
-StatusOr<std::unique_ptr<FlatBufferFile<T>>> FlatBufferFile<T>::FromString(
-    Identifier identifier, std::string buffer_data) {
+template <typename Container>
+StatusOr<std::unique_ptr<FlatBufferFile<T>>> FlatBufferFile<T>::FromContainer(
+    Identifier identifier, Container buffer_data) {
   std::unique_ptr<FlatBufferFile<T>> flat_buffer_file{new FlatBufferFile<T>};
   auto* base_file = static_cast<FlatBufferFileBase*>(flat_buffer_file.get());
-  RETURN_IF_ERROR(base_file->FromString(identifier, std::move(buffer_data),
-                                        sizeof(T), VerifierFnT));
+  RETURN_IF_ERROR(base_file->FromContainer(identifier, std::move(buffer_data),
+                                           sizeof(T), VerifierFnT));
   return std::move(flat_buffer_file);
 }
 

iree/base/memory.h

diff --git a/iree/base/memory.h b/iree/base/memory.h
index 5b1d01d..cc223cf 100644
--- a/iree/base/memory.h
+++ b/iree/base/memory.h
@@ -75,7 +75,7 @@
 // Usage:
 //   auto baton = MoveToLambda(my_ref);
 //   DoSomething([baton] () { baton.value; });
-#define MoveToLambda(p) ::iree::move_on_copy<decltype(p)>(std::move(p))
+#define IreeMoveToLambda(p) ::iree::move_on_copy<decltype(p)>(std::move(p))
 
 // TODO(benvanik): replace with an absl version when it exists.
 // A move-only RAII object that calls a stored cleanup functor when

iree/vm/debug/debug_server_flags.cc

diff --git a/iree/vm/debug/debug_server_flags.cc b/iree/vm/debug/debug_server_flags.cc
index 3997438..9720260 100644
--- a/iree/vm/debug/debug_server_flags.cc
+++ b/iree/vm/debug/debug_server_flags.cc
@@ -55,7 +55,7 @@
     RETURN_IF_ERROR(debug_server->WaitUntilSessionReady());
     LOG(INFO) << "Debugger attached";
     // TODO(benvanik): C++14 to avoid this.
-    auto debugger_baton = MoveToLambda(debugger);
+    auto debugger_baton = IreeMoveToLambda(debugger);
     debug_server->AtExit([debugger_baton]() { debugger_baton.value.reset(); });
   }
 #else