<commit_before>#include "pbf_writer.hpp"
#include "config.h"
#include "writer_common.hpp"

#include <google/protobuf/io/gzip_stream.h>
#include <google/protobuf/io/zero_copy_stream_impl_lite.h>
#include <osmpbf/osmpbf.h>

#include <boost/unordered_map.hpp>
#include <boost/foreach.hpp>

#include <zlib.h>
#include <arpa/inet.h>
#include <fstream>

namespace bt = boost::posix_time;

namespace {

struct string_table {
  typedef boost::unordered_map<std::string, int> string_map_t;

  string_table() : m_strings(), m_indexed_strings(), m_next_id(1), m_approx_size(0) {}

  int operator()(const std::string &s) {
    string_map_t::iterator itr = m_strings.find(s);
    if (itr == m_strings.end()) {
      int key = m_next_id;
      ++m_next_id;
      m_strings.insert(std::make_pair(s, key));
      m_indexed_strings.push_back(s);
      // keep track of approximate size - strings are stored as the bytes of
      // the string, length prefixed with varint which should be 1 byte if
      // the value is < 128, and more if it's more. OSM strings are generally
      // limited to 255 unicode chars, so unlikely to be > 16k bytes.
      m_approx_size += s.size() + (s.size() > 128 ? 2 : 1);
      return key;

    } else {
      return itr->second;
    }
  }

  size_t approx_size() const { return m_approx_size; }

  void clear() {
    m_strings.clear();
    m_indexed_strings.clear();
    m_next_id = 1;
    m_approx_size = 0;
  }

  void write(OSMPBF::StringTable *st) const {
    // id 0 string is reserved for dense nodes, so just put an empty one in here.
    st->add_s("");
    for (std::vector<std::string>::const_iterator itr = m_indexed_strings.begin();
         itr != m_indexed_strings.end(); ++itr) {
      st->add_s(*itr);
    }
  }

  string_map_t m_strings;
  std::vector<std::string> m_indexed_strings;
  int m_next_id;
  size_t m_approx_size;
};

// simple function to calculate the delta between the last value of
// something and the current value, returning the difference and
// updating the last value.
template <typename T>
inline T delta(T &last, T cur) {
  T d = cur - last;
  last = cur;
  return d;
}

} // anonymous namespace

struct pbf_writer::pimpl {
  enum element_type {
    element_NULL,
    element_CHANGESET,
    element_NODE,
    element_WAY,
    element_RELATION
  };

  pimpl(const std::string &out_name, const bt::ptime &now, bool history_format,
        const user_map_t &user_map, const boost::program_options::variables_map &options) 
    : num_elements(0), buffer(), out(out_name.c_str()), str_table(),
      pblock(), m_dense_section(), pgroup(pblock.add_primitivegroup()), 
      current_node(NULL), current_way(NULL), current_relation(NULL),
      m_byte_limit(int(0.25 * OSMPBF::max_uncompressed_blob_size)),
      m_current_element(element_NULL),
      m_last_way_node_ref(0),
      m_last_relation_member_ref(0),
      m_est_pblock_size(0),
      m_history_format(history_format),
      m_user_map(user_map),
      m_changeset_user_map(),
      m_recheck_elements(int(element_RELATION) + 1) {

    // different re-check limits per type so that we can better
    // adapt to the different sizes of elements, and hit the
    // byte limit without overflowing it.
    m_recheck_elements[element_NULL] = 1;
    m_recheck_elements[element_CHANGESET] = 1;
    m_recheck_elements[element_NODE] = 16000;
    m_recheck_elements[element_WAY] = 8000;
    m_recheck_elements[element_RELATION] = 4000;

    m_dense_nodes = options["dense-nodes"].as<bool>();

    write_header_block(now);
  }

  ~pimpl() {
  }

  void write_header_block(const bt::ptime &now) {
    using namespace OSMPBF;

    HeaderBlock header;
    HeaderBBox *bbox = header.mutable_bbox();
    bbox->set_left  (-180L * lonlat_resolution);
    bbox->set_right ( 180L * lonlat_resolution);
    bbox->set_top   (  90L * lonlat_resolution);
    bbox->set_bottom( -90L * lonlat_resolution);

    header.add_required_features("OsmSchema-V" OSM_VERSION_TEXT);
    if (m_history_format) { 
      header.add_required_features("HistoricalInformation");
    }
    if (m_dense_nodes) {
      header.add_required_features("DenseNodes");
    }
    header.add_optional_features("Has_Metadata");
    header.add_optional_features("Sort.Type_then_ID");
    header.set_writingprogram(PACKAGE_STRING);
    header.set_source(OSM_API_ORIGIN);
#ifndef WITH_OLD_OSMPBF
    header.set_osmosis_replication_timestamp((now - bt::from_time_t(time_t(0))).total_seconds());
#endif

    write_blob(header, "OSMHeader");
  }

  void write_blob(const google::protobuf::MessageLite &message, const std::string &type) {
    using namespace OSMPBF;
    using google::protobuf::io::StringOutputStream;
    using google::protobuf::io::GzipOutputStream;

    Blob blob;
    size_t uncompressed_size = message.ByteSize();
    // sanity check - if we're about to violate the OSMPBF format rules
    // then we'd rather stop than ship an invalid file.
    if (uncompressed_size >= OSMPBF::max_uncompressed_blob_size) {
      std::ostringstream ostr;
      ostr << "Unable to write block of type " << type << ", uncompressed size " << uncompressed_size
           << " because it is larger than the maximum allowed " << OSMPBF::max_uncompressed_blob_size
           << "." << std::endl;
      BOOST_THROW_EXCEPTION(std::runtime_error(ostr.str()));
    }
    blob.set_raw_size(uncompressed_size);

    std::string str;
    StringOutputStream string_stream(&str);
    GzipOutputStream::Options options;
    options.format = GzipOutputStream::ZLIB;
    options.compression_level = 9;
    GzipOutputStream gzip_stream(&string_stream, options);
    message.SerializeToZeroCopyStream(&gzip_stream);
    gzip_stream.Close();
    blob.set_zlib_data(str);
    str.clear();

    BlobHeader blob_header;
    blob_header.set_type(type);
    blob_header.set_datasize(blob.ByteSize());

    int blob_header_size = blob_header.ByteSize();
    if (blob_header_size < 0) {
      std::ostringstream ostr;
      ostr << "Unable to write blob header size " << blob_header_size 
           << " because it will not correctly cast to uint32_t.";
      BOOST_THROW_EXCEPTION(std::runtime_error(ostr.str()));
    }
    uint32_t bh_size = htonl(uint32_t(blob_header_size));
    out.write((char *)&bh_size, sizeof bh_size);
    out << blob_header.SerializeAsString();
    out << blob.SerializeAsString();
    out.flush();
  }

  void check_overflow(element_type type) {
    if ((m_current_element == element_NULL) ||
        (m_current_element == element_CHANGESET)) {  // <- to deal with OSMPBF brokenness...
      m_current_element = type;
    }

    if ((m_current_element != type) || 
        (num_elements >= m_recheck_elements[m_current_element])) {
      m_est_pblock_size += pgroup->ByteSize();
      bool new_block = ((m_current_element != type) || 
                        ((m_est_pblock_size + str_table.approx_size()) >= m_byte_limit));

      if (new_block) {
        str_table.write(pblock.mutable_stringtable());
        write_blob(pblock, "OSMData");
        pblock.Clear();
        str_table.clear();
        
        m_current_element = type;
        m_est_pblock_size = 0;
      }

      m_last_dense_id = 0;
      m_last_dense_lat = 0;
      m_last_dense_lon = 0;
      m_last_dense_timestamp = 0;
      m_last_dense_changeset = 0;
      m_last_dense_uid = 0;
      m_last_dense_user_sid = 0;
      pgroup = pblock.add_primitivegroup();
      num_elements = 0;
      current_node = NULL;
      current_way = NULL;
      current_relation = NULL;
    }
  }

  template <typename T>
  void set_info(const T &t, OSMPBF::Info *info) {
    static bt::ptime epoch = bt::from_time_t(time_t(0));
    
    info->set_version(t.version);
    info->set_timestamp((t.timestamp - epoch).total_seconds());
    info->set_changeset(t.changeset_id);
    // if we are doing a history file, and the default of visible=true
    // doesn't apply, then we need to explicitly set visible=false.
    if (m_history_format && !t.visible) {
      info->set_visible(t.visible);
    }
    // set the uid and user information, if the user is public
    std::map<int64_t, int64_t>::const_iterator itr = m_changeset_user_map.find(t.changeset_id);
    if (itr != m_changeset_user_map.end()) {
      user_map_t::const_iterator jtr = m_user_map.find(itr->second);
      if (jtr != m_user_map.end()) {
        info->set_uid(jtr->first);
        info->set_user_sid(str_table(jtr->second));
      }
    }
  }

  void add_changeset(const changeset &cs) {
    // looks like OSMPBF is broken and doesn't really support this.
    check_overflow(element_CHANGESET);
  }

  void add_node(const node &n) {
    check_overflow(element_NODE);
    if (m_dense_nodes) return add_dense_node(n);

    current_node = pgroup->add_nodes();
    current_node->set_id(n.id);
    // deleted nodes don't have lat/lon attributes
    if (n.visible) {
      current_node->set_lat(n.latitude);
      current_node->set_lon(n.longitude);
    } else {
      // however, PBF doesn't allow you not to set these attributes,
      // so we have to set them to some null value. (0, 0) is, sadly,
      // the traditional value for these, even though it is valid.
      current_node->set_lat(0);
      current_node->set_lon(0);
    }
    set_info(n, current_node->mutable_info());

    ++num_elements;
  }

  void add_dense_node(const node &n) {
    static bt::ptime epoch = bt::from_time_t(time_t(0));
    current_node = NULL;
    m_dense_section = pgroup->mutable_dense();
    m_dense_section->add_id(delta(m_last_dense_id, n.id));
    m_dense_section->add_lon(delta(m_last_dense_lon, int64_t(n.visible ? n.longitude : 0)));
    m_dense_section->add_lat(delta(m_last_dense_lat, int64_t(n.visible ? n.latitude : 0)));
    OSMPBF::DenseInfo* info = m_dense_section->mutable_denseinfo();
    info->add_version(n.version);
    info->add_timestamp(delta(m_last_dense_timestamp, int64_t((n.timestamp - epoch).total_seconds())));
    info->add_changeset(delta(m_last_dense_changeset, n.changeset_id));
    // if we are doing a history file, and the default of visible=true
    // doesn't apply, then we need to explicitly set visible=false.
    if (m_history_format && !n.visible) {<commit_msg>Add visible flag for all dense nodes when writing history file.<commit_after>#include "pbf_writer.hpp"
#include "config.h"
#include "writer_common.hpp"

#include <google/protobuf/io/gzip_stream.h>
#include <google/protobuf/io/zero_copy_stream_impl_lite.h>
#include <osmpbf/osmpbf.h>

#include <boost/unordered_map.hpp>
#include <boost/foreach.hpp>

#include <zlib.h>
#include <arpa/inet.h>
#include <fstream>

namespace bt = boost::posix_time;

namespace {

struct string_table {
  typedef boost::unordered_map<std::string, int> string_map_t;

  string_table() : m_strings(), m_indexed_strings(), m_next_id(1), m_approx_size(0) {}

  int operator()(const std::string &s) {
    string_map_t::iterator itr = m_strings.find(s);
    if (itr == m_strings.end()) {
      int key = m_next_id;
      ++m_next_id;
      m_strings.insert(std::make_pair(s, key));
      m_indexed_strings.push_back(s);
      // keep track of approximate size - strings are stored as the bytes of
      // the string, length prefixed with varint which should be 1 byte if
      // the value is < 128, and more if it's more. OSM strings are generally
      // limited to 255 unicode chars, so unlikely to be > 16k bytes.
      m_approx_size += s.size() + (s.size() > 128 ? 2 : 1);
      return key;

    } else {
      return itr->second;
    }
  }

  size_t approx_size() const { return m_approx_size; }

  void clear() {
    m_strings.clear();
    m_indexed_strings.clear();
    m_next_id = 1;
    m_approx_size = 0;
  }

  void write(OSMPBF::StringTable *st) const {
    // id 0 string is reserved for dense nodes, so just put an empty one in here.
    st->add_s("");
    for (std::vector<std::string>::const_iterator itr = m_indexed_strings.begin();
         itr != m_indexed_strings.end(); ++itr) {
      st->add_s(*itr);
    }
  }

  string_map_t m_strings;
  std::vector<std::string> m_indexed_strings;
  int m_next_id;
  size_t m_approx_size;
};

// simple function to calculate the delta between the last value of
// something and the current value, returning the difference and
// updating the last value.
template <typename T>
inline T delta(T &last, T cur) {
  T d = cur - last;
  last = cur;
  return d;
}

} // anonymous namespace

struct pbf_writer::pimpl {
  enum element_type {
    element_NULL,
    element_CHANGESET,
    element_NODE,
    element_WAY,
    element_RELATION
  };

  pimpl(const std::string &out_name, const bt::ptime &now, bool history_format,
        const user_map_t &user_map, const boost::program_options::variables_map &options) 
    : num_elements(0), buffer(), out(out_name.c_str()), str_table(),
      pblock(), m_dense_section(), pgroup(pblock.add_primitivegroup()), 
      current_node(NULL), current_way(NULL), current_relation(NULL),
      m_byte_limit(int(0.25 * OSMPBF::max_uncompressed_blob_size)),
      m_current_element(element_NULL),
      m_last_way_node_ref(0),
      m_last_relation_member_ref(0),
      m_est_pblock_size(0),
      m_history_format(history_format),
      m_user_map(user_map),
      m_changeset_user_map(),
      m_recheck_elements(int(element_RELATION) + 1) {

    // different re-check limits per type so that we can better
    // adapt to the different sizes of elements, and hit the
    // byte limit without overflowing it.
    m_recheck_elements[element_NULL] = 1;
    m_recheck_elements[element_CHANGESET] = 1;
    m_recheck_elements[element_NODE] = 16000;
    m_recheck_elements[element_WAY] = 8000;
    m_recheck_elements[element_RELATION] = 4000;

    m_dense_nodes = options["dense-nodes"].as<bool>();

    write_header_block(now);
  }

  ~pimpl() {
  }

  void write_header_block(const bt::ptime &now) {
    using namespace OSMPBF;

    HeaderBlock header;
    HeaderBBox *bbox = header.mutable_bbox();
    bbox->set_left  (-180L * lonlat_resolution);
    bbox->set_right ( 180L * lonlat_resolution);
    bbox->set_top   (  90L * lonlat_resolution);
    bbox->set_bottom( -90L * lonlat_resolution);

    header.add_required_features("OsmSchema-V" OSM_VERSION_TEXT);
    if (m_history_format) { 
      header.add_required_features("HistoricalInformation");
    }
    if (m_dense_nodes) {
      header.add_required_features("DenseNodes");
    }
    header.add_optional_features("Has_Metadata");
    header.add_optional_features("Sort.Type_then_ID");
    header.set_writingprogram(PACKAGE_STRING);
    header.set_source(OSM_API_ORIGIN);
#ifndef WITH_OLD_OSMPBF
    header.set_osmosis_replication_timestamp((now - bt::from_time_t(time_t(0))).total_seconds());
#endif

    write_blob(header, "OSMHeader");
  }

  void write_blob(const google::protobuf::MessageLite &message, const std::string &type) {
    using namespace OSMPBF;
    using google::protobuf::io::StringOutputStream;
    using google::protobuf::io::GzipOutputStream;

    Blob blob;
    size_t uncompressed_size = message.ByteSize();
    // sanity check - if we're about to violate the OSMPBF format rules
    // then we'd rather stop than ship an invalid file.
    if (uncompressed_size >= OSMPBF::max_uncompressed_blob_size) {
      std::ostringstream ostr;
      ostr << "Unable to write block of type " << type << ", uncompressed size " << uncompressed_size
           << " because it is larger than the maximum allowed " << OSMPBF::max_uncompressed_blob_size
           << "." << std::endl;
      BOOST_THROW_EXCEPTION(std::runtime_error(ostr.str()));
    }
    blob.set_raw_size(uncompressed_size);

    std::string str;
    StringOutputStream string_stream(&str);
    GzipOutputStream::Options options;
    options.format = GzipOutputStream::ZLIB;
    options.compression_level = 9;
    GzipOutputStream gzip_stream(&string_stream, options);
    message.SerializeToZeroCopyStream(&gzip_stream);
    gzip_stream.Close();
    blob.set_zlib_data(str);
    str.clear();

    BlobHeader blob_header;
    blob_header.set_type(type);
    blob_header.set_datasize(blob.ByteSize());

    int blob_header_size = blob_header.ByteSize();
    if (blob_header_size < 0) {
      std::ostringstream ostr;
      ostr << "Unable to write blob header size " << blob_header_size 
           << " because it will not correctly cast to uint32_t.";
      BOOST_THROW_EXCEPTION(std::runtime_error(ostr.str()));
    }
    uint32_t bh_size = htonl(uint32_t(blob_header_size));
    out.write((char *)&bh_size, sizeof bh_size);
    out << blob_header.SerializeAsString();
    out << blob.SerializeAsString();
    out.flush();
  }

  void check_overflow(element_type type) {
    if ((m_current_element == element_NULL) ||
        (m_current_element == element_CHANGESET)) {  // <- to deal with OSMPBF brokenness...
      m_current_element = type;
    }

    if ((m_current_element != type) || 
        (num_elements >= m_recheck_elements[m_current_element])) {
      m_est_pblock_size += pgroup->ByteSize();
      bool new_block = ((m_current_element != type) || 
                        ((m_est_pblock_size + str_table.approx_size()) >= m_byte_limit));

      if (new_block) {
        str_table.write(pblock.mutable_stringtable());
        write_blob(pblock, "OSMData");
        pblock.Clear();
        str_table.clear();
        
        m_current_element = type;
        m_est_pblock_size = 0;
      }

      m_last_dense_id = 0;
      m_last_dense_lat = 0;
      m_last_dense_lon = 0;
      m_last_dense_timestamp = 0;
      m_last_dense_changeset = 0;
      m_last_dense_uid = 0;
      m_last_dense_user_sid = 0;
      pgroup = pblock.add_primitivegroup();
      num_elements = 0;
      current_node = NULL;
      current_way = NULL;
      current_relation = NULL;
    }
  }

  template <typename T>
  void set_info(const T &t, OSMPBF::Info *info) {
    static bt::ptime epoch = bt::from_time_t(time_t(0));
    
    info->set_version(t.version);
    info->set_timestamp((t.timestamp - epoch).total_seconds());
    info->set_changeset(t.changeset_id);
    // if we are doing a history file, and the default of visible=true
    // doesn't apply, then we need to explicitly set visible=false.
    if (m_history_format && !t.visible) {
      info->set_visible(t.visible);
    }
    // set the uid and user information, if the user is public
    std::map<int64_t, int64_t>::const_iterator itr = m_changeset_user_map.find(t.changeset_id);
    if (itr != m_changeset_user_map.end()) {
      user_map_t::const_iterator jtr = m_user_map.find(itr->second);
      if (jtr != m_user_map.end()) {
        info->set_uid(jtr->first);
        info->set_user_sid(str_table(jtr->second));
      }
    }
  }

  void add_changeset(const changeset &cs) {
    // looks like OSMPBF is broken and doesn't really support this.
    check_overflow(element_CHANGESET);
  }

  void add_node(const node &n) {
    check_overflow(element_NODE);
    if (m_dense_nodes) return add_dense_node(n);

    current_node = pgroup->add_nodes();
    current_node->set_id(n.id);
    // deleted nodes don't have lat/lon attributes
    if (n.visible) {
      current_node->set_lat(n.latitude);
      current_node->set_lon(n.longitude);
    } else {
      // however, PBF doesn't allow you not to set these attributes,
      // so we have to set them to some null value. (0, 0) is, sadly,
      // the traditional value for these, even though it is valid.
      current_node->set_lat(0);
      current_node->set_lon(0);
    }
    set_info(n, current_node->mutable_info());

    ++num_elements;
  }

  void add_dense_node(const node &n) {
    static bt::ptime epoch = bt::from_time_t(time_t(0));
    current_node = NULL;
    m_dense_section = pgroup->mutable_dense();
    m_dense_section->add_id(delta(m_last_dense_id, n.id));
    m_dense_section->add_lon(delta(m_last_dense_lon, int64_t(n.visible ? n.longitude : 0)));
    m_dense_section->add_lat(delta(m_last_dense_lat, int64_t(n.visible ? n.latitude : 0)));
    OSMPBF::DenseInfo* info = m_dense_section->mutable_denseinfo();
    info->add_version(n.version);
    info->add_timestamp(delta(m_last_dense_timestamp, int64_t((n.timestamp - epoch).total_seconds())));
    info->add_changeset(delta(m_last_dense_changeset, n.changeset_id));
    // if we are doing a history file, we need to set the visible flag
    // for all entries in the dense node table, as this array is indexed
    // into by position to get the visibility flag.
    if (m_history_format) {
{'fraction_non_alphanumeric': 0.11951255322272794, 'fraction_numerical': 0.008613517349385797, 'mean_word_length': 3.520796460176991, 'pattern_counts': {'":': 0, '<': 91, '<?xml version=': 0, '>': 107, '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': '6770783', 'n_tokens_mistral': 7140, 'n_tokens_neox': 6879, 'n_words': 1465}