Refactor mjVFS and resource management.

Resource operations (e.g. mju_openResource, mju_readResource, and
mju_closeResource, etc.) are now all handled by a VFS instance. It
is now up to the VFS to determine which provider to use in order to
handle those operations.

This allows us to dynamically add/remove (aka "mount") providers to
a VFS to handle special requests. mj_addFileVFS and mj_addBufferVFS
have been reimplemented as two such use-cases. Moreover, we expose
the mounting behaviour with two new functions: mj_mountVFS and
mj_unmountVFS.

PiperOrigin-RevId: 861550939
Change-Id: I070eb4bcc2982466c8f368f7918005538baa5185

src/user/user_vfs.h

@@ -17,38 +17,115 @@
 #ifndef MUJOCO_SRC_USER_USER_VFS_H_
 #define MUJOCO_SRC_USER_USER_VFS_H_
 
-#include <stddef.h>
+#include <functional>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <string_view>
+#include <unordered_map>
 
-#include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjplugin.h>
+#include <mujoco/mujoco.h>
+#include "user/user_util.h"
 
-#ifdef __cplusplus
-extern "C" {
-#endif
+namespace mujoco::user {
 
-// Initialize an empty VFS, mj_deleteVFS must be called to deallocate the VFS
-MJAPI void mj_defaultVFS(mjVFS* vfs);
+// Underlying Virtual File System implementation for opaque mjVFS struct.
+//
+// This class owns and manages all the mjResource instances that are created
+// using the mju_openResource. Its main job is to find the correct
+// mjpResourceProvider to handle the mju open/read/close operations. It does
+// this by "mounting" mjpResourceProviders at specific paths such that any
+// operation within that path will be handled by the mjpResourceProvider.
+//
+// Mounting can be done explicitly (using mj_mountVFS) or implicitly (using
+// mjp_registerResourceProvider). If no provider is found for a given path,
+// then a "default" provider is used that uses normal C file operations to
+// open/read/close files.
+//
+// To support legacy use-cases (where the VFS is an optional argument), a
+// "self-destruct" mode can be configured so that a temporary VFS instance can
+// be created with a lifetime tied to the opened mjResource instance.
+//
+// This class itself is thread-safe, but it makes no guarantees about the
+// thread-safety of the underlying mjResourceProviders.
+class VFS {
+ public:
+  explicit VFS(mjVFS* vfs);
+  ~VFS();
 
-// add file to VFS, return 0: success, 2: repeated name, -1: not found on disk
-MJAPI int mj_addFileVFS(mjVFS* vfs, const char* directory, const char* filename);
+  VFS(const VFS&) = delete;
+  VFS& operator=(const VFS&) = delete;
 
-// add file from buffer into VFS, return 0: success, 2: repeated name, -1: failed to load
-MJAPI int mj_addBufferVFS(mjVFS* vfs, const char* filename, const void* buffer, int nbuffer);
+  // Status codes for VFS operations. Values are based on the mj_VFS APIs.
+  enum Status {
+    kSuccess = 0,
+    kFailedToLoad = -1,
+    kFailedToRead = -1,
+    kNotFound = -1,
+    kRepeatedName = 2,
+    kInvalidVfs = -1,
+    kInvalidResource = -1,
+    kInvalidResourceProvider = -1,
+  };
 
-// return file index in VFS, or -1 if not found in VFS
-MJAPI int mj_findFileVFS(const mjVFS* vfs, const char* filename);
+  // 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);
 
-// delete file from VFS, return 0: success, -1: not found in VFS
-MJAPI int mj_deleteFileVFS(mjVFS* vfs, const char* filename);
+  // 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);
 
-// delete all files from VFS
-MJAPI void mj_deleteVFS(mjVFS* vfs);
+  // Closes the resource by invoking the 'close' callback for the
+  // mjpResourceProvider associated with the resource.
+  Status Close(mjResource* resource);
 
-#ifdef __cplusplus
-}
-#endif
+  // Mounts a ResourceProvider at the given path. All subsequent operations
+  // under `path` will be delegated to the `provider` until it is unmounted.
+  Status Mount(const FilePath& path, const mjpResourceProvider* provider);
 
-const mjpResourceProvider* GetVfsResourceProvider();
+  // Unmounts the ResourceProvider from the given path.
+  Status Unmount(const FilePath& path);
+
+  // Sets a destructor to be called when the VFS has no more open resources.
+  // Assumes that `destructor` will delete `this`.
+  //
+  // 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);
+
+  // Converts the public C-API pointer to the internal C++ class.
+  static VFS* Upcast(mjVFS* vfs);
+  static const VFS* Upcast(const mjVFS* vfs);
+
+ private:
+  using ResourcePtr = std::unique_ptr<mjResource, void (*)(mjResource*)>;
+  ResourcePtr CreateResource(std::string_view name,
+                             const mjpResourceProvider* provider);
+
+  // Returns a mounted mjResource* that matches the given path. If no explicitly
+  // mounted mjResource* is found, returns a "default" mounting that uses the
+  // C file system.
+  mjResource* FindMount(const std::string& fullpath);
+
+  // Invokes the `destructor_, but only if it has been set previously. This
+  // should be only called when resources_ is empty and callers should assume
+  // that `this` will be invalidated after this call.
+  void MaybeSelfDestruct();
+
+  mjVFS* self_;
+  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_;
+};
+
+}  // namespace mujoco::user
 
 #endif  // MUJOCO_SRC_USER_USER_VFS_H_