File indexing completed on 2024-05-12 04:42:18

 /*
     SPDX-FileCopyrightText: 2023 Volker Krause <vkrause@kde.org>
     SPDX-License-Identifier: LGPL-2.0-or-later
 */
 
 #include "osmpbfwriter.h"
 #include "datatypes.h"
 
 #include "fileformat.pb.h"
 #include "osmformat.pb.h"
 
 #include <QIODevice>
 #include <QtEndian>
 
 #include <zlib.h>
 
 using namespace OSM;
 
 static constexpr const std::size_t BLOCK_SIZE_LIMIT = 16'777'216;
 
 OsmPbfWriter::OsmPbfWriter() = default;
 OsmPbfWriter::~OsmPbfWriter() = default;
 
 void OsmPbfWriter::writeToIODevice(const OSM::DataSet &dataSet, QIODevice *io)
 {
     m_io = io;
     writeNodes(dataSet);
     writeWays(dataSet);
     writeRelations(dataSet);
     if (m_block) {
         writeBlob();
     }
     m_io =  nullptr;
 }
 
 void OsmPbfWriter::writeNodes(const OSM::DataSet &dataSet)
 {
     if (dataSet.nodes.empty()) {
         return;
     }
 
     int64_t prevId = 0;
     int64_t prevLat = 900'000'000ll;
     int64_t prevLon = 1'800'000'000ll;
 
     OSMPBF::DenseNodes *denseBlock = nullptr;
     for(auto const &node: dataSet.nodes) {
         createBlockIfNeeded();
         if (!denseBlock) {
             denseBlock = m_block->add_primitivegroup()->mutable_dense();
         }
 
         denseBlock->add_id(node.id - prevId);
         prevId = node.id;
 
         denseBlock->add_lat((int64_t)node.coordinate.latitude - prevLat);
         prevLat = node.coordinate.latitude;
         denseBlock->add_lon((int64_t)node.coordinate.longitude - prevLon);
         prevLon = node.coordinate.longitude;
 
         for (const auto &tag : node.tags) {
             denseBlock->add_keys_vals(stringTableEntry(tag.key.name()));
             denseBlock->add_keys_vals(stringTableEntry(tag.value.constData()));
             m_blockSizeEstimate += 2 * sizeof(int32_t);
         }
 
         denseBlock->add_keys_vals(0);
         m_blockSizeEstimate += 3 * sizeof(int64_t) + sizeof(int32_t);
 
         if (blockSizeLimitReached()) {
             denseBlock = nullptr;
             writeBlob();
         }
     }
 }
 
 void OsmPbfWriter::writeWays(const OSM::DataSet &dataSet)
 {
     OSMPBF::PrimitiveGroup *group = nullptr;
     for(auto const &way : dataSet.ways) {
         createBlockIfNeeded();
         if (!group) {
             group = m_block->add_primitivegroup();
         }
 
         auto w = group->add_ways();
         w->set_id(way.id);
         m_blockSizeEstimate += sizeof(int64_t);
 
         int64_t prevId = 0;
         for (const auto &id : way.nodes) {
             w->add_refs(id - prevId);
             prevId = id;
             m_blockSizeEstimate += sizeof(int64_t);
         }
         for (const auto &tag : way.tags) {
             w->add_keys(stringTableEntry(tag.key.name()));
             w->add_vals(stringTableEntry(tag.value.constData()));
             m_blockSizeEstimate += 2 * sizeof(int32_t);
         }
 
         if (blockSizeLimitReached()) {
             group = nullptr;
             writeBlob();
         }
     }
 }
 
 static OSMPBF::Relation_MemberType pbfMemberType(OSM::Type t)
 {
     switch (t) {
         case OSM::Type::Null:
             Q_UNREACHABLE();
         case OSM::Type::Node:
             return OSMPBF::Relation_MemberType::Relation_MemberType_NODE;
         case OSM::Type::Way:
             return OSMPBF::Relation_MemberType::Relation_MemberType_WAY;
         case OSM::Type::Relation:
             return OSMPBF::Relation_MemberType::Relation_MemberType_RELATION;
     }
     return OSMPBF::Relation_MemberType::Relation_MemberType_NODE;
 }
 
 void OsmPbfWriter::writeRelations(const OSM::DataSet &dataSet)
 {
     OSMPBF::PrimitiveGroup *group = nullptr;
     for (const auto &rel :dataSet.relations) {
         createBlockIfNeeded();
         if (!group) {
             group = m_block->add_primitivegroup();
         }
 
         auto r = group->add_relations();
         r->set_id(rel.id);
         m_blockSizeEstimate += sizeof(int64_t);
 
         for (const auto &tag : rel.tags) {
             r->add_keys(stringTableEntry(tag.key.name()));
             r->add_vals(stringTableEntry(tag.value.constData()));
             m_blockSizeEstimate += 2 * sizeof(int32_t);
         }
 
         int64_t prevMemId = 0;
         for (const auto &mem : rel.members) {
             r->add_roles_sid(stringTableEntry(mem.role().name()));
             r->add_memids(mem.id - prevMemId);
             prevMemId = mem.id;
             r->add_types(pbfMemberType(mem.type()));
             m_blockSizeEstimate += 2* sizeof(int32_t) + sizeof(int64_t);
         }
 
         if (blockSizeLimitReached()) {
             group = nullptr;
             writeBlob();
         }
     }
 }
 
 int32_t OsmPbfWriter::stringTableEntry(const char *s)
 {
     assert(m_block);
     const auto it = m_stringTable.find(s);
     if (it == m_stringTable.end()) {
         auto st = m_block->mutable_stringtable();
         st->add_s(s);
         m_stringTable[s] = st->s_size() - 1;
         m_blockSizeEstimate += std::strlen(s) + 1 + sizeof(int32_t);
         return st->s_size() - 1;
     }
 
     return (*it).second;
 }
 
 void OsmPbfWriter::createBlockIfNeeded()
 {
     if (!m_block) {
         m_block = std::make_unique<OSMPBF::PrimitiveBlock>();
         m_blockSizeEstimate = 0;
         m_block->mutable_stringtable()->add_s(""); // dense node block tag separation marker
     }
 }
 
 bool OsmPbfWriter::blockSizeLimitReached() const
 {
     return m_blockSizeEstimate >BLOCK_SIZE_LIMIT;
 }
 
 void OsmPbfWriter::writeBlob()
 {
     OSMPBF::Blob blob;
     auto rawBlobData = m_block->SerializeAsString();
     auto zlibBlobData = blob.mutable_zlib_data();
     zlibBlobData->resize(32 * 1024ul * 1024ul);
 
     z_stream zStream;
     zStream.next_in = (uint8_t*)rawBlobData.data();
     zStream.avail_in = rawBlobData.size();
     zStream.next_out = (uint8_t*)zlibBlobData->data();
     zStream.avail_out = zlibBlobData->size();
     zStream.zalloc = nullptr;
     zStream.zfree = nullptr;
     zStream.opaque = nullptr;
     deflateInit(&zStream, Z_DEFAULT_COMPRESSION);
     while (true) {
         const auto ret = deflate(&zStream, Z_FINISH);
         if (ret == Z_STREAM_END) {
             break;
         }
         if (ret != Z_OK) {
             qWarning() << "zlib compression error!" << ret;
             break;
         }
         if (zStream.avail_out == 0) {
             // we could dynamically resize the output buffer here, but that
             // is already about twice the size of the data we want to compress
             // so we really shouldn't end up here
             qFatal("zlib output buffer underrun");
             break;
         }
     }
     zlibBlobData->resize(zlibBlobData->size() - zStream.avail_out);
 
     deflateEnd(&zStream);
     blob.set_raw_size((int32_t)rawBlobData.size());
 
     OSMPBF::BlobHeader header;
     header.set_type("OSMData");
     header.set_datasize((int32_t)blob.ByteSizeLong());
 
     auto blobHeaderSize = (int32_t)header.ByteSizeLong();
     blobHeaderSize = qToBigEndian(blobHeaderSize);
     m_io->write(reinterpret_cast<const char*>(&blobHeaderSize), sizeof(blobHeaderSize));
     m_io->write(header.SerializeAsString().c_str(), header.ByteSizeLong()); // TODO do this copy-free
     m_io->write(blob.SerializeAsString().c_str(), blob.ByteSizeLong()); // TODO do this copy-free
 
     m_block.reset();
     m_blockSizeEstimate = 0;
     m_stringTable.clear();
 }