#include <protozero/pbf_message.hpp>

#include <iostream>

#include <vector>

#include "pbf_reader.h"

#include "helpers.h"

// Where pbf_processor.cpp has higher-level routines that populate our structures,

// pbf_reader.cpp has low-level tools that interact with the protobuf.

//

// The lifetime of an object is only until someone calls a readXyz function at

// the same or higher level.

// - e.g. readPrimitiveGroup invalidates the result of a prior readPrimitiveGroup call,

// but not the result of a prior readBlob call

//

// This allows us to re-use buffers to minimize heap churn and allocation cost.

//

// If you want to persist the data beyond that, you must make a copy in memory

// that you own.

PbfReader::BlobHeader PbfReader::PbfReader::readBlobHeader(std::istream& input) {

// See https://wiki.openstreetmap.org/wiki/PBF_Format#File_format

unsigned int size;

input.read((char*)&size, sizeof(size));

if (input.eof()) {

return {"eof", -1};

}

endian_swap(size);

std::vector<char> data;

data.resize(size);

input.read(&data[0], size);

if (input.eof())

throw std::runtime_error("readBlobHeader: unexpected eof");

protozero::pbf_message<Schema::BlobHeader> message{&data[0], data.size()};

std::string type;

int32_t datasize = -1;

while (message.next()) {

switch (message.tag()) {

case Schema::BlobHeader::required_string_type:

type = message.get_string();

break;

case Schema::BlobHeader::required_int32_datasize:

datasize = message.get_int32();

break;

default:

// ignore data for unknown tags to allow for future extensions

// std::cout << "BlobHeader: unknown tag: " << std::to_string(static_cast<uint32_t>(message.tag())) << std::endl;

message.skip();

}

}

if (type.empty())

throw std::runtime_error("BlobHeader type is missing");

if (datasize == -1)

throw std::runtime_error("BlobHeader datasize is missing");

return { type, datasize };

}

protozero::data_view PbfReader::PbfReader::readBlob(int32_t datasize, std::istream& input) {

blobStorage.resize(datasize);

input.read(&blobStorage[0], datasize);

if (input.eof())

throw std::runtime_error("readBlob: unexpected eof");

enum class BlobDataType { None, Raw, Zlib };

int32_t rawSize = -1;

BlobDataType dataType = BlobDataType::None;

protozero::data_view view;

protozero::pbf_message<Schema::Blob> message{&blobStorage[0], blobStorage.size()};

while (message.next()) {

switch (message.tag()) {

case Schema::Blob::optional_int32_raw_size:

rawSize = message.get_int32();

break;

case Schema::Blob::oneof_data_bytes_raw:

dataType = BlobDataType::Raw;

view = message.get_view();

break;

case Schema::Blob::oneof_data_bytes_zlib_data:

dataType = BlobDataType::Zlib;

view = message.get_view();

break;

case Schema::Blob::oneof_data_bytes_lzma_data:

case Schema::Blob::oneof_data_bytes_lz4_data:

case Schema::Blob::oneof_data_bytes_zstd_data:

throw std::runtime_error("Blob: unsupported compression tag: " + std::to_string(static_cast<uint32_t>(message.tag())));

default:

throw std::runtime_error("Blob: unknown tag: " + std::to_string(static_cast<uint32_t>(message.tag())));

}

}

if (dataType == BlobDataType::None)

throw std::runtime_error("Blob: missing data");

if (dataType == BlobDataType::Raw)

return view;

if (rawSize != -1)

blobStorage2.resize(rawSize);

else

blobStorage2.clear();

decompress_string(blobStorage2, view.data(), view.size(), false);

return { blobStorage2.data(), blobStorage2.size() };

}

PbfReader::HeaderBBox PbfReader::PbfReader::readHeaderBBox(protozero::data_view data) {

HeaderBBox box{0, 0, 0, 0};

protozero::pbf_message<Schema::HeaderBBox> message{data};

while (message.next()) {

switch (message.tag()) {

case Schema::HeaderBBox::required_sint64_left:

box.minLon = message.get_sint64() / 1000000000.0;

break;

case Schema::HeaderBBox::required_sint64_right:

box.maxLon = message.get_sint64() / 1000000000.0;

break;

case Schema::HeaderBBox::required_sint64_bottom:

box.minLat = message.get_sint64() / 1000000000.0;

break;

case Schema::HeaderBBox::required_sint64_top:

box.maxLat = message.get_sint64() / 1000000000.0;

break;

default:

throw std::runtime_error("HeaderBBox: unknown tag: " + std::to_string(static_cast<uint32_t>(message.tag())));

}

}

return box;

}

PbfReader::HeaderBlock PbfReader::PbfReader::readHeaderBlock(protozero::data_view data) {

HeaderBlock block{false};

protozero::pbf_message<Schema::HeaderBlock> message{data};

while (message.next()) {

switch (message.tag()) {

case Schema::HeaderBlock::optional_HeaderBBox_bbox:

block.hasBbox = true;

block.bbox = PbfReader::readHeaderBBox(message.get_view());

break;

case Schema::HeaderBlock::repeated_string_optional_features: {

const auto feature = message.get_string();

block.optionalFeatures.insert(feature);

break;

}

default:

// ignore data for unknown tags to allow for future extensions

//std::cout << "HeaderBlock: unknown tag: " << std::to_string(static_cast<uint32_t>(message.tag())) << std::endl;

message.skip();

}

}

return block;

}

void PbfReader::PbfReader::readStringTable(protozero::data_view data, std::vector<protozero::data_view>& stringTable) {

protozero::pbf_message<Schema::StringTable> message{data};

while (message.next()) {

switch (message.tag()) {

case Schema::StringTable::repeated_bytes_s:

stringTable.push_back(message.get_view());

break;

default:

throw std::runtime_error("StringTable: unknown tag: " + std::to_string(static_cast<uint32_t>(message.tag())));

}

}

}

PbfReader::PrimitiveBlock& PbfReader::PbfReader::readPrimitiveBlock(protozero::data_view data) {

pb.stringTable.clear();

pb.internalGroups.clear();

protozero::pbf_message<Schema::PrimitiveBlock> message{data};

while (message.next()) {

switch (message.tag()) {

case Schema::PrimitiveBlock::required_StringTable_stringtable:

// Most of our use cases require the string table, so we eagerly

// initialize it.

PbfReader::readStringTable(message.get_view(), pb.stringTable);

break;

case Schema::PrimitiveBlock::repeated_PrimitiveGroup_primitivegroup: {

pb.internalGroups.push_back(PrimitiveGroup(

message.get_view(),

denseNodes,

way,

relation

));

break;

}

default:

// ignore data for unknown tags to allow for future extensions

//std::cout << "HeaderBlock: unknown tag: " << std::to_string(static_cast<uint32_t>(message.tag())) << std::endl;

message.skip();

}

}

pb.groupsImpl = PrimitiveBlock::PrimitiveGroups(pb.internalGroups);

return pb;

}

void PbfReader::DenseNodes::readDenseNodes(protozero::data_view data) {

protozero::pbf_message<Schema::DenseNodes> message{data};

uint64_t id = 0;

int32_t lon = 0, lat = 0;

while (message.next()) {

switch (message.tag()) {

case Schema::DenseNodes::repeated_sint64_id: {

auto pi = message.get_packed_sint64();

for (auto i : pi) {

id += i;

ids.push_back(id);

}

break;

} case Schema::DenseNodes::repeated_sint64_lat: {

auto pi = message.get_packed_sint64();

for (auto i : pi) {

lat += i;

lats.push_back(lat);

}

break;

}

case Schema::DenseNodes::repeated_sint64_lon: {

auto pi = message.get_packed_sint64();

for (auto i : pi) {

lon += i;

lons.push_back(lon);

}

break;

}

case Schema::DenseNodes::repeated_int32_keys_vals: {

auto pi = message.get_packed_int32();

for (auto kv : pi) {

keyValues.push_back(kv);

}

break;

}

default:

// ignore data for unknown tags to allow for future extensions

//std::cout << "HeaderBlock: unknown tag: " << std::to_string(static_cast<uint32_t>(message.tag())) << std::endl;

message.skip();

}

}

for (uint32_t cur = 0, prev = 0; cur < keyValues.size(); cur++) {

if (keyValues[cur] == 0) {

tagStart.push_back(prev);

tagEnd.push_back(cur);

prev = cur + 1;

}

}

while(tagStart.size() < ids.size()) {

tagStart.push_back(0);

tagEnd.push_back(0);

}

}

PbfReader::PrimitiveGroup::PrimitiveGroup(

protozero::data_view data,

DenseNodes& denseNodes,

Way& way,

Relation& relation

):

data(data),

denseNodes(denseNodes),

internalWays({this, way}),

internalRelations({this, relation}),

denseNodesInitialized(false) {

}

int32_t PbfReader::PrimitiveGroup::translateNodeKeyValue(int32_t i) const {

return denseNodes.keyValues.at(i);

}

protozero::data_view PbfReader::PrimitiveGroup::getDataView() {

return data;

}

void PbfReader::PrimitiveGroup::ensureData() {

// Reset our thread locals.

denseNodes.clear();

internalWays.pg = this;

internalRelations.pg = this;

protozero::pbf_message<Schema::PrimitiveGroup> message{data};

if (message.next()) {

switch (message.tag()) {

case Schema::PrimitiveGroup::repeated_Node_nodes:

throw std::runtime_error("PrimitiveGroup: non-dense Nodes are not supported");

break;

case Schema::PrimitiveGroup::optional_DenseNodes_dense:

internalType = PrimitiveGroupType::DenseNodes;

denseNodes.readDenseNodes(message.get_view());

break;

case Schema::PrimitiveGroup::repeated_Way_ways:

internalType = PrimitiveGroupType::Way;

break;

case Schema::PrimitiveGroup::repeated_Relation_relations:

internalType = PrimitiveGroupType::Relation;

break;

case Schema::PrimitiveGroup::repeated_ChangeSet_changesets:

internalType = PrimitiveGroupType::ChangeSet;

break;

default:

throw std::runtime_error("PrimitiveGroup: unknown tag: " + std::to_string(static_cast<uint32_t>(message.tag())));

}

}

}

PbfReader::DenseNodes& PbfReader::PrimitiveGroup::nodes() const { return denseNodes; };

PbfReader::PrimitiveBlock::PrimitiveGroups& PbfReader::PrimitiveBlock::groups() { return groupsImpl; };

void PbfReader::DenseNodes::clear() {

ids.clear();

lons.clear();

lats.clear();

tagStart.clear();

tagEnd.clear();

keyValues.clear();

}

bool PbfReader::DenseNodes::Iterator::operator!=(Iterator& other) const {

return offset != other.offset;

}

void PbfReader::DenseNodes::Iterator::operator++() {

offset++;

if (offset < nodes.ids.size()) {

node.id = nodes.ids[offset];

node.lon = nodes.lons[offset];

node.lat = nodes.lats[offset];

node.tagStart = nodes.tagStart[offset];

node.tagEnd = nodes.tagEnd[offset];

}

}

PbfReader::DenseNodes::Node& PbfReader::DenseNodes::Iterator::operator*() {

return node;

}

bool PbfReader::DenseNodes::empty() {

return ids.empty();

}

PbfReader::DenseNodes::Iterator PbfReader::DenseNodes::begin() {

auto it = Iterator {-1, Node{}, *this};

++it;

return it;

}

PbfReader::DenseNodes::Iterator PbfReader::DenseNodes::end() {

return Iterator {static_cast<int32_t>(ids.size()), Node{}, *this};

}

bool PbfReader::PrimitiveBlock::PrimitiveGroups::Iterator::operator!=(Iterator& other) const {

return offset != other.offset;

}

void PbfReader::PrimitiveBlock::PrimitiveGroups::Iterator::operator++() {

offset++;

if (offset < groups->size()) {

(*groups)[offset].ensureData();

}

}

PbfReader::PrimitiveGroup& PbfReader::PrimitiveBlock::PrimitiveGroups::Iterator::operator*() {

return (*groups)[offset];

}

PbfReader::PrimitiveBlock::PrimitiveGroups::Iterator PbfReader::PrimitiveBlock::PrimitiveGroups::begin() {

auto it = PrimitiveBlock::PrimitiveGroups::Iterator {-1, *groups };

++it;

return it;

}

PbfReader::PrimitiveBlock::PrimitiveGroups::Iterator PbfReader::PrimitiveBlock::PrimitiveGroups::end() {

return PrimitiveBlock::PrimitiveGroups::Iterator {static_cast<int32_t>(groups->size()), *groups };

}

PbfReader::PrimitiveGroupType PbfReader::PrimitiveGroup::type() const {

return internalType;

}

void PbfReader::Ways::Iterator::readWay(protozero::data_view data) {

protozero::pbf_message<Schema::Way> message{data};

way.id = 0;

way.keys.clear();

way.vals.clear();

way.refs.clear();

way.lats.clear();

way.lons.clear();

uint64_t ref = 0;

uint32_t lat = 0, lon = 0;

while (message.next()) {

switch (message.tag()) {

case Schema::Way::required_int64_id:

way.id = message.get_int64();

break;

case Schema::Way::repeated_uint32_keys: {

auto pi = message.get_packed_uint32();

for (auto i : pi) {

way.keys.push_back(i);

}

break;

}

case Schema::Way::repeated_uint32_vals: {

auto pi = message.get_packed_uint32();

for (auto i : pi) {

way.vals.push_back(i);

}

break;

}

case Schema::Way::repeated_sint64_refs: {

auto pi = message.get_packed_sint64();

for (auto i : pi) {

ref += i;

way.refs.push_back(ref);

}

break;

}

case Schema::Way::repeated_sint64_lats: {

auto pi = message.get_packed_sint64();

for (auto i : pi) {

lat += i;

way.lats.push_back(lat);

}

break;

}

case Schema::Way::repeated_sint64_lons: {

auto pi = message.get_packed_sint64();

for (auto i : pi) {

lon += i;

way.lons.push_back(lon);

}

break;

}

default:

// ignore data for unknown tags to allow for future extensions

//std::cout << "Way: unknown tag: " << std::to_string(static_cast<uint32_t>(message.tag())) << std::endl;

message.skip();

}

}

}

PbfReader::Ways& PbfReader::PrimitiveGroup::ways() const {

return internalWays;

}

bool PbfReader::Ways::Iterator::operator!=(Ways::Iterator& other) const {

return offset != other.offset;

}

void PbfReader::Ways::Iterator::operator++() {

if (message.next()) {

readWay(message.get_view());

offset++;

} else {

offset = -1;

}

}

PbfReader::Way& PbfReader::Ways::Iterator::operator*() {

return way;

}

bool PbfReader::Ways::empty() {

return pg->type() != PrimitiveGroupType::Way;

}

PbfReader::Ways::Iterator PbfReader::Ways::begin() {

if (pg->type() != PrimitiveGroupType::Way)

return Ways::Iterator{protozero::pbf_message<Schema::PrimitiveGroup>{nullptr, 0ul}, -1, way};

protozero::pbf_message<Schema::PrimitiveGroup> message{pg->getDataView()};

if (message.next()) {

protozero::pbf_message<Schema::PrimitiveGroup> message{pg->getDataView()};

auto it = Ways::Iterator{message, -1, way};

++it;

return it;

}

return Ways::Iterator{message, -1, way};

}

PbfReader::Ways::Iterator PbfReader::Ways::end() {

return Ways::Iterator{protozero::pbf_message<Schema::PrimitiveGroup>{nullptr, 0ul}, -1, way};

}

void PbfReader::Relations::Iterator::readRelation(protozero::data_view data) {

protozero::pbf_message<Schema::Relation> message{data};

relation.id = 0;

relation.keys.clear();

relation.vals.clear();

relation.memids.clear();

relation.roles_sid.clear();

relation.types.clear();

uint64_t memid = 0;

while (message.next()) {

switch (message.tag()) {

case Schema::Relation::required_int64_id:

relation.id = message.get_int64();

break;

case Schema::Relation::repeated_uint32_keys: {

auto pi = message.get_packed_uint32();

for (auto i : pi) {

relation.keys.push_back(i);

}

break;

}

case Schema::Relation::repeated_uint32_vals: {

auto pi = message.get_packed_uint32();

for (auto i : pi) {

relation.vals.push_back(i);

}

break;

}

case Schema::Relation::repeated_int32_roles_sid: {

auto pi = message.get_packed_int32();

for (auto i : pi) {

relation.roles_sid.push_back(i);

}

break;

}

case Schema::Relation::repeated_sint64_memids: {

auto pi = message.get_packed_sint64();

for (auto i : pi) {

memid += i;

relation.memids.push_back(memid);

}

break;

}

case Schema::Relation::repeated_MemberType_types: {

auto pi = message.get_packed_int32();

for (auto i : pi) {

relation.types.push_back(i);

}

break;

}

default:

// ignore data for unknown tags to allow for future extensions

//std::cout << "Way: unknown tag: " << std::to_string(static_cast<uint32_t>(message.tag())) << std::endl;

message.skip();

}

}

}

PbfReader::Relations& PbfReader::PrimitiveGroup::relations() const {

return internalRelations;

}

bool PbfReader::Relations::Iterator::operator!=(Relations::Iterator& other) const {

return offset != other.offset;

}

void PbfReader::Relations::Iterator::operator++() {

if (message.next()) {

readRelation(message.get_view());

offset++;

} else {

offset = -1;

}

}

PbfReader::Relation& PbfReader::Relations::Iterator::operator*() {

return relation;

}

bool PbfReader::Relations::empty() {

return pg->type() != PrimitiveGroupType::Relation;

}

PbfReader::Relations::Iterator PbfReader::Relations::begin() {

if (pg->type() != PrimitiveGroupType::Relation)

return Relations::Iterator{protozero::pbf_message<Schema::PrimitiveGroup>{nullptr, 0ul}, -1, relation};

protozero::pbf_message<Schema::PrimitiveGroup> message{pg->getDataView()};

if (message.next()) {

protozero::pbf_message<Schema::PrimitiveGroup> message{pg->getDataView()};

auto it = Relations::Iterator{message, -1, relation};

++it;

return it;

}

return Relations::Iterator{message, -1, relation};

}

PbfReader::Relations::Iterator PbfReader::Relations::end() {

return Relations::Iterator{protozero::pbf_message<Schema::PrimitiveGroup>{nullptr, 0ul}, -1, relation};

}

PbfReader::HeaderBlock PbfReader::PbfReader::readHeaderFromFile(std::istream& input) {

BlobHeader bh = readBlobHeader(input);

protozero::data_view blob = readBlob(bh.datasize, input);

HeaderBlock header = readHeaderBlock(blob);

return header;

}