pbf_craft/writers/

raw_writer.rs

use std::fs::File;
use std::io::{BufWriter, Write};
use std::mem;
use std::path::Path;

use anyhow;
use byteorder::{self, WriteBytesExt};
use flate2::write::ZlibEncoder;
use flate2::Compression;
use protobuf::Message;

use crate::codecs::block_builder::PrimitiveBuilder;
use crate::models::{Bound, Element};
use crate::proto::{fileformat, osmformat};

const MAX_BLOCK_ITEM_LENGTH: usize = 8000;

/// A writer for creating PBF files.
///
/// The `PbfWriter` struct provides functionality to write PBF data to an underlying writer.
/// It supports writing elements in either dense or non-dense format and can include optional
/// bounding box information.
///
/// Elements are buffered and flushed as blocks of 8000; the output blobs are compressed with
/// zlib. Elements marked `visible = false` that are seen before the header is emitted (the
/// first automatic flush or `finish()`) cause the header to declare the required
/// `HistoricalInformation` feature; callers streaming historical data whose invisible
/// elements may arrive after the first block must declare it up front with
/// [`PbfWriter::set_historical_data`]. `finish()` must be called to flush the last partial
/// block; dropping the writer flushes it best-effort (errors are only surfaced by
/// `finish()`).
///
/// Please note: the PBF format does not require sorted elements, but the conventional layout
/// (all nodes by id, then all ways by id, then all relations by id) is assumed by
/// `IndexedReader` and most other tools. `PbfWriter` stores elements in the order in which
/// `write` is called, so it is up to the caller to provide them in the desired order.
///
/// # Type Parameters
///
/// * `W` - A type that implements the `Write` trait, which is used to write the PBF data.
///
/// # Example
///
/// ```rust
/// use pbf_craft::models::{Element, Node};
/// use pbf_craft::writers::PbfWriter;
///
/// let mut writer = PbfWriter::from_path(std::env::temp_dir().join("output.pbf"), true).unwrap();
/// writer.write(Element::Node(Node::default())).unwrap();
/// writer.finish().unwrap();
/// ```
pub struct PbfWriter<W: Write> {
    writer: W,
    use_dense: bool,
    bbox: Option<Bound>,
    cache: Vec<Element>,
    has_written_header: bool,
    has_invisible_elements: bool,
    declared_historical_data: bool,
}

impl PbfWriter<BufWriter<File>> {
    /// Creates a new `PbfWriter` from a file path.
    ///
    /// # Parameters
    ///
    /// * `path` - The path to the file to write the PBF data to.
    /// * `use_dense` - A boolean value indicating whether to use dense format for writing nodes.
    ///
    pub fn from_path<P: AsRef<Path>>(path: P, use_dense: bool) -> anyhow::Result<Self> {
        let f = File::create(path)?;
        let writer = BufWriter::new(f);
        Ok(Self::new(writer, use_dense))
    }
}

impl<W: Write> PbfWriter<W> {
    /// Creates a new `PbfWriter` from an existing writer.
    ///
    /// # Parameters
    ///
    /// * `writer` - The writer to use for writing the PBF data. It should implement the `Write`
    ///   trait, which is used to write the PBF data.
    /// * `use_dense` - A boolean value indicating whether to use dense format for writing nodes.
    ///
    pub fn new(writer: W, use_dense: bool) -> PbfWriter<W> {
        Self {
            writer,
            use_dense,
            bbox: None,
            cache: Vec::new(),
            has_written_header: false,
            has_invisible_elements: false,
            declared_historical_data: false,
        }
    }

    fn build_raw_blob(&mut self, raw: Vec<u8>) -> anyhow::Result<fileformat::Blob> {
        let raw_size = raw.len();
        let mut zlib_encoder = ZlibEncoder::new(Vec::new(), Compression::default());
        zlib_encoder.write_all(raw.as_slice())?;
        let compressed = zlib_encoder.finish()?;

        let mut blob = fileformat::Blob::new();
        blob.set_zlib_data(compressed);
        blob.set_raw_size(raw_size as i32);
        Ok(blob)
    }

    /// Sets the bounding box for the PBF file.
    ///
    /// If you want to include a bounding box in the PBF file, you set it before writing any elements.
    ///
    pub fn set_bbox(&mut self, bbox: Bound) {
        self.bbox = Some(bbox);
    }

    /// Declares that the data being written is historical, i.e. contains elements with
    /// `visible = false`, so the header block declares the required `HistoricalInformation`
    /// feature.
    ///
    /// The writer also auto-detects `visible = false` elements, but only those seen before
    /// the header is emitted — the header is written together with the first flushed block
    /// (8000 elements) or at `finish()` and can never be amended afterwards. Callers
    /// streaming historical data that cannot guarantee an invisible element inside the first
    /// block must call this before writing, like [`PbfWriter::set_bbox`].
    ///
    /// Returns an error if the header has already been written (i.e. the first block was
    /// already flushed) and the declaration can no longer take effect.
    pub fn set_historical_data(&mut self, historical: bool) -> anyhow::Result<()> {
        if self.has_written_header {
            bail!(
                "cannot declare historical data after the header has already been written \
                 (the first block was flushed)"
            );
        }
        self.declared_historical_data = historical;
        Ok(())
    }

    fn write_header(&mut self) -> anyhow::Result<()> {
        let mut header_block = osmformat::HeaderBlock::new();
        header_block
            .required_features
            .push("OsmSchema-V0.6".to_string());
        if self.use_dense {
            header_block
                .required_features
                .push("DenseNodes".to_string());
        }
        // Per the PBF spec, a writer that emits `visible = false` (historical data) MUST
        // declare the HistoricalInformation feature. The flag combines every element seen so
        // far (auto-detected) with any explicit up-front declaration made via
        // `set_historical_data`; once the header is flushed it can never be amended, so
        // callers whose invisible elements may arrive after the first block must declare the
        // data historical before writing.
        if self.has_invisible_elements || self.declared_historical_data {
            header_block
                .required_features
                .push("HistoricalInformation".to_string());
        }

        if let Some(bbox) = &self.bbox {
            let mut header_bbox = osmformat::HeaderBBox::new();
            header_bbox.set_left(bbox.left);
            header_bbox.set_right(bbox.right);
            header_bbox.set_top(bbox.top);
            header_bbox.set_bottom(bbox.bottom);
            header_block.set_bbox(header_bbox);
            header_block.set_source(bbox.origin.clone());
        }

        let blob = self.build_raw_blob(header_block.write_to_bytes()?)?;
        self.write_blob(blob, "OSMHeader")?;
        self.has_written_header = true;
        Ok(())
    }

    /// Writes an element.
    ///
    /// Please note: the PBF format does not require sorted elements, but `IndexedReader` and
    /// most other tools assume the conventional ordering (all nodes by id, then all ways by
    /// id, then all relations by id). The writer stores elements in the order they are
    /// written — the caller is responsible for providing them in the desired order.
    ///
    pub fn write(&mut self, element: Element) -> anyhow::Result<()> {
        // Track whether any element is marked invisible so the header can declare the
        // required HistoricalInformation feature (see `write_header`).
        match &element {
            Element::Node(node) => self.has_invisible_elements |= !node.visible,
            Element::Way(way) => self.has_invisible_elements |= !way.visible,
            Element::Relation(relation) => self.has_invisible_elements |= !relation.visible,
        }
        self.cache.push(element);
        if self.cache.len() >= MAX_BLOCK_ITEM_LENGTH {
            self.write_to_block()?;
        }
        Ok(())
    }

    fn write_to_block(&mut self) -> anyhow::Result<()> {
        if !self.has_written_header {
            self.write_header()?;
        }
        if self.cache.is_empty() {
            // Nothing buffered: emit no empty data block (the header alone already forms a
            // valid file for a writer with no elements).
            return Ok(());
        }
        let block_builder = PrimitiveBuilder::new();
        let cache = mem::take(&mut self.cache);
        let block = block_builder.build(cache, self.use_dense);

        let blob = self.build_raw_blob(block.write_to_bytes()?)?;
        self.write_blob(blob, "OSMData")?;
        Ok(())
    }

    fn write_blob(&mut self, blob: fileformat::Blob, blob_type: &str) -> anyhow::Result<()> {
        let blob_bytes = blob.write_to_bytes()?;

        let mut header = fileformat::BlobHeader::new();
        header.set_datasize(blob_bytes.len() as i32);
        header.set_field_type(blob_type.to_owned());
        let header_bytes = header.write_to_bytes()?;

        self.writer
            .write_u32::<byteorder::BigEndian>(header_bytes.len() as u32)?;
        self.writer.write_all(header_bytes.as_slice())?;
        self.writer.write_all(blob_bytes.as_slice())?;

        Ok(())
    }

    /// Finishes writing the PBF file.
    ///
    /// This method should be called after writing all elements to the PBF file. It writes the
    /// header (even for an empty file) and flushes any buffered elements.
    ///
    pub fn finish(&mut self) -> anyhow::Result<()> {
        self.write_to_block()?;
        self.writer.flush()?;
        Ok(())
    }
}

impl<W: Write> Drop for PbfWriter<W> {
    fn drop(&mut self) {
        // Best-effort flush of buffered elements so a forgotten `finish()` does not silently
        // produce an empty file. Errors cannot be returned from `drop`; call `finish()`
        // explicitly to surface them.
        if !self.cache.is_empty() {
            let _ = self.write_to_block();
        }
        let _ = self.writer.flush();
    }
}