validate_overlap.cpp

#define SAFETENSORS_CPP_IMPLEMENTATION

#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <iostream>
#include <string>

#include "safetensors.hh"

int main(int argc, char **argv) {
  if (argc != 2) {
    return 2;
  }

  std::ifstream f(argv[1], std::ios::binary);
  std::string contents((std::istreambuf_iterator<char>(f)),
                       std::istreambuf_iterator<char>());
  if (contents.empty()) {
    return 2;
  }

  uint8_t *buf = static_cast<uint8_t *>(std::malloc(contents.size()));
  if (!buf) {
    return 2;
  }
  std::memcpy(buf, contents.data(), contents.size());

  safetensors::safetensors_t st;
  std::string warn;
  std::string err;
  bool ok = safetensors::mmap_from_memory(buf, contents.size(), argv[1], &st,
                                          &warn, &err);
  std::cerr << "mmap_from_memory=" << ok << " err=" << err << "\n";
  if (!ok) {
    return 1;
  }

  err.clear();
  bool valid = safetensors::validate_data_offsets(st, err);
  std::cerr << "validate_data_offsets=" << valid << " err=" << err << "\n";

  safetensors::tensor_t a;
  safetensors::tensor_t b;
  if (!st.tensors.at("a", &a) || !st.tensors.at("b", &b)) {
    return 1;
  }

  const uint8_t *base = st.databuffer_addr;
  std::cerr << "databuffer_size=" << st.databuffer_size << "\n";
  std::cerr << "a=[" << a.data_offsets[0] << "," << a.data_offsets[1]
            << "] bytes=" << char(base[a.data_offsets[0]])
            << char(base[a.data_offsets[1] - 1]) << "\n";
  std::cerr << "b=[" << b.data_offsets[0] << "," << b.data_offsets[1]
            << "] byte=" << char(base[b.data_offsets[0]]) << "\n";
  std::cerr << "alias=" << (base + b.data_offsets[0] == base + a.data_offsets[1] - 1)
            << "\n";

  std::free(buf);
  return valid ? 0 : 1;
}