Add pluggable resource writing to MuJoCo

Extend mjpResourceProvider with an optional write callback (write)
so that mj_encode, mj_saveXML, and mj_saveModel can write to any
registered provider.

PiperOrigin-RevId: 945202741
Change-Id: I37903425260932e555f4a8c2392c4ff8c2e6cc06

src/user/user_api.cc

@@ -261,29 +261,18 @@ mjtSize mj_encode(const mjSpec* s, const mjModel* m, const char* filename,
     return -1;
   }
 
-  FILE* fp = fopen(filename, "wb");
-  if (!fp) {
-    std::free(resource.data);
-    if (error) {
-      strncpy(error, "could not open file for writing", error_sz);
-      error[error_sz - 1] = '\0';
-    }
-    return -1;
-  }
-
-  const std::size_t written = fwrite(resource.data, 1, nbytes, fp);
-  fclose(fp);
+  mjtSize written = mju_writeResource(filename, resource.data, nbytes, vfs, error, error_sz);
   encoder->close_resource(&resource);
 
-  if (static_cast<mjtSize>(written) != nbytes) {
-    if (error) {
-      strncpy(error, "failed to write all bytes to file", error_sz);
+  if (written != nbytes) {
+    if (error && error[0] == '\0') {
+      strncpy(error, "write failed", error_sz);
       error[error_sz - 1] = '\0';
     }
     return -1;
   }
 
-  return nbytes;
+  return written;
 }
 
 // helper function to log compile time diagnostics

src/user/user_resource.cc

@@ -30,8 +30,13 @@
 #include "user/user_util.h"
 #include "user/user_vfs.h"
 
-mjResource* mju_openResource(const char* dir, const char* name,
-                             const mjVFS* vfs, char* error, size_t nerror) {
+static mjResource* openResourceInternal(const char* dir, const char* name,
+                                        const mjVFS* vfs, char* error,
+                                        size_t nerror) {
+  if (error && nerror > 0) {
+    error[0] = '\0';
+  }
+
   // TODO: Update API to use non-const pointer. Unfortunately, while this is
   // ABI stable, it will cause compiler errors in user code that is const
   // correct.
@@ -52,20 +57,21 @@ mjResource* mju_openResource(const char* dir, const char* name,
     non_const_vfs = local_vfs;
   }
 
-  mjResource* resource =
-      mujoco::user::VFS::Upcast(non_const_vfs)->Open(dir ? dir : "", name);
+  mjResource* resource = mujoco::user::VFS::Upcast(non_const_vfs)
+                             ->Open(dir ? dir : "", name, error, nerror);
 
-  if (error) {
-    if (resource) {
-      error[0] = '\0';
-    } else {
-      std::snprintf(error, nerror, "Error opening file '%s'", name);
-    }
+  if (error && nerror > 0 && !resource && error[0] == '\0') {
+    std::snprintf(error, nerror, "Error opening file '%s'", name);
   }
 
   return resource;
 }
 
+mjResource* mju_openResource(const char* dir, const char* name,
+                             const mjVFS* vfs, char* error, size_t nerror) {
+  return openResourceInternal(dir, name, vfs, error, nerror);
+}
+
 void mju_closeResource(mjResource* resource) {
   if (resource && resource->vfs) {
     mujoco::user::VFS::Upcast(resource->vfs)->Close(resource);
@@ -79,6 +85,46 @@ int mju_readResource(mjResource* resource, const void** buffer) {
   return -1;  // default (error reading bytes)
 }
 
+mjtSize mju_writeResource(const char* name, const void* buffer, mjtSize nbytes,
+                          const mjVFS* vfs, char* error, size_t nerror) {
+  if (error && nerror > 0) {
+    error[0] = '\0';
+  }
+
+  if (!name) {
+    if (error && nerror > 0) {
+      std::snprintf(error, nerror, "Resource name is NULL");
+    }
+    return -1;
+  }
+
+  mjResource resource = {};
+  resource.name = const_cast<char*>(name);
+
+  mjVFS* non_const_vfs = const_cast<mjVFS*>(vfs);
+  bool local_vfs_created = false;
+  if (non_const_vfs == nullptr) {
+    non_const_vfs = (mjVFS*)mju_malloc(sizeof(mjVFS));
+    mj_defaultVFS(non_const_vfs);
+    local_vfs_created = true;
+  }
+
+  resource.vfs = non_const_vfs;
+
+  mjtSize written = mujoco::user::VFS::Upcast(non_const_vfs)->Write(&resource, buffer, nbytes);
+
+  if (local_vfs_created) {
+    mj_deleteVFS(non_const_vfs);
+    mju_free(non_const_vfs);
+  }
+
+  if (written < 0 && error && nerror > 0 && error[0] == '\0') {
+    std::snprintf(error, nerror, "Error writing resource '%s'", name);
+  }
+
+  return written;
+}
+
 void mju_getResourceDir(mjResource* resource, const char** dir, int* ndir) {
   *dir = nullptr;
   *ndir = 0;

src/user/user_resource.h

@@ -38,6 +38,9 @@ MJAPI void mju_closeResource(mjResource* resource);
 // return negative value if error
 MJAPI int mju_readResource(mjResource* resource, const void** buffer);
 
+MJAPI mjtSize mju_writeResource(const char* name, const void* buffer, mjtSize nbytes,
+                                const mjVFS* vfs, char* error, size_t nerror);
+
 // set for a resource with a name partitioned as {dir}{filename}, the dir and ndir pointers
 MJAPI void mju_getResourceDir(mjResource* resource, const char** dir, int* ndir);
 

src/user/user_vfs.cc

@@ -23,6 +23,7 @@
 #include <cctype>
 #include <cstddef>
 #include <cstdint>
+#include <cstdio>
 #include <cstring>
 #include <ctime>
 #include <functional>
@@ -120,6 +121,15 @@ VFS::VFS(mjVFS* vfs) {
   default_provider_.modified = [](const mjResource* res, const char* time) {
     return FileModified(res, time);
   };
+  default_provider_.write = [](mjResource* res, const void* buffer, mjtSize nbytes) -> mjtSize {
+    FILE* fp = fopen(res->name, "wb");
+    if (!fp) {
+      return -1;
+    }
+    mjtSize written = static_cast<mjtSize>(fwrite(buffer, 1, nbytes, fp));
+    fclose(fp);
+    return written == nbytes ? written : -1;
+  };
 
   default_provider_.prefix = nullptr;
   default_mount_.vfs = &wrapped_vfs_;
@@ -149,22 +159,28 @@ VFS::~VFS() {
   mounts_.clear();
 }
 
-mjResource* VFS::Open(const char* dir, const char* name) {
+mjResource* VFS::Open(const char* dir, const char* name, char* error,
+                      size_t nerror) {
   const std::string path = FilePath(dir, name).Str();
 
   const mjResource* mount = FindMount(path);
-  if (!mount) {
+  if (!mount || !mount->provider) {
+    if (error) {
+      std::snprintf(error, nerror, "No provider found for '%s'", name);
+    }
     MaybeSelfDestruct();
     return nullptr;
   }
 
-  ResourcePtr res = CreateResource(path.c_str(), mount->provider);
+  const mjpResourceProvider* provider = mount->provider;
+
+  ResourcePtr res = CreateResource(path.c_str(), provider);
 
   // Smuggle the mounted resource provider's mount-specific data pointer in the
   // requested resource's data pointer. This allows the provider to access its
   // own per-mount data without any intrusive changes to the provider interface.
   res->data = mount->data;
-  const int result = mount->provider->open(res.get());
+  const int result = provider->open ? provider->open(res.get()) : 0;
   // If the data pointer was not modified, then that means the resource did not
   // set its own data pointer. So, we need to set it back to nullptr.
   if (res->data == mount->data) {
@@ -261,6 +277,25 @@ int VFS::Read(mjResource* resource, const void** buffer) {
   return kFailedToRead;
 }
 
+mjtSize VFS::Write(mjResource* resource, const void* buffer, mjtSize nbytes) {
+  if (resource) {
+    if (!resource->provider) {
+      const mjResource* mount = FindMount(resource->name);
+      if (mount) {
+        resource->provider = mount->provider;
+        // Smuggle mount data if not already set.
+        if (!resource->data) {
+          resource->data = mount->data;
+        }
+      }
+    }
+    if (resource->provider && resource->provider->write) {
+      return resource->provider->write(resource, buffer, nbytes);
+    }
+  }
+  return -1;
+}
+
 VFS::ResourcePtr VFS::CreateResource(std::string_view name,
                                      const mjpResourceProvider* provider) {
   mjResource* res = new mjResource();

src/user/user_vfs.h

@@ -72,15 +72,21 @@ class VFS {
   // Opens a mjResource for the given path, or nullptr on error. If successful,
   // will invoke the 'open' callback for the mjpResourceProvider associated with
   // the path/dir.
-  mjResource* Open(const char* dir, const char* name);
+  mjResource* Open(const char* dir, const char* name, char* error = nullptr,
+                   size_t nerror = 0);
 
   // Sets `buffer` to the contents of the resource and returns the number of
   // bytes of the content. This is done by invoking the 'read' callback for the
   // mjpResourceProvider associated with the resource. Returns -1 on error.
   int Read(mjResource* resource, const void** buffer);
 
-  // Closes the resource by invoking the 'close' callback for the
-  // mjpResourceProvider associated with the resource.
+  // Writes the resource data. This is done by invoking the 'write' callback for
+  // the mjpResourceProvider associated with the resource.
+  // Returns bytes written, or -1 on error.
+  mjtSize Write(mjResource* resource, const void* buffer, mjtSize nbytes);
+
+  // Closes the resource by invoking the 'close' callback for
+  // the mjpResourceProvider associated with the resource.
   Status Close(mjResource* resource);
 
   // Mounts a ResourceProvider at the given path. All subsequent operations