Add AddVector and GetVector to user cache.

PiperOrigin-RevId: 601485067
Change-Id: I1a31f15c2495de4f057316022dcddf4df39afd64

src/user/user_cache.cc

@@ -26,35 +26,54 @@
 #include <utility>
 #include <vector>
 
-// adds a block of data for the asset and returns number of bytes stored
+// copies a block of data into the asset and returns number of bytes stored
-template<typename T>
+// loading data into an asset should happen in a single thread
-std::size_t mjCAsset::Add(const std::string& name, const std::vector<T>& v) {
+template<typename T> std::size_t mjCAsset::Add(const std::string& name,
+                                               const T* data, std::size_t n) {
   auto [it, inserted] = blocks_.insert({name, mjCAssetData()});
   if (!inserted) {
     return 0;
   }
 
-  std::size_t n = v.size() * sizeof(T);
+  std::size_t nbytes = n * sizeof(T);
-  const uint8_t* ptr = reinterpret_cast<const uint8_t*>(v.data());
+  const uint8_t* ptr = reinterpret_cast<const uint8_t*>(data);
   mjCAssetData& block = it->second;
 
-  block.bytes = std::make_shared<uint8_t[]>(n);
+  block.bytes = std::make_shared<uint8_t[]>(nbytes);
-  std::copy(ptr, ptr + n, block.bytes.get());
+  std::copy(ptr, ptr + nbytes, block.bytes.get());
-  block.nbytes = n;
+  block.nbytes = nbytes;
-  nbytes_ += n;
+  nbytes_ += nbytes;
-  return n;
+  return nbytes;
 }
-
 template std::size_t mjCAsset::Add(const std::string& name,
-                                   const std::vector<int>& v);
+                                   const uint8_t* data, std::size_t n);
 template std::size_t mjCAsset::Add(const std::string& name,
-                                   const std::vector<float>& v);
+                                   const int* data, std::size_t n);
 template std::size_t mjCAsset::Add(const std::string& name,
-                                   const std::vector<double>& v);
+                                   const float* data, std::size_t n);
+template std::size_t mjCAsset::Add(const std::string& name,
+                                   const double* data, std::size_t n);
 
 
 
-// fetches a block of data, sets n to size of data
+// copies a vector into the asset and returns number of bytes stored
+// loading data into an asset should happen in a single thread
+template<typename T> std::size_t mjCAsset::AddVector(const std::string& name,
+                                                     const std::vector<T>& v) {
+  return Add<T>(name, v.data(), v.size());
+}
+template std::size_t mjCAsset::AddVector(const std::string& name,
+                                         const std::vector<uint8_t>& v);
+template std::size_t mjCAsset::AddVector(const std::string& name,
+                                         const std::vector<int>& v);
+template std::size_t mjCAsset::AddVector(const std::string& name,
+                                         const std::vector<float>& v);
+template std::size_t mjCAsset::AddVector(const std::string& name,
+                                         const std::vector<double>& v);
+
+
+
+// returns a pointer to a block of data, sets n to size of data
 template<typename T>
 const T* mjCAsset::Get(const std::string& name, std::size_t* n) const {
   auto it = blocks_.find(name);
@@ -76,6 +95,8 @@ const T* mjCAsset::Get(const std::string& name, std::size_t* n) const {
   return reinterpret_cast<T*>(data.bytes.get());
 }
 
+template
+const uint8_t* mjCAsset::Get(const std::string& name, std::size_t* n) const;
 template
 const int* mjCAsset::Get(const std::string& name, std::size_t* n) const;
 template
@@ -85,6 +106,28 @@ const double* mjCAsset::Get(const std::string& name, std::size_t* n) const;
 
 
 
+// copies a block of data into a vector
+template<typename T> std::optional<std::vector<T>>
+mjCAsset::GetVector(const std::string& name) const {
+  std::size_t n;
+  const T* ptr = Get<T>(name, &n);
+  if (ptr == nullptr) {
+    return std::nullopt;
+  }
+  return std::vector<T>(ptr, ptr + n);
+}
+
+template std::optional<std::vector<uint8_t>>
+mjCAsset::GetVector(const std::string& name) const;
+template std::optional<std::vector<int>>
+mjCAsset::GetVector(const std::string& name) const;
+template std::optional<std::vector<float>>
+mjCAsset::GetVector(const std::string& name) const;
+template std::optional<std::vector<double>>
+mjCAsset::GetVector(const std::string& name) const;
+
+
+
 // replaces blocks data in asset
 void mjCAsset::ReplaceBlocks(
     const std::unordered_map<std::string, mjCAssetData>& blocks,

src/user/user_cache.h

@@ -56,17 +56,30 @@ class mjCAsset {
   mjCAsset(const mjCAsset& other) = default;
   mjCAsset& operator=(const mjCAsset& other) = default;
 
-  // adds a block of data for the asset and returns number of bytes stored
+  // copies a block of data into the asset and returns number of bytes stored
-  // loading data into an asset should happen single thread
+  // loading data into an asset should happen in a single thread
   template<typename T> std::size_t Add(const std::string& name,
-                                       const std::vector<T>& v);
+                                       const T* data, std::size_t n);
 
-  // fetches a block of data, sets n to size of data
+
-  // TODO(kylebayes): The C++ span utility doesn't seem to be supported by
+  // copies a vector into the asset and returns number of bytes stored
-  // Google C++ coding standards.  For now, we fallback to an C style API.
+  // loading data into an asset should happen in a single thread
+  template<typename T> std::size_t AddVector(const std::string& name,
+                                             const std::vector<T>& v);
+
+  // returns a pointer to a block of data, sets n to size of data
   template<typename T>
   const T* Get(const std::string& name, std::size_t* n) const;
 
+  // copies a block of data into a vector
+  template<typename T>
+  std::optional<std::vector<T>> GetVector(const std::string& name) const;
+
+  // returns true if a block of data by the given name is stored in the asset
+  bool HasData(const std::string& name) const {
+    return blocks_.find(name) != blocks_.end();
+  }
+
  private:
   mjCAsset() = default;
 

test/user/user_cache_test.cc

@@ -66,8 +66,8 @@ TEST(CacheTest, AddSuccessTest) {
   std::vector<double> v2 = {1.0, 2.0, 3.0};
 
   mjCAsset asset("file.xml", "foo.obj", "now");
-  std::size_t nbytes1 = asset.Add("v1", v1);
+  std::size_t nbytes1 = asset.AddVector("v1", v1);
-  std::size_t nbytes2 = asset.Add("v2", v2);
+  std::size_t nbytes2 = asset.AddVector("v2", v2);
   cache.Insert(asset);
 
   ASSERT_EQ(nbytes1, 12);
@@ -81,8 +81,8 @@ TEST(CacheTest, AddFailureTest) {
   std::vector<double> v2 = {1.0, 2.0, 3.0};
 
   mjCAsset asset("file.xml", "foo.obj", "now");
-  asset.Add("v1", v1);
+  asset.AddVector("v1", v1);
-  std::size_t nbytes = asset.Add("v1", v2);
+  std::size_t nbytes = asset.AddVector("v1", v2);
   cache.Insert(asset);
 
   ASSERT_EQ(nbytes, 0);
@@ -95,18 +95,16 @@ TEST(CacheTest, InsertReplaceTest) {
   std::vector<double> v2 = {1.0, 2.0, 3.0};
 
   mjCAsset asset("file.xml", "foo.obj", "now");
-  asset.Add("v", v1);
+  asset.AddVector("v", v1);
   cache.Insert(asset);
 
   mjCAsset asset2("file.xml", "foo.obj", "nower");
-  asset2.Add("v", v2);
+  asset2.AddVector("v", v2);
   bool inserted = cache.Insert(asset2);
   EXPECT_TRUE(inserted);
 
   mjCAsset asset3 = *(cache.Get("foo.obj"));
-  std::size_t n = 0;
+  std::vector<double> v3 = asset3.GetVector<double>("v").value();
-  const double* ptr = asset3.Get<double>("v", &n);
-  std::vector<double> v3 = std::vector<double>(ptr, ptr + n);
   EXPECT_THAT(v3, ElementsAreArray(v2));
 
   ASSERT_EQ(cache.Size(), 24);
@@ -118,8 +116,8 @@ TEST(CacheTest, MoveInsertNewTest) {
   std::vector<double> v2 = {1.0, 2.0, 3.0};
 
   mjCAsset asset("file.xml", "foo.obj", "now");
-  std::size_t nbytes1 = asset.Add("v1", v1);
+  std::size_t nbytes1 = asset.AddVector("v1", v1);
-  std::size_t nbytes2 = asset.Add("v2", v2);
+  std::size_t nbytes2 = asset.AddVector("v2", v2);
   cache.Insert(std::move(asset));
 
   ASSERT_EQ(nbytes1, 12);
@@ -133,18 +131,16 @@ TEST(CacheTest, MoveInsertReplaceTest) {
   std::vector<double> v2 = {1.0, 2.0, 3.0};
 
   mjCAsset asset("file.xml", "foo.obj", "now");
-  asset.Add("v", v1);
+  asset.AddVector("v", v1);
   cache.Insert(std::move(asset));
 
   mjCAsset asset2("file.xml", "foo.obj", "nower");
-  asset2.Add("v", v2);
+  asset2.AddVector("v", v2);
   bool inserted = cache.Insert(std::move(asset2));
   EXPECT_TRUE(inserted);
 
   mjCAsset asset3 = *(cache.Get("foo.obj"));
-  std::size_t n = 0;
+  std::vector<double> v3 = asset3.GetVector<double>("v").value();
-  const double* ptr = asset3.Get<double>("v", &n);
-  std::vector<double> v3 = std::vector<double>(ptr, ptr + n);
   EXPECT_THAT(v3, ElementsAreArray(v2));
 
   ASSERT_EQ(cache.Size(), 24);
@@ -153,40 +149,33 @@ TEST(CacheTest, MoveInsertReplaceTest) {
 TEST(CacheTest, GetSuccessTest) {
   mjCCache cache(kMaxSize);
   std::vector<int> v1 = {1, 2, 3};
-  std::vector<double> v2 = {1.0, 2.0, 3.0};
+  std::vector<double> v2 = {4.0, 5.0, 6.0};
   mjCAsset asset("file.xml", "foo.obj", "now");
-  asset.Add("v1", v1);
+  asset.AddVector("v1", v1);
-  asset.Add("v2", v2);
+  asset.AddVector("v2", v2);
   cache.Insert(asset);
 
-  std::size_t n = 0;
-
   mjCAsset asset2 = *(cache.Get("foo.obj"));
 
-  const int* ptr1 = asset2.Get<int>("v1", &n);
+  std::vector<int> v3 = asset2.GetVector<int>("v1").value();
-  EXPECT_EQ(n, 3);
-  std::vector<int> v3 = std::vector<int>(ptr1, ptr1 + n);
   EXPECT_THAT(v3, ElementsAreArray(v1));
 
-  const double* ptr2 = asset2.Get<double>("v2", &n);
+  std::vector<double> v4 = asset2.GetVector<double>("v2").value();
-  EXPECT_EQ(n, 3);
+  EXPECT_THAT(v4, ElementsAreArray(v2));
-  std::vector<double> v4 = std::vector<double>(ptr2, ptr2 + n);
-  EXPECT_THAT(v4, ElementsAreArray(v3));
 }
 
 TEST(CacheTest, GetFailueTest) {
   mjCCache cache(kMaxSize);
   std::vector<int> v = {1, 2, 3};
   mjCAsset asset("file.xml", "foo.obj", "now");
-  asset.Add("v", v);
+  asset.AddVector("v", v);
   cache.Insert(asset);
 
   mjCAsset asset2 = *(cache.Get("foo.obj"));
   EXPECT_EQ(cache.Get("bar.obj").has_value(), false);
 
-  std::size_t n = 0;
+  auto v2 = asset2.GetVector<int>("v2");
-  const int* ptr = asset.Get<int>("v2", &n);
+  EXPECT_EQ(v2.has_value(), false);
-  EXPECT_THAT(ptr, IsNull());
 }
 
 // Trim cache based off of access count
@@ -197,8 +186,8 @@ TEST(CacheTest, LimitTest1) {
 
   mjCAsset asset1 = mjCAsset("file.xml", "foo.obj", "now");
   mjCAsset asset2 = mjCAsset("file.xml", "bar.obj", "now");
-  asset1.Add("foo.obj", v);
+  asset1.AddVector("foo.obj", v);
-  asset2.Add("bar.obj", v);
+  asset2.AddVector("bar.obj", v);
   cache.Insert(asset1);
   cache.Insert(asset2);
 
@@ -223,8 +212,8 @@ TEST(CacheTest, LimitTest2) {
   std::vector<int> v = {1, 2, 3};
   mjCAsset asset1("file.xml", "foo.obj", "now");
   mjCAsset asset2("file.xml", "bar.obj", "now");
-  asset1.Add("v", v);
+  asset1.AddVector("v", v);
-  asset2.Add("v", v);
+  asset2.AddVector("v", v);
 
   cache.Insert(asset1);
   cache.Insert(asset2);