mujoco

// 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 <mutex>
#include <string>
#include <string_view>
#include <span>
#include <unordered_map>
#include <utility>
#include <vector>

#if defined(__clang__)
  #pragma clang diagnostic push
  #pragma clang diagnostic ignored "-Wunused-function"
#elif defined(__GNUC__)
  #pragma GCC diagnostic push
  #pragma GCC diagnostic ignored "-Wunused-function"
#endif
#include <miniz.h>
#if defined(__clang__)
  #pragma clang diagnostic pop
#elif defined(__GNUC__)
  #pragma GCC diagnostic pop
#endif

#include <mujoco/mjspec.h>
#include <mujoco/mujoco.h>
#include "user/user_resource.h"
#include "user/user_util.h"
#include "user/user_vfs.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_(mujoco::user::FilePath(name).Str()), 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, {}};
    }

    // Look for the root XML model in the archive. We try the following
    // locations:
    // 1. model.xml
    // 2. [archive_name].xml
    // 3. [archive_name]/[archive_name].xml
    // 4. [archive_name]/model.xml
    const mujoco::user::FilePath path(name_);
    const std::string            stem       = path.StripExt().StripPath().Str();
    std::vector<std::string>     candidates = {
        name_ + "/model.xml",
        name_ + "/" + stem + ".xml",
        name_ + "/" + stem + "/" + stem + ".xml",
        name_ + "/" + stem + "/model.xml",
    };

    bool found = false;
    for (const auto& candidate : candidates) {
      if (Contains(candidate)) {
        root_model_ = candidate;
        found       = true;
        break;
      }
    }

    if (!found) {
      SetError(error, error_sz, "Zip error: no root XML file found.");
      return;
    }

    // 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 std::string   path = mujoco::user::FilePath(resource->name).Str();
      if (path == self->name_) { return 1; }
      const bool found = self->Contains(path);
      return found ? 1 : 0;
    };
    read = [](mjResource* resource, const void** buffer) {
      ZipArchiveProvider* self = (ZipArchiveProvider*)resource->provider;
      const std::string   path = mujoco::user::FilePath(resource->name).Str();
      if (path == self->name_) {
        *buffer = self->buffer_.data();
        return static_cast<int>(self->buffer_.size());
      }
      std::span<char> bytes = self->Read(path);
      *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 { return root_model_; }

  // Returns true if the archive contains a file with the given name/path.
  bool Contains(std::string_view name) const {
    // The searched name should be longer than the archive name, it should
    // start with the archive name and have a separator after the archive name.
    if (name.size() <= name_.size() || !name.starts_with(name_) || name[name_.size()] != '/') {
      return false;
    }

    const std::string_view filename = name.substr(name_.size() + 1);
    // Explicitly construct std::string since filename is a string_view slice
    // and not null-terminated.
    return files_.find(std::string(filename)) != 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) {
    if (!Contains(name)) { return {}; }
    const std::string filename = name.substr(name_.size() + 1);
    auto              it       = files_.find(filename);
    FileInfo&         info     = it->second;

    // Lazily read and store the file contents from the archive.
    // The mutex is needed because mz_zip_archive is not thread-safe, and
    // multiple threads may call Read concurrently during parallel compilation.
    std::lock_guard<std::mutex> lock(mutex_);
    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_;
  std::string                               root_model_;
  mz_zip_archive                            archive_;
  std::vector<char>                         buffer_;
  std::unordered_map<std::string, FileInfo> files_;
  mutable std::mutex                        mutex_;
};


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]) {
    delete provider;
    return nullptr;
  }

  const std::string root   = provider->GetRootModelPath();
  const int         status = mj_mountVFS(vfs, name, provider);
  if (status == mujoco::user::VFS::kRepeatedName) {
    // Archive is already mounted in this VFS; reuse the existing provider.
    delete provider;
  } else if (status != 0) {
    delete provider;
    SetError(error, error_sz, "Failed to mount zip archive: %s", name);
    return nullptr;
  }

  return mj_parseXML(root.c_str(), vfs, error, error_sz);
}


mjPLUGIN_LIB_INIT(mjz_decoder) {
  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; }
    char    error[1024] = "";
    mjSpec* spec =
        ParseZipBuffer(buffer, size, resource->name, const_cast<mjVFS*>(vfs), error, sizeof(error));
    if (!spec && error[0]) { mju_warning("%s", error); }
    return spec;
  };
  mjp_registerDecoder(&decoder);
}