From 6f9725a1daa14ff10afa5447f2190ad283f5036c Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Tue, 26 May 2026 15:05:54 -0700 Subject: [PATCH 1/2] Make mju_getXMLDependencies read XMLs through the resource provider. The function called tinyxml2's LoadFile directly, which only works on the OS file system. Reading through mju_openResource lets it work against any registered backend (VFS, HTTP, github:, ...). --- src/xml/xml_util.cc | 60 +++++++++++++++++++++--------------- test/xml/xml_utils_test.cc | 62 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 98 insertions(+), 24 deletions(-) diff --git a/src/xml/xml_util.cc b/src/xml/xml_util.cc index b224a372..e5446da3 100644 --- a/src/xml/xml_util.cc +++ b/src/xml/xml_util.cc @@ -222,34 +222,46 @@ void mjCopyError(char* dst, const char* src, int maxlen) { } void mju_getXMLDependencies(const char* filename, mjStringVec* dependencies) { - // load XML file or parse string - tinyxml2::XMLDocument doc; - doc.LoadFile(filename); - - // error checking - if (doc.Error()) { - mju_error("Problem reading XML file '%s': %s", filename, doc.ErrorStr()); + // Open through the resource provider so the function works against any + // registered backend (OS file system, VFS, HTTP, "github:", ...) rather + // than only the OS file system. + mjResource* resource = mju_openResource("", filename, nullptr, nullptr, 0); + if (resource == nullptr) { + mju_error("Could not open '%s'", filename); } - // get top-level element - tinyxml2::XMLElement *root = doc.RootElement(); - if (!root) { - mju_error("XML root element not found"); - } - std::unordered_set files = {filename}; - - std::optional model_dir = std::nullopt; - mjResource *resource = mju_openResource("", filename, nullptr, - nullptr, 0); - if (resource != nullptr) { - const char* dir; - int ndir; - mju_getResourceDir(resource, &dir, &ndir); - model_dir = FilePath(std::string(dir, ndir)); + // Read the XML bytes from the resource. + const void* buffer = nullptr; + int size = mju_readResource(resource, &buffer); + if (size < 0 || !size) { mju_closeResource(resource); + mju_error("Could not read '%s'", filename); } - // Get file references from include and model tags. - AccumulateFiles(files, root, model_dir.value()); + + // Capture the model directory while the resource is still open. + const char* dir = nullptr; + int ndir = 0; + mju_getResourceDir(resource, &dir, &ndir); + FilePath model_dir(std::string(dir, ndir)); + + // Parse from buffer and close (the parsed DOM is independent of the + // resource buffer once Parse returns). + tinyxml2::XMLDocument doc; + tinyxml2::XMLError err = + doc.Parse(static_cast(buffer), static_cast(size)); + mju_closeResource(resource); + + if (err != tinyxml2::XML_SUCCESS) { + mju_error("Problem reading XML file '%s': %s", filename, + doc.ErrorStr() ? doc.ErrorStr() : ""); + } + tinyxml2::XMLElement* root = doc.RootElement(); + if (!root) { + mju_error("XML root element not found in '%s'", filename); + } + + std::unordered_set files = {filename}; + AccumulateFiles(files, root, model_dir); *dependencies = {files.begin(), files.end()}; } diff --git a/test/xml/xml_utils_test.cc b/test/xml/xml_utils_test.cc index c820753d..5d49053c 100644 --- a/test/xml/xml_utils_test.cc +++ b/test/xml/xml_utils_test.cc @@ -13,11 +13,14 @@ // limitations under the License. +#include #include #include #include +#include #include +#include #include #include "test/fixture.h" @@ -42,5 +45,64 @@ TEST_F(MujocoTest, GetXMLDependenciesTest) { kModelPaths.end()}; EXPECT_EQ(dependency_set, expected_dependency_set); } + +// Custom resource provider that serves XML strings from an in-memory map, +// used to exercise the non-filesystem code path. +namespace memxml { +static const std::map* g_files = nullptr; + +int Open(mjResource* resource) { + return g_files && g_files->count(resource->name) ? 1 : 0; +} +int Read(mjResource* resource, const void** buffer) { + auto it = g_files->find(resource->name); + if (it == g_files->end()) return -1; + *buffer = it->second.data(); + return static_cast(it->second.size()); +} +void Close(mjResource* resource) {} +} // namespace memxml + +// Verifies mju_getXMLDependencies works against a non-filesystem resource +// provider (the case the WASM/HTTP build relies on). +TEST_F(MujocoTest, GetXMLDependenciesViaResourceProvider) { + const std::map files = { + {"memxml:/scene.xml", + "\n" + " \n" + " \n" + " \n" + " \n" + " \n" + ""}, + {"memxml:/child.xml", + "\n" + " \n" + " \n" + " \n" + ""}, + }; + memxml::g_files = &files; + + mjpResourceProvider provider = { + .prefix = "memxml", + .open = memxml::Open, + .read = memxml::Read, + .close = memxml::Close, + }; + ASSERT_GT(mjp_registerResourceProvider(&provider), 0); + + mjStringVec dependencies; + mju_getXMLDependencies("memxml:/scene.xml", &dependencies); + std::set dep_set(dependencies.begin(), dependencies.end()); + + EXPECT_THAT(dep_set, testing::UnorderedElementsAre( + "memxml:/scene.xml", + "memxml:/child.xml", + "memxml:/meshes/m.obj", + "memxml:/t.png")); + + memxml::g_files = nullptr; +} } // namespace } // namespace mujoco From 249e9c08b8c04fcc1285418ea0475167caf594f2 Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Tue, 26 May 2026 15:21:01 -0700 Subject: [PATCH 2/2] MuJoCo Live: pre-fetch model assets in parallel before compilation. The prefetcher asks mju_getXMLDependencies for the transitive asset list, fetches every URL with Promise.all, and primes the bytes into the WASM-side FetchCache. When the compiler runs, every resource it opens is already in memory. Cold-load on a typical Menagerie model drops from ~9 s to ~2 s. --- src/experimental/studio/emscripten.cc | 82 ++++++++++++++++-------- src/experimental/studio/live.js | 91 +++++++++++++++++++++++++++ 2 files changed, 146 insertions(+), 27 deletions(-) diff --git a/src/experimental/studio/emscripten.cc b/src/experimental/studio/emscripten.cc index 644c9ea6..b69fae6b 100644 --- a/src/experimental/studio/emscripten.cc +++ b/src/experimental/studio/emscripten.cc @@ -27,6 +27,7 @@ #include #include #include +#include #include #include "experimental/platform/hal/graphics_mode.h" @@ -108,8 +109,9 @@ EM_ASYNC_JS(int, FetchUrl, } }); -// Cache for data fetched via HTTP/HTTPS. Stores the downloaded bytes keyed by -// the resource name (URL) so that read() can return a pointer to the data. +// Cache for data fetched via HTTP/HTTPS. Stores downloaded (or pre-primed) +// bytes keyed by the resource URL so that the resource provider's read() +// can return a pointer into stable storage. class FetchCache { public: static FetchCache& Instance() { @@ -117,52 +119,74 @@ class FetchCache { return instance; } - // Fetches the URL and stores the result. Returns the size (>0) on success. + // Returns the size of the entry for `url`, fetching it over the network + // if necessary. Entries that have been pre-populated via Prime() are + // returned without touching the network. Returns 0 on fetch failure. int Fetch(const char* url) { - char* data = nullptr; + if (auto it = entries_.find(url); it != entries_.end()) { + return static_cast(it->second.size()); + } + char* buf = nullptr; std::int32_t size = 0; - if (!FetchUrl(url, &data, &size)) { + if (!FetchUrl(url, &buf, &size)) { return 0; } - entries_[url] = Entry{UniquePtrWasm(data), size}; - return size; + // FetchUrl hands us a malloc'd buffer; take ownership of the bytes via + // a std::string and release the original allocation. + std::string bytes(buf, size); + std::free(buf); + return static_cast(entries_.emplace(url, std::move(bytes)) + .first->second.size()); } - // Returns pointer and size for a previously fetched URL. - int Read(const char* url, const void** buffer) { + // Pre-populates the cache so a subsequent Fetch() of the same URL hits + // memory instead of the network. First writer wins. + void Prime(std::string url, std::string bytes) { + entries_.try_emplace(std::move(url), std::move(bytes)); + } + + // Hands back a pointer into the cached bytes for the given URL. + // Returns -1 if the URL has not been fetched or primed. + int Read(const char* url, const void** buffer) const { auto it = entries_.find(url); - if (it == entries_.end()) { - return -1; - } - *buffer = it->second.data.get(); - return it->second.size; + if (it == entries_.end()) return -1; + *buffer = it->second.data(); + return static_cast(it->second.size()); } - // Frees the data for a URL. void Close(const char* url) { entries_.erase(url); } private: - struct FreeDeleter { - void operator()(void* p) const { std::free(p); } - }; - template - using UniquePtrWasm = std::unique_ptr; - - struct Entry { - UniquePtrWasm data; - int size; - }; - std::unordered_map entries_; + std::unordered_map entries_; }; // --------------------------------------------------------------------------- -// Javascript-facing function to register an asset. +// Javascript-facing function to register a build-time asset (font, IBL, +// .filamat material) that is fetched via the static AssetRegistry rather +// than over HTTP. void RegisterAsset(std::string filename, std::string contents) { AssetRegistry::Instance().RegisterAsset(std::move(filename), std::move(contents)); } +// Javascript-facing function to pre-populate the HTTP/HTTPS fetch cache. +// Used by the JS-side parallel prefetcher to feed already-downloaded asset +// bytes to MuJoCo's compiler before LoadUrl is invoked, sidestepping the +// otherwise strictly-serial EM_ASYNC_JS fetch chain. +void PrimeFetchCache(std::string url, std::string bytes) { + FetchCache::Instance().Prime(std::move(url), std::move(bytes)); +} + +// Javascript-facing wrapper around mju_getXMLDependencies. Returns the +// list of resource URLs that the given root XML transitively references +// (meshes, textures, includes, ...), per MuJoCo's compiler rules. +std::vector GetXMLDependencies(const std::string& root_url) { + mjStringVec deps; + mju_getXMLDependencies(root_url.c_str(), &deps); + return deps; +} + // Javascript-facing function to initialize the app. void Init(const std::string& title, bool dark_theme) { // Note: dimensions do not matter as window will be resized to fit canvas. @@ -305,7 +329,11 @@ void Deinit() { } EMSCRIPTEN_BINDINGS(studio_bindings) { + // Matches the C-level `mjStringVec` typedef in mjspec.h. + emscripten::register_vector("mjStringVec"); emscripten::function("registerAsset", &RegisterAsset); + emscripten::function("primeFetchCache", &PrimeFetchCache); + emscripten::function("getXMLDependencies", &GetXMLDependencies); emscripten::function("init", &Init); emscripten::function("loadFile", &LoadFile); emscripten::function("loadUrl", &LoadUrl); diff --git a/src/experimental/studio/live.js b/src/experimental/studio/live.js index 8de5230f..13bd874f 100644 --- a/src/experimental/studio/live.js +++ b/src/experimental/studio/live.js @@ -21,6 +21,85 @@ function hideLoading() { loadingOverlay.style.display = 'none'; } +// --------------------------------------------------------------------------- +// Parallel asset prefetcher. +// +// MuJoCo's XML compiler synchronously opens each referenced asset via +// mjpResourceProvider. In WASM that becomes a chain of ASYNCIFY-suspended +// fetch() calls, strictly serial, one RTT per asset. For Menagerie models +// that's ~80 round trips and ~10 s of avoidable wall time on a cold cache. +// +// To sidestep this we let MuJoCo itself enumerate the dependencies (via +// mju_getXMLDependencies, exposed to JS as Module.getXMLDependencies), then +// fetch them all in parallel and feed the bytes into the WASM-side +// FetchCache via Module.primeFetchCache. When Module.loadUrl subsequently +// runs, every resource the compiler opens is already in memory. +// +// Discovery rules live in src/xml/xml_util.cc; the JS only has to know +// the scheme rewrites used by the resource providers in emscripten.cc. +// --------------------------------------------------------------------------- + +// MuJoCo Live custom URL schemes that the C++ resource providers rewrite +// before fetching. JS must do the same so primed cache keys match what +// the providers eventually request. +function resolveScheme(url) { + if (url.startsWith('github:')) { + return 'https://raw.githubusercontent.com/' + url.slice('github:'.length); + } + return url; +} + +async function prefetchModelAssets(rootUrl, onProgress) { + const primed = new Set(); // URLs we've already pushed into FetchCache + const stats = { files: 0, bytes: 0, errors: 0 }; + + // Fetch + prime, dedup against `primed`. Returns true on success. + async function fetchAndPrime(httpUrl) { + if (primed.has(httpUrl)) return true; + primed.add(httpUrl); + try { + const resp = await fetch(httpUrl); + if (!resp.ok) { + console.warn('[prefetch] HTTP', resp.status, httpUrl); + stats.errors++; + return false; + } + const bytes = new Uint8Array(await resp.arrayBuffer()); + Module.primeFetchCache(httpUrl, bytes); + stats.files++; + stats.bytes += bytes.length; + onProgress?.(stats); + return true; + } catch (err) { + console.warn('[prefetch]', httpUrl, err); + stats.errors++; + return false; + } + } + + // Step 1: fetch the root XML so mju_getXMLDependencies (which reads it + // through the resource provider) hits cache instead of the network. + // Also gives the loading screen an immediate progress tick. + const rootHttp = resolveScheme(rootUrl); + if (!await fetchAndPrime(rootHttp)) { + throw new Error(`could not fetch root model ${rootHttp}`); + } + + // Step 2: ask MuJoCo for the transitive dependency list. The call is + // async because reading recursively included XMLs may suspend on + // ASYNCIFY fetches; those XMLs end up in FetchCache as a side effect. + // The result is an embind mjStringVec we copy out and release. + const depsVec = await Module.getXMLDependencies(rootUrl); + const deps = []; + for (let i = 0; i < depsVec.size(); ++i) deps.push(depsVec.get(i)); + depsVec.delete(); + + // Step 3: fetch every dependency in parallel, deduped against URLs + // already primed in step 1 (and earlier deps within this batch). + await Promise.all(deps.map((depUrl) => fetchAndPrime(resolveScheme(depUrl)))); + return stats; +} + var Module = { preRun: [], postRun: [], @@ -111,13 +190,25 @@ var Module = { // loop until the load completes, otherwise renderFrame will // hit "Cannot have multiple async operations in flight". showLoading(); + // Surface prefetch progress on the loading screen. + const dlEl = document.getElementById('loadingDownload'); + const dlFileEl = document.getElementById('loadingDownloadFile'); + if (dlEl) dlEl.style.display = 'block'; + const onProgress = (s) => { + if (dlFileEl) { + dlFileEl.textContent = + `${s.files} files ยท ${(s.bytes/1024/1024).toFixed(1)} MB`; + } + }; requestAnimationFrame(() => { requestAnimationFrame(async () => { try { + await prefetchModelAssets(modelUrl, onProgress); await Module.loadUrl(modelUrl); } catch (error) { console.error('Failed to load model from URL:', error); } finally { + if (dlEl) dlEl.style.display = 'none'; hideLoading(); requestAnimationFrame(Module.animate); }