First-pass implementation of a spec loader for zipped assets.

PiperOrigin-RevId: 861625990
Change-Id: I65cf601cd9aa3aab4915a5a247dda00a6a9aa3db

src/experimental/mjz/mjz_decoder.cc

@@ -0,0 +1,202 @@
+// Copyright 2026 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <cstdarg>
+#include <cstdio>
+#include <cstring>
+#include <filesystem>
+#include <string>
+#include <string_view>
+#include <span>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include <miniz_zip.h>
+#include <mujoco/mjspec.h>
+#include <mujoco/mujoco.h>
+#include "user/user_resource.h"
+
+static void mjPRINTFLIKE(3, 4)
+    SetError(char* error, int error_sz, const char* format, ...) {
+  if (error) {
+    va_list args;
+    va_start(args, format);
+    std::vsnprintf(error, error_sz, format, args);
+    va_end(args);
+  }
+}
+
+// A mjpResourceProvider that reads files from a zip archive.
+//
+// The zip archive itself is provided as a byte buffer in the constructor. This
+// class can then be used to lazily read the contents of the individual files
+// from within the archive as needed.
+class ZipArchiveProvider : public mjpResourceProvider {
+ public:
+  ZipArchiveProvider(std::string name, const void* buffer, int nbuffer,
+                     char* error, int error_sz)
+      : name_(std::move(name)), buffer_((char*)buffer, (char*)buffer + nbuffer) {
+    mjp_defaultResourceProvider(this);
+
+    std::memset(&archive_, 0, sizeof(archive_));
+    if (!mz_zip_reader_init_mem(&archive_, buffer_.data(), nbuffer, 0)) {
+      SetError(error, error_sz, "Zip error: invalid zip archive");
+      return;
+    }
+
+    // Create an index of the files in the archive.
+    const int num_files = mz_zip_reader_get_num_files(&archive_);
+    for (int i = 0; i < num_files; ++i) {
+      mz_zip_archive_file_stat stat;
+      if (!mz_zip_reader_file_stat(&archive_, i, &stat)) {
+        SetError(error, error_sz, "Zip error: failed to stat item %d.", i);
+        files_.clear();
+        return;
+      }
+      const int size = static_cast<int>(stat.m_uncomp_size);
+      if (size == 0) {
+        continue;
+      }
+      files_[stat.m_filename] = FileInfo{i, size, {}};
+    }
+
+    // Setup mjpResourceProvider callbacks.
+    mount = [](mjResource* resource) {
+      return 0;
+    };
+    unmount = [](mjResource* resource) {
+      ZipArchiveProvider* self = (ZipArchiveProvider*)resource->provider;
+      delete self;
+      return 0;
+    };
+    open = [](mjResource* resource) {
+      ZipArchiveProvider* self = (ZipArchiveProvider*)resource->provider;
+      const bool found = self->Contains(resource->name);
+      return found ? 1 : 0;
+    };
+    read = [](mjResource* resource, const void** buffer) {
+      ZipArchiveProvider* self = (ZipArchiveProvider*)resource->provider;
+      std::span<char> bytes = self->Read(resource->name);
+      *buffer = bytes.data();
+      return static_cast<int>(bytes.size());
+    };
+    close = [](mjResource* resource) {
+      // no-op
+    };
+  }
+
+  ~ZipArchiveProvider() {
+    mz_zip_reader_end(&archive_);
+  }
+
+  ZipArchiveProvider(const ZipArchiveProvider&) = delete;
+  ZipArchiveProvider& operator=(const ZipArchiveProvider&) = delete;
+
+  // Returns the path to the root XML model in the archive.
+  std::string GetRootModelPath() const {
+    // Assumes the root has the same name as the archive itself, but with a .xml
+    // extension.
+    const std::filesystem::path path(name_);
+    return (path / path.stem()).string() + ".xml";
+  }
+
+  // Returns true if the archive contains a file with the given name/path.
+  bool Contains(std::string_view name) const {
+    const std::string_view filename = name.substr(name_.size() + 1);
+    return files_.find(filename.data()) != files_.end();
+  }
+
+  // Reads the contents of the file with the given name/path. The contents are
+  // cached internally so that subsequent reads for the same file do not need to
+  // re-read the file from the archive.
+  std::span<char> Read(const std::string& name) {
+    const std::string filename = name.substr(name_.size() + 1);
+    auto it = files_.find(filename);
+    if (it == files_.end()) {
+      return {};
+    }
+
+    FileInfo& info = it->second;
+
+    // Lazily read and store the file contents from the archive.
+    if (info.contents.empty()) {
+      info.contents.resize(info.size);
+      if (!mz_zip_reader_extract_to_mem(&archive_, info.index,
+                                        info.contents.data(), info.size, 0)) {
+        return {};
+      }
+    }
+    return info.contents;
+  }
+
+ private:
+  struct FileInfo {
+    // Index of the file in the archive.
+    int index = 0;
+
+    // Size (in bytes) of the uncompressed file
+    int size = 0;
+
+    // Contents of the uncompressed file.
+    std::vector<char> contents;
+  };
+
+  std::string name_;
+  mz_zip_archive archive_;
+  std::vector<char> buffer_;
+  std::unordered_map<std::string, FileInfo> files_;
+};
+
+static mjSpec* ParseZipBuffer(const void* buffer, int nbuffer, const char* name,
+                              mjVFS* vfs, char* error, int error_sz) {
+  if (error) {
+    error[0] = 0;
+  }
+
+  ZipArchiveProvider* provider =
+      new ZipArchiveProvider(name, buffer, nbuffer, error, error_sz);
+  if (error && error[0]) {
+    return nullptr;
+  }
+
+  const int status = mj_mountVFS(vfs, name, provider);
+  if (status != 0) {
+    SetError(error, error_sz, "Failed to mount zip archive: %s", name);
+    return nullptr;
+  }
+
+  const std::string root = provider->GetRootModelPath();
+  return mj_parseXML(root.c_str(), vfs, error, error_sz);
+}
+
+mjPLUGIN_LIB_INIT {
+  mjpDecoder decoder;
+  decoder.content_type = "application/zip";
+  decoder.extension = ".mjz|.zip";
+  decoder.can_decode = +[](const mjResource* resource) {
+    const char* ext = strrchr(resource->name, '.');
+    return ext ? (!strcmp(ext, ".mjz") || !strcmp(ext, ".zip")) : 0;
+  };
+  decoder.decode = +[](mjResource* resource, const mjVFS* vfs) -> mjSpec* {
+    const char* buffer = nullptr;
+    const int size = mju_readResource(resource, (const void**)&buffer);
+    if (size <= 0) {
+      return nullptr;
+    }
+    return ParseZipBuffer(buffer, size, resource->name, const_cast<mjVFS*>(vfs),
+                          nullptr, 0);
+  };
+  mjp_registerDecoder(&decoder);
+}

test/experimental/mjz/mjz_decoder_test.cc

@@ -0,0 +1,60 @@
+// Copyright 2026 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <mujoco/mjspec.h>
+#include <mujoco/mujoco.h>
+#include "test/fixture.h"
+
+namespace mujoco {
+namespace {
+
+using MjzTest = MujocoTest;
+
+using ::testing::IsNull;
+using ::testing::Not;
+using ::testing::NotNull;
+using ::testing::StrEq;
+
+TEST_F(MjzTest, Parse) {
+  mjVFS vfs;
+  mj_defaultVFS(&vfs);
+  std::string filepath = GetTestDataFilePath("testdata/model.mjz");
+  char err[1000] = "";
+  mjSpec* spec = mj_parse(filepath.c_str(), "", &vfs, err, sizeof(err));
+  EXPECT_THAT(spec, NotNull());
+  EXPECT_THAT(err, StrEq(""));
+  mjModel* model = mj_compile(spec, &vfs);
+  EXPECT_THAT(model, NotNull()) << mjs_getError(spec);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+  mj_deleteVFS(&vfs);
+}
+
+TEST_F(MjzTest, InvalidPath) {
+  mjVFS vfs;
+  mj_defaultVFS(&vfs);
+  std::string filepath = GetTestDataFilePath("testdata/not_a_file.mjz");
+  char err[1000] = "";
+  mjSpec* spec = mj_parse(filepath.c_str(), "", nullptr, err, sizeof(err));
+  EXPECT_THAT(spec, IsNull());
+  EXPECT_THAT(err, Not(StrEq("")));
+  mj_deleteVFS(&vfs);
+}
+
+}  // namespace
+}  // namespace mujoco