Add support for vfs assets in mj_compile function to MuJoCo WASM.

Regarding mj_compile VFS argument:
I was encountering a memory error when using `std::optional<MjVFS>` with `mj_compile` (because vfs should be an optional argument of the function). This is likely because `MjVFS` doesn't have copy/move constructors, which are required for types used with `std::optional`, causing issues with memory management when std::optional copies the MjVFS object.
To solve this and to make vfs an optional argument to `mj_compile`, I replaced `mj_compile_wrapper` with two overloaded functions. This avoids using `std::optional<MjVFS>` and resolves the memory issue, while still allowing mj_compile to be called with or without the vfs argument from JavaScript.

But what if we instead add those copy/move constructors so it works with `std::optional`?:
While making `MjVFS` copyable or movable would theoretically allow it to work with std::optional, it's not a feasible approach in this case.
The MjVFS struct manages file data in memory, allocating buffers when files are added and freeing them when `mj_deleteVFS` is called in its destructor. A copy would require a deep copy of the entire virtual filesystem, but MuJoCo does not provide a public API function (like a  "mj_copyVFS") to do this. Implementing a deep copy manually would be complex and fragile. A shallow copy would result in multiple MjVFS objects pointing to the same memory, leading to double-free errors when their destructors run.
Given these constraints, MjVFS cannot safely be made copyable. The approach of providing two overloaded C++ functions is better.

PiperOrigin-RevId: 855740541
Change-Id: Ia08b6db8868ad245f07145db3fb0322cb9f737cb

wasm/codegen/generated/bindings.cc

@@ -31,9 +31,9 @@
 #include <vector>
 
 #include <mujoco/mjmodel.h>
+#include <mujoco/mjspec.h>
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mujoco.h>
-#include <mujoco/mjspec.h>
 #include "engine/engine_util_errmem.h"
 #include "wasm/unpack.h"
 
@@ -7638,6 +7638,27 @@ struct MjvScene {
   std::vector<MjvGLCamera> camera;
 };
 
+struct MjVFS {
+  MjVFS() : ptr_(new mjVFS) { mj_defaultVFS(ptr_); }
+  ~MjVFS() {
+    mj_deleteVFS(ptr_);
+  }
+  void AddBuffer(const std::string& name, const emscripten::val& buffer) {
+    std::vector<uint8_t> vec = emscripten::vecFromJSArray<uint8_t>(buffer);
+    int result = mj_addBufferVFS(ptr_, name.c_str(), vec.data(), vec.size());
+    if (result != 0) {
+      mju_error("Could not add buffer to VFS: %d", result);
+    }
+  }
+  void DeleteFile(const std::string& filename) {
+    mj_deleteFileVFS(ptr_, filename.c_str());
+  }
+  mjVFS* get() const { return ptr_; }
+
+ private:
+  mjVFS* ptr_;
+};
+
 MjModel::MjModel(mjModel* ptr)
     : ptr_(ptr), opt(&ptr->opt), vis(&ptr->vis), stat(&ptr->stat) {}
 
@@ -7833,7 +7854,7 @@ std::unique_ptr<MjSpec> parseXMLString_wrapper(const std::string &xml) {
   return std::unique_ptr<MjSpec>(new MjSpec(ptr));
 }
 
-std::unique_ptr<MjModel> mj_compile_wrapper(const MjSpec& spec) {
+std::unique_ptr<MjModel> mj_compile_wrapper_1(const MjSpec& spec) {
   mjSpec* spec_ptr = spec.get();
   mjModel* model = mj_compile(spec_ptr, nullptr);
   if (!model || mjs_isWarning(spec_ptr)) {
@@ -7842,6 +7863,16 @@ std::unique_ptr<MjModel> mj_compile_wrapper(const MjSpec& spec) {
   return std::unique_ptr<MjModel>(new MjModel(model));
 }
 
+std::unique_ptr<MjModel> mj_compile_wrapper_2(const MjSpec& spec, const MjVFS& vfs) {
+  mjSpec* spec_ptr = spec.get();
+  mjVFS* vfs_ptr = vfs.get();
+  mjModel* model = mj_compile(spec_ptr, vfs_ptr);
+  if (!model || mjs_isWarning(spec_ptr)) {
+    mju_error("%s", mjs_getError(spec_ptr));
+  }
+  return std::unique_ptr<MjModel>(new MjModel(model));
+}
+
 void error_wrapper(const String& msg) { mju_error("%s\n", msg.as<const std::string>().data()); }
 
 int mj_saveLastXML_wrapper(const String& filename, const MjModel& m) {
@@ -12334,6 +12365,11 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
   emscripten::register_optional<MjsTuple>();
   emscripten::register_optional<MjsWrap>();
 
+  emscripten::class_<MjVFS>("MjVFS")
+      .constructor<>()
+      .function("addBuffer", &MjVFS::AddBuffer)
+      .function("deleteFile", &MjVFS::DeleteFile);
+
   function("mj_Euler", &mj_Euler_wrapper);
   function("mj_RungeKutta", &mj_RungeKutta_wrapper);
   function("mj_addContact", &mj_addContact_wrapper);
@@ -12348,7 +12384,6 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
   function("mj_comPos", &mj_comPos_wrapper);
   function("mj_comVel", &mj_comVel_wrapper);
   function("mj_compareFwdInv", &mj_compareFwdInv_wrapper);
-  function("mj_compile", &mj_compile_wrapper);
   function("mj_constraintUpdate", &mj_constraintUpdate_wrapper);
   function("mj_contactForce", &mj_contactForce_wrapper);
   function("mj_copyBack", &mj_copyBack_wrapper);
@@ -12427,13 +12462,11 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
   function("mj_resetDataKeyframe", &mj_resetDataKeyframe_wrapper);
   function("mj_rne", &mj_rne_wrapper);
   function("mj_rnePostConstraint", &mj_rnePostConstraint_wrapper);
-  function("mj_saveLastXML", &mj_saveLastXML_wrapper);
   function("mj_sensorAcc", &mj_sensorAcc_wrapper);
   function("mj_sensorPos", &mj_sensorPos_wrapper);
   function("mj_sensorVel", &mj_sensorVel_wrapper);
   function("mj_setConst", &mj_setConst_wrapper);
   function("mj_setKeyframe", &mj_setKeyframe_wrapper);
-  function("mj_setLengthRange", &mj_setLengthRange_wrapper);
   function("mj_setState", &mj_setState_wrapper);
   function("mj_setTotalmass", &mj_setTotalmass_wrapper);
   function("mj_sizeModel", &mj_sizeModel_wrapper);
@@ -12710,9 +12743,14 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
   function("mjv_updateCamera", &mjv_updateCamera_wrapper);
   function("mjv_updateScene", &mjv_updateScene_wrapper);
   function("mjv_updateSkin", &mjv_updateSkin_wrapper);
-
   function("parseXMLString", &parseXMLString_wrapper, take_ownership());
   function("error", &error_wrapper);
+  function("mj_saveLastXML", &mj_saveLastXML_wrapper);
+  function("mj_setLengthRange", &mj_setLengthRange_wrapper);
+  // mj_compile is bound using two overloads to handle the optional MjVFS argument,
+  // as using std::optional<MjVFS> caused memory errors due to missing copy/move constructors.
+  function("mj_compile", emscripten::select_overload<std::unique_ptr<MjModel>(const MjSpec&)>(&mj_compile_wrapper_1));
+  function("mj_compile", emscripten::select_overload<std::unique_ptr<MjModel>(const MjSpec&, const MjVFS&)>(&mj_compile_wrapper_2));
 
   emscripten::class_<WasmBuffer<float>>("FloatBuffer")
       .constructor<int>()

wasm/codegen/generators/constants.py

@@ -64,12 +64,16 @@ _SKIPPED_PLUGIN_FUNCTIONS: tuple[str, ...] = (
 # Functions that are bound as class methods
 _SKIPPED_CLASS_METHODS: tuple[str, ...] = (
     # go/keep-sorted start
+    "mj_addBufferVFS",
     "mj_copyData",
     "mj_copyModel",
     "mj_copySpec",
+    "mj_defaultVFS",
     "mj_deleteData",
+    "mj_deleteFileVFS",
     "mj_deleteModel",
     "mj_deleteSpec",
+    "mj_deleteVFS",
     "mj_loadXML",
     "mj_makeData",
     "mj_makeSpec",
@@ -113,11 +117,9 @@ _SKIPPED_ASSET_CACHE_FUNCTIONS: tuple[str, ...] = (
 # Omitted Virtual Filesystem (VFS) functions
 _SKIPPED_VFS_FUNCTIONS: tuple[str, ...] = (
     # go/keep-sorted start
-    "mj_addBufferVFS",
     "mj_addFileVFS",
-    "mj_defaultVFS",
-    "mj_deleteFileVFS",
-    "mj_deleteVFS",
+    "mj_mountVFS",
+    "mj_unmountVFS",
     # go/keep-sorted end
 )
 
@@ -187,17 +189,6 @@ _SKIPPED_UTILITY_FUNCTIONS: tuple[str, ...] = (
     # go/keep-sorted end
 )
 
-# Functions that require special wrappers.
-# These functions are not bound automatically but are written by hand instead.
-MANUAL_WRAPPER_FUNCTIONS: tuple[str, ...] = (
-    # go/keep-sorted start
-    "mj_compile",
-    "mj_saveLastXML",
-    "mj_setLengthRange",
-    "mju_error",
-    # go/keep-sorted end
-)
-
 # List of functions that should be skipped during the code generation process.
 SKIPPED_FUNCTIONS: tuple[str, ...] = (
     _SKIPPED_CLASS_METHODS
@@ -232,12 +223,24 @@ SKIPPED_STRUCTS: tuple[str, ...] = (
     # go/keep-sorted end
 )
 
+# Functions that require special wrappers.
+# These functions are not bound automatically but are written by hand instead.
+MANUAL_WRAPPER_FUNCTIONS: tuple[str, ...] = (
+    # go/keep-sorted start
+    "mj_compile",
+    "mj_saveLastXML",
+    "mj_setLengthRange",
+    "mju_error",
+    # go/keep-sorted end
+)
+
 # Structs for which header generation is done manually.
 # mjvScene is included here because buffer sizes need to be calculated based on
 # introspect doc strings, which was considered a brittle unreliable solution in
 # the past.
 MANUAL_STRUCTS_HEADERS: tuple[str, ...] = (
     "mjvScene",
+    "mjVFS",
 )
 # Structs for which source code generation is done manually.
 MANUAL_STRUCTS_SOURCES: tuple[str, ...] = (

wasm/codegen/generators/functions.py

@@ -272,8 +272,9 @@ def generate(
 
   function_bindings = []
   for func in sorted(functions, key=lambda f: f.name):
-    w = common.wrapped_function_name(func)
-    function_bindings.append(f'function("{func.name}", &{w});')
+    if func.name not in constants.MANUAL_WRAPPER_FUNCTIONS:
+      w = common.wrapped_function_name(func)
+      function_bindings.append(f'function("{func.name}", &{w});')
   bindings_content = "\n".join(function_bindings)
 
   return [

wasm/codegen/templates/bindings.cc

@@ -31,9 +31,9 @@
 #include <vector>
 
 #include <mujoco/mjmodel.h>
+#include <mujoco/mjspec.h>
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mujoco.h>
-#include <mujoco/mjspec.h>
 #include "engine/engine_util_errmem.h"
 #include "wasm/unpack.h"
 
@@ -212,6 +212,27 @@ struct MjvScene {
   std::vector<MjvGLCamera> camera;
 };
 
+struct MjVFS {
+  MjVFS() : ptr_(new mjVFS) { mj_defaultVFS(ptr_); }
+  ~MjVFS() {
+    mj_deleteVFS(ptr_);
+  }
+  void AddBuffer(const std::string& name, const emscripten::val& buffer) {
+    std::vector<uint8_t> vec = emscripten::vecFromJSArray<uint8_t>(buffer);
+    int result = mj_addBufferVFS(ptr_, name.c_str(), vec.data(), vec.size());
+    if (result != 0) {
+      mju_error("Could not add buffer to VFS: %d", result);
+    }
+  }
+  void DeleteFile(const std::string& filename) {
+    mj_deleteFileVFS(ptr_, filename.c_str());
+  }
+  mjVFS* get() const { return ptr_; }
+
+ private:
+  mjVFS* ptr_;
+};
+
 MjModel::MjModel(mjModel* ptr)
     : ptr_(ptr), opt(&ptr->opt), vis(&ptr->vis), stat(&ptr->stat) {}
 
@@ -407,7 +428,7 @@ std::unique_ptr<MjSpec> parseXMLString_wrapper(const std::string &xml) {
   return std::unique_ptr<MjSpec>(new MjSpec(ptr));
 }
 
-std::unique_ptr<MjModel> mj_compile_wrapper(const MjSpec& spec) {
+std::unique_ptr<MjModel> mj_compile_wrapper_1(const MjSpec& spec) {
   mjSpec* spec_ptr = spec.get();
   mjModel* model = mj_compile(spec_ptr, nullptr);
   if (!model || mjs_isWarning(spec_ptr)) {
@@ -416,6 +437,16 @@ std::unique_ptr<MjModel> mj_compile_wrapper(const MjSpec& spec) {
   return std::unique_ptr<MjModel>(new MjModel(model));
 }
 
+std::unique_ptr<MjModel> mj_compile_wrapper_2(const MjSpec& spec, const MjVFS& vfs) {
+  mjSpec* spec_ptr = spec.get();
+  mjVFS* vfs_ptr = vfs.get();
+  mjModel* model = mj_compile(spec_ptr, vfs_ptr);
+  if (!model || mjs_isWarning(spec_ptr)) {
+    mju_error("%s", mjs_getError(spec_ptr));
+  }
+  return std::unique_ptr<MjModel>(new MjModel(model));
+}
+
 void error_wrapper(const String& msg) { mju_error("%s\n", msg.as<const std::string>().data()); }
 
 int mj_saveLastXML_wrapper(const String& filename, const MjModel& m) {
@@ -444,10 +475,20 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
 
   // {{ STRUCTS_BINDINGS }}
 
-  // {{ FUNCTION_BINDINGS }}
+  emscripten::class_<MjVFS>("MjVFS")
+      .constructor<>()
+      .function("addBuffer", &MjVFS::AddBuffer)
+      .function("deleteFile", &MjVFS::DeleteFile);
 
+  // {{ FUNCTION_BINDINGS }}
   function("parseXMLString", &parseXMLString_wrapper, take_ownership());
   function("error", &error_wrapper);
+  function("mj_saveLastXML", &mj_saveLastXML_wrapper);
+  function("mj_setLengthRange", &mj_setLengthRange_wrapper);
+  // mj_compile is bound using two overloads to handle the optional MjVFS argument,
+  // as using std::optional<MjVFS> caused memory errors due to missing copy/move constructors.
+  function("mj_compile", emscripten::select_overload<std::unique_ptr<MjModel>(const MjSpec&)>(&mj_compile_wrapper_1));
+  function("mj_compile", emscripten::select_overload<std::unique_ptr<MjModel>(const MjSpec&, const MjVFS&)>(&mj_compile_wrapper_2));
 
   emscripten::class_<WasmBuffer<float>>("FloatBuffer")
       .constructor<int>()

wasm/tests/bindings_test.ts

@@ -17,7 +17,7 @@ import 'jasmine';
 import {MainModule, MjContact, MjContactVec, MjData, MjLROpt, MjModel,
 MjOption, MjsGeom, MjSolverStat, MjSpec, MjStatistic, MjTimerStat, MjvCamera,
 MjvFigure, MjvGeom, MjvGLCamera, MjvLight, MjvOption, MjvPerturb, MjvScene,
-MjWarningStat} from '../dist/mujoco_wasm.js';
+MjWarningStat, MjVFS} from '../dist/mujoco_wasm.js';
 
 import loadMujoco from '../dist/mujoco_wasm.js'
 
@@ -1803,7 +1803,7 @@ describe('MuJoCo WASM Bindings', () => {
     }
   });
 
-  it('should compile a spec from XML string', () => {
+  it('should compile a spec from XML string with no assets', () => {
     let spec = null;
     let model = null;
     try {
@@ -1826,4 +1826,50 @@ describe('MuJoCo WASM Bindings', () => {
       }
     }
   });
+
+  it('should compile a spec from XML with .obj asset', () => {
+    const xml = `
+    <mujoco>
+      <asset>
+        <mesh file="cube.obj"/>
+      </asset>
+      <worldbody>
+        <geom type="mesh" mesh="cube"/>
+      </worldbody>
+    </mujoco>`;
+
+    const cube1 = `
+    v -1 -1  1
+    v  1 -1  1
+    v -1  1  1
+    v  1  1  1
+    v -1  1 -1
+    v  1  1 -1
+    v -1 -1 -1
+    v  1 -1 -1`;
+
+    let spec: MjSpec|null = null;
+    let model: MjModel|null = null;
+    let vfs: MjVFS|null = null;
+    try {
+      spec = mujoco.parseXMLString(xml);
+      assertExists(spec);
+
+      vfs = new mujoco.MjVFS();
+      vfs.addBuffer('cube.obj', new TextEncoder().encode(cube1));
+      assertExists(vfs);
+
+      model = mujoco.mj_compile(spec, vfs);
+      assertExists(model);
+      expect(model.nmesh).toBe(1);
+
+      const meshId =
+          mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_MESH.value, 'cube');
+      expect(meshId).toBeGreaterThanOrEqual(0);
+    } finally {
+      spec?.delete();
+      vfs?.delete();
+      model?.delete();
+    }
+  });
 });