<commit_before>#include "jnode.h"

#include <cstring>

#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>

/* 
 * CONSTRUCTORS AND I/O
 */
JNodeTable::JNodeTable(jnid_t max_jnids, bool init_kids, size_t memory_limit) :
  nodes_state(State::ALLOCATED), end_id(0),
  max_id(max_jnids), nodes((JNode*)malloc(sizeof(JNode) * max_id)),
  kid_data(init_kids ? max_id : 0), pst_data(max_id, memory_limit), jxn_data(max_id, memory_limit),
  roots(max_id)
{
  if (nodes == nullptr)
    throw std::bad_alloc();
}

JNodeTable::JNodeTable(char const *filename, jnid_t max_jnids, bool init_kids, size_t memory_limit) :
  nodes_state(State::MAPPED), end_id(0), max_id(max_jnids), nodes(nullptr),
  kid_data(init_kids ? max_id : 0), pst_data(max_id, memory_limit), jxn_data(max_id, memory_limit),
  roots(max_id)
{
  int fd = open(filename, O_RDWR | O_CREAT | O_TRUNC, 00666);
  if (fd == -1)
    throw std::bad_alloc();

  if (posix_fallocate(fd, 0, sizeof(jnid_t) + sizeof(JNode) * max_id) != 0) {
    close(fd);
    throw std::bad_alloc();
  }

  char *const nodes_map = (char*)mmap(nullptr, sizeof(jnid_t) + sizeof(JNode) * max_id,
        PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
  if (nodes_map == MAP_FAILED) {
    close(fd);
    throw std::bad_alloc();
  }
  close(fd);
  nodes = (JNode*)(nodes_map + sizeof(jnid_t));
}

JNodeTable::JNodeTable(char const *filename) :
  nodes_state(State::MAPPED), end_id(0), max_id(0), nodes(nullptr),
  kid_data(0), pst_data(0), jxn_data(0), roots(0)
{
  int fd = open(filename, O_RDWR);
  if (fd == -1)
    throw std::bad_alloc();

  struct stat buf;
  if (fstat(fd, &buf) == -1) {
    close(fd);
    throw std::bad_alloc();
  }
  max_id = (buf.st_size - sizeof(jnid_t)) / sizeof(JNode);

  char *const nodes_map = (char*)mmap(nullptr, sizeof(jnid_t) + sizeof(JNode) * max_id,
        PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
  if (nodes_map == MAP_FAILED) {
    close(fd);
    throw std::bad_alloc();
  }
  close(fd);
  end_id = *((jnid_t*)nodes_map);
  nodes = (JNode*)(nodes_map + sizeof(jnid_t));

  kid_data = std::move(JDataTable<jnid_t>(max_id));
  makeKids();

  //XXX Currently there's nothing you can do with these, so why bother?
  #if 0
  pst_data = std::move(JDataTable<vid_t>(max_id, memory_limit));
  jxn_data = std::move(JDataTable<vid_t>(max_id, memory_limit));

  roots = UnionFind(max_id);
  for (jnid_t id = 0; id < size(); ++id)
    if (parent(id) != INVALID_JNID)
      roots.unify(id, parent(id));
  #endif
}

JNodeTable::JNodeTable(JNode *n, jnid_t end) :
  nodes_state(State::TEMPORARY), end_id(end), max_id(end_id), nodes(n),
  kid_data(max_id), pst_data(0), jxn_data(0), roots(0)
{
  if (nodes != nullptr)
    makeKids();
}

JNodeTable::JNodeTable(JNodeTable &&other) :
  nodes_state(other.nodes_state), end_id(other.end_id), max_id(other.max_id), nodes(other.nodes),
  kid_data(std::move(other.kid_data)), pst_data(std::move(other.pst_data)), jxn_data(std::move(other.jxn_data)),
  roots(std::move(other.roots))
{
  other.nodes_state = State::TEMPORARY;
  other.end_id = 0;
  other.max_id = 0;
  other.nodes = nullptr;
}

JNodeTable::JNodeTable(JNodeTable const &other) :
  nodes_state(State::ALLOCATED), end_id(other.end_id),
  max_id(other.max_id), nodes((JNode*)malloc(sizeof(JNode) * max_id)),
  kid_data(other.kid_data), pst_data(other.pst_data), jxn_data(other.jxn_data), roots(other.roots)
{
  if (nodes == nullptr)
    throw std::bad_alloc();
  std::memcpy(nodes, other.nodes, end_id);
}

JNodeTable::JNodeTable(JNodeTable const &other, jnid_t partial_end) :
  nodes_state(State::ALLOCATED), end_id(other.end_id),
  max_id(other.max_id), nodes((JNode*)malloc(sizeof(JNode) * max_id)),
  kid_data(other.kid_data, partial_end),
  pst_data(other.pst_data, partial_end),
  jxn_data(other.jxn_data, partial_end),
  roots(max_id)
{
  if (nodes == nullptr)
    throw std::bad_alloc();
  std::memcpy(nodes, other.nodes, partial_end);
  
  for (jnid_t id = 0; id < size(); ++id) {
    if (parent(id) < end_id)
      roots.unify(id, parent(id));
    else
      parent(id) = INVALID_JNID;
  }
}

JNodeTable::~JNodeTable() {
  if (nodes_state == State::ALLOCATED)
    free(nodes);
  else if (nodes_state == State::MAPPED) {
    char *const nodes_map = ((char*)nodes) - sizeof(jnid_t);
    *((jnid_t*)nodes_map) = end_id;
    munmap(nodes_map, sizeof(jnid_t) + sizeof(JNode*) * max_id);
  }
}

void JNodeTable::save(char const *const filename) {
  std::ofstream stream(filename, std::ios::binary | std::ios::trunc);
  stream.write((char*)&end_id, sizeof(jnid_t));
  stream.write((char*)nodes, sizeof(JNode) * max_id);
}



/*
 * TREE MERGING METHODS
 */
void JNodeTable::merge(JNodeTable const &lhs, JNodeTable const &rhs, bool const make_kids)
{
  assert(lhs.size() == rhs.size());

  for (jnid_t current = 0; current < lhs.size(); ++current) {
    jnid_t const tmp = newJNode();
    assert(tmp == current);

    if (make_kids) newKids(current, lhs.kids(current).size() + rhs.kids(current).size());

    //XXX For k-way merge, generalize to a JNodeTable list
    auto insert_kids = [current,make_kids,this](JNodeTable const &src)
    {
      for (jnid_t const kid : src.kids(current)) {
        if (!make_kids)
          adopt(kid, current);
        else
          meetKid(kid, current, src.pre_weight(kid));
      }
      pst_weight(current) += src.pst_weight(current);
    };
    insert_kids(lhs);
    insert_kids(rhs);

    if (make_kids)
      adoptKids(current);
  }
}

void mpi_merge_reduction(void *in, void *inout, int *len, MPI_Datatype *datatype) {
  JNodeTable lhs((JNodeTable::JNode*)in, *len);
  JNodeTable rhs((JNodeTable::JNode*)inout, *len);

  JNodeTable tmp(*len, false, 0);
  tmp.merge(lhs, rhs, false);

  assert(*len = tmp.end_id);
  memcpy(inout, tmp.nodes, sizeof(JNodeTable::JNode*) * tmp.end_id);<commit_msg>fixed a very serious sizeof bug that made MPI only work when sizeof(JNode) == 8. thankfully, this was true for the paper<commit_after>#include "jnode.h"

#include <cstring>

#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>

/* 
 * CONSTRUCTORS AND I/O
 */
JNodeTable::JNodeTable(jnid_t max_jnids, bool init_kids, size_t memory_limit) :
  nodes_state(State::ALLOCATED), end_id(0),
  max_id(max_jnids), nodes((JNode*)malloc(sizeof(JNode) * max_id)),
  kid_data(init_kids ? max_id : 0), pst_data(max_id, memory_limit), jxn_data(max_id, memory_limit),
  roots(max_id)
{
  if (nodes == nullptr)
    throw std::bad_alloc();
}

JNodeTable::JNodeTable(char const *filename, jnid_t max_jnids, bool init_kids, size_t memory_limit) :
  nodes_state(State::MAPPED), end_id(0), max_id(max_jnids), nodes(nullptr),
  kid_data(init_kids ? max_id : 0), pst_data(max_id, memory_limit), jxn_data(max_id, memory_limit),
  roots(max_id)
{
  int fd = open(filename, O_RDWR | O_CREAT | O_TRUNC, 00666);
  if (fd == -1)
    throw std::bad_alloc();

  if (posix_fallocate(fd, 0, sizeof(jnid_t) + sizeof(JNode) * max_id) != 0) {
    close(fd);
    throw std::bad_alloc();
  }

  char *const nodes_map = (char*)mmap(nullptr, sizeof(jnid_t) + sizeof(JNode) * max_id,
        PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
  if (nodes_map == MAP_FAILED) {
    close(fd);
    throw std::bad_alloc();
  }
  close(fd);
  nodes = (JNode*)(nodes_map + sizeof(jnid_t));
}

JNodeTable::JNodeTable(char const *filename) :
  nodes_state(State::MAPPED), end_id(0), max_id(0), nodes(nullptr),
  kid_data(0), pst_data(0), jxn_data(0), roots(0)
{
  int fd = open(filename, O_RDWR);
  if (fd == -1)
    throw std::bad_alloc();

  struct stat buf;
  if (fstat(fd, &buf) == -1) {
    close(fd);
    throw std::bad_alloc();
  }
  max_id = (buf.st_size - sizeof(jnid_t)) / sizeof(JNode);

  char *const nodes_map = (char*)mmap(nullptr, sizeof(jnid_t) + sizeof(JNode) * max_id,
        PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
  if (nodes_map == MAP_FAILED) {
    close(fd);
    throw std::bad_alloc();
  }
  close(fd);
  end_id = *((jnid_t*)nodes_map);
  nodes = (JNode*)(nodes_map + sizeof(jnid_t));

  kid_data = std::move(JDataTable<jnid_t>(max_id));
  makeKids();

  //XXX Currently there's nothing you can do with these, so why bother?
  #if 0
  pst_data = std::move(JDataTable<vid_t>(max_id, memory_limit));
  jxn_data = std::move(JDataTable<vid_t>(max_id, memory_limit));

  roots = UnionFind(max_id);
  for (jnid_t id = 0; id < size(); ++id)
    if (parent(id) != INVALID_JNID)
      roots.unify(id, parent(id));
  #endif
}

JNodeTable::JNodeTable(JNode *n, jnid_t end) :
  nodes_state(State::TEMPORARY), end_id(end), max_id(end_id), nodes(n),
  kid_data(max_id), pst_data(0), jxn_data(0), roots(0)
{
  if (nodes != nullptr)
    makeKids();
}

JNodeTable::JNodeTable(JNodeTable &&other) :
  nodes_state(other.nodes_state), end_id(other.end_id), max_id(other.max_id), nodes(other.nodes),
  kid_data(std::move(other.kid_data)), pst_data(std::move(other.pst_data)), jxn_data(std::move(other.jxn_data)),
  roots(std::move(other.roots))
{
  other.nodes_state = State::TEMPORARY;
  other.end_id = 0;
  other.max_id = 0;
  other.nodes = nullptr;
}

JNodeTable::JNodeTable(JNodeTable const &other) :
  nodes_state(State::ALLOCATED), end_id(other.end_id),
  max_id(other.max_id), nodes((JNode*)malloc(sizeof(JNode) * max_id)),
  kid_data(other.kid_data), pst_data(other.pst_data), jxn_data(other.jxn_data), roots(other.roots)
{
  if (nodes == nullptr)
    throw std::bad_alloc();
  std::memcpy(nodes, other.nodes, end_id);
}

JNodeTable::JNodeTable(JNodeTable const &other, jnid_t partial_end) :
  nodes_state(State::ALLOCATED), end_id(other.end_id),
  max_id(other.max_id), nodes((JNode*)malloc(sizeof(JNode) * max_id)),
  kid_data(other.kid_data, partial_end),
  pst_data(other.pst_data, partial_end),
  jxn_data(other.jxn_data, partial_end),
  roots(max_id)
{
  if (nodes == nullptr)
    throw std::bad_alloc();
  std::memcpy(nodes, other.nodes, partial_end);
  
  for (jnid_t id = 0; id < size(); ++id) {
    if (parent(id) < end_id)
      roots.unify(id, parent(id));
    else
      parent(id) = INVALID_JNID;
  }
}

JNodeTable::~JNodeTable() {
  if (nodes_state == State::ALLOCATED)
    free(nodes);
  else if (nodes_state == State::MAPPED) {
    char *const nodes_map = ((char*)nodes) - sizeof(jnid_t);
    *((jnid_t*)nodes_map) = end_id;
    munmap(nodes_map, sizeof(jnid_t) + sizeof(JNode*) * max_id);
  }
}

void JNodeTable::save(char const *const filename) {
  std::ofstream stream(filename, std::ios::binary | std::ios::trunc);
  stream.write((char*)&end_id, sizeof(jnid_t));
  stream.write((char*)nodes, sizeof(JNode) * max_id);
}



/*
 * TREE MERGING METHODS
 */
void JNodeTable::merge(JNodeTable const &lhs, JNodeTable const &rhs, bool const make_kids)
{
  assert(lhs.size() == rhs.size());

  for (jnid_t current = 0; current < lhs.size(); ++current) {
    jnid_t const tmp = newJNode();
    assert(tmp == current);

    if (make_kids) newKids(current, lhs.kids(current).size() + rhs.kids(current).size());

    //XXX For k-way merge, generalize to a JNodeTable list
    auto insert_kids = [current,make_kids,this](JNodeTable const &src)
    {
      for (jnid_t const kid : src.kids(current)) {
        if (!make_kids)
          adopt(kid, current);
        else
          meetKid(kid, current, src.pre_weight(kid));
      }
      pst_weight(current) += src.pst_weight(current);
    };
    insert_kids(lhs);
    insert_kids(rhs);

    if (make_kids)
      adoptKids(current);
  }
}

void mpi_merge_reduction(void *in, void *inout, int *len, MPI_Datatype *datatype) {
  JNodeTable lhs((JNodeTable::JNode*)in, *len);
  JNodeTable rhs((JNodeTable::JNode*)inout, *len);

  JNodeTable tmp(*len, false, 0);
  tmp.merge(lhs, rhs, false);

  assert(*len = tmp.end_id);
  memcpy(inout, tmp.nodes, sizeof(JNodeTable::JNode) * tmp.end_id);
{'fraction_non_alphanumeric': 0.166156202143951, 'fraction_numerical': 0.005530032329419772, 'mean_word_length': 4.648726573762614, 'pattern_counts': {'":': 0, '<': 27, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3894430', 'n_tokens_mistral': 4875, 'n_tokens_neox': 4692, 'n_words': 865}