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

@@ -27,6 +27,7 @@
 #include <string_view>
 #include <unordered_map>
 #include <utility>
+#include <vector>
 
 #include <mujoco/mujoco.h>
 #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
+// Cache for data fetched via HTTP/HTTPS. Stores downloaded (or pre-primed)
-// the resource name (URL) so that read() can return a pointer to the data.
+// 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<int>(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<char[]>(data), size};
+    // FetchUrl hands us a malloc'd buffer; take ownership of the bytes via
-    return size;
+    // a std::string and release the original allocation.
+    std::string bytes(buf, size);
+    std::free(buf);
+    return static_cast<int>(entries_.emplace(url, std::move(bytes))
+                                .first->second.size());
   }
 
-  // Returns pointer and size for a previously fetched URL.
+  // Pre-populates the cache so a subsequent Fetch() of the same URL hits
-  int Read(const char* url, const void** buffer) {
+  // 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()) {
+    if (it == entries_.end()) return -1;
-      return -1;
+    *buffer = it->second.data();
-    }
+    return static_cast<int>(it->second.size());
-    *buffer = it->second.data.get();
-    return it->second.size;
   }
 
-  // Frees the data for a URL.
   void Close(const char* url) { entries_.erase(url); }
 
  private:
-  struct FreeDeleter {
+  std::unordered_map<std::string, std::string> entries_;
-    void operator()(void* p) const { std::free(p); }
-  };
-  template <typename T>
-  using UniquePtrWasm = std::unique_ptr<T, FreeDeleter>;
-
-  struct Entry {
-    UniquePtrWasm<char[]> data;
-    int size;
-  };
-  std::unordered_map<std::string, Entry> 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<std::string> 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<std::string>("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);

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);
                 }