Merge pull request #3105 from NandiniDhanrale:fix/vfs-move-3098

PiperOrigin-RevId: 881946056
Change-Id: Ide041ff1dcdae00133251d928936920b47cf6c0d

src/user/user_resource.cc

@@ -44,9 +44,9 @@ mjResource* mju_openResource(const char* dir, const char* name,
   if (non_const_vfs == nullptr) {
     mjVFS* local_vfs = (mjVFS*)mju_malloc(sizeof(mjVFS));
     mj_defaultVFS(local_vfs);
-    mujoco::user::VFS::Upcast(local_vfs)->SetToSelfDestruct([](mjVFS* ptr) {
-      mj_deleteVFS(ptr);
-      mju_free(ptr);
+    mujoco::user::VFS::Upcast(local_vfs)->SetToSelfDestruct([=]() {
+      mj_deleteVFS(local_vfs);
+      mju_free(local_vfs);
     });
 
     non_const_vfs = local_vfs;

src/user/user_vfs.cc

@@ -105,7 +105,8 @@ std::string StripPathAndLower(std::string path) {
 
 namespace mujoco::user {
 
-VFS::VFS(mjVFS* vfs) : self_(vfs) {
+VFS::VFS(mjVFS* vfs) {
+  wrapped_vfs_.impl_ = this;
   mjp_defaultResourceProvider(&default_provider_);
   default_provider_.open = [](mjResource* res) {
     return OpenFile(res->name, res);
@@ -121,7 +122,7 @@ VFS::VFS(mjVFS* vfs) : self_(vfs) {
   };
 
   default_provider_.prefix = nullptr;
-  default_mount_.vfs = self_;
+  default_mount_.vfs = &wrapped_vfs_;
   default_mount_.provider = &default_provider_;
   default_mount_.data = nullptr;
   default_mount_.name = nullptr;
@@ -245,7 +246,7 @@ int VFS::Read(mjResource* resource, const void** buffer) {
 VFS::ResourcePtr VFS::CreateResource(std::string_view name,
                                      const mjpResourceProvider* provider) {
   mjResource* res = new mjResource();
-  res->vfs = self_;
+  res->vfs = &wrapped_vfs_;
   res->provider = provider;
   res->data = nullptr;
   res->name = new char[name.size() + 1];
@@ -312,11 +313,14 @@ mjResource* VFS::FindMount(const std::string& fullpath) {
 
 void VFS::MaybeSelfDestruct() {
   if (destructor_) {
-    destructor_(self_);
+    // Copy the destructor to a local variable so that we can destroy `this`
+    // object within the destructor.
+    auto fn = std::move(destructor_);
+    fn();
   }
 }
 
-void VFS::SetToSelfDestruct(std::function<void(mjVFS*)> destructor) {
+void VFS::SetToSelfDestruct(std::function<void()> destructor) {
   destructor_ = std::move(destructor);
 }
 

src/user/user_vfs.h

@@ -96,7 +96,7 @@ class VFS {
   // This is useful for when you want to create a temporary VFS instance with
   // a lifetime tied to a single mjResource to be opened. The `destructor`
   // should be set to `delete this` and any other cleanup that needs to happen.
-  void SetToSelfDestruct(std::function<void(mjVFS*)> destructor);
+  void SetToSelfDestruct(std::function<void()> destructor);
 
   // Converts the public C-API pointer to the internal C++ class.
   static VFS* Upcast(mjVFS* vfs);
@@ -117,13 +117,13 @@ class VFS {
   // that `this` will be invalidated after this call.
   void MaybeSelfDestruct();
 
-  mjVFS* self_;
+  mjVFS wrapped_vfs_;
   std::mutex mutex_;  // Protects open_resources_ and mounts_.
   std::unordered_map<mjResource*, ResourcePtr> open_resources_;
   std::unordered_map<std::string, ResourcePtr> mounts_;
   mjResource default_mount_;
   mjpResourceProvider default_provider_;
-  std::function<void(mjVFS*)> destructor_;
+  std::function<void()> destructor_;
 };
 
 }  // namespace mujoco::user

test/user/user_vfs_test.cc

@@ -407,5 +407,29 @@ TEST_F(UserVfsTest, StackedMounts) {
   EXPECT_EQ(test2, expect2);
   EXPECT_EQ(test3, expect3);
 }
+
+TEST_F(UserVfsTest, MoveVfs) {
+  // Create and move a VFS to another address.
+  mjVFS* original = new mjVFS();
+  mj_defaultVFS(original);
+  mjVFS vfs = *original;
+  delete original;
+
+  std::string buffer = "<mujoco/>";
+  mj_addBufferVFS(&vfs, "model", static_cast<const void*>(buffer.c_str()),
+                  buffer.size());
+
+  mjResource* resource = mju_openResource("", "model", &vfs, nullptr, 0);
+  ASSERT_THAT(resource, NotNull());
+
+  const void* out = nullptr;
+  const int size = mju_readResource(resource, &out);
+  EXPECT_GT(size, 0);
+  EXPECT_THAT(out, NotNull());
+
+  mju_closeResource(resource);
+  mj_deleteVFS(&vfs);
+}
+
 }  // namespace
 }  // namespace mujoco