pbfhogg 0.5.0

//! Speed up searches by using an index

use super::blob::{BlobReader, BlobReaderSource, BlobType, ByteOffset};
use super::block::PrimitiveBlock;
use super::elements::{Element, Way};
use crate::error::Result;
use std::collections::BTreeSet;
use std::fs::File;
use std::ops::RangeInclusive;
use std::path::Path;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum SimpleBlobType {
    Header,
    Primitive,
    Unknown,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum ElementsAvailable {
    Yes,
    No,
    Unknown,
}

/// Returns true if the given set contains at least one value that is inside the given range.
// Takes RangeInclusive by value intentionally - it's 16 bytes (two i64s), copy is
// cheaper than indirection. BTreeSet::range accepts both owned and borrowed.
fn range_included(range: RangeInclusive<i64>, node_ids: &BTreeSet<i64>) -> bool {
    node_ids.range(range).next().is_some()
}

#[derive(Clone, Debug, Eq, PartialEq)]
enum RangeIncluded {
    Yes(RangeInclusive<i64>),
    No,
    Unknown,
}

/// Stores the minimum and maximum id of every element type.
#[derive(Debug)]
pub struct IdRanges {
    node_ids: Option<RangeInclusive<i64>>,
    way_ids: Option<RangeInclusive<i64>>,
}

/// A part of the index that stores information about a specific blob.
///
/// `id_ranges: None` means the blob has not been fully decoded yet -
/// it is distinct from `Some(IdRanges { node_ids: None, .. })` which
/// means the blob was decoded and contained no nodes. The
/// `ElementsAvailable::{Yes, No, Unknown}` return of the helpers
/// below encodes this three-valued state explicitly; callers must
/// not collapse `None`/`Unknown` into "no elements", or relation-only
/// blobs and not-yet-decoded blobs become indistinguishable.
#[derive(Debug)]
struct BlobInfo {
    offset: ByteOffset,
    blob_type: SimpleBlobType,
    id_ranges: Option<IdRanges>,
}

impl BlobInfo {
    /// Is there at least one node in this blob?
    fn nodes_available(&self) -> ElementsAvailable {
        match self.id_ranges {
            Some(IdRanges {
                node_ids: Some(_), ..
            }) => ElementsAvailable::Yes,
            Some(IdRanges { node_ids: None, .. }) => ElementsAvailable::No,
            None => ElementsAvailable::Unknown,
        }
    }

    /// Is there at least one way in this blob?
    fn ways_available(&self) -> ElementsAvailable {
        match self.id_ranges {
            Some(IdRanges {
                way_ids: Some(_), ..
            }) => ElementsAvailable::Yes,
            Some(IdRanges { way_ids: None, .. }) => ElementsAvailable::No,
            None => ElementsAvailable::Unknown,
        }
    }

    /// Compute if the range of node IDs of this blob (min and max ID value) is included in the
    /// given set of IDs with at least one ID inside of this range.
    fn node_range_included(&self, node_ids: &BTreeSet<i64>) -> RangeIncluded {
        match self.id_ranges.as_ref() {
            None => RangeIncluded::Unknown,
            Some(IdRanges { node_ids: None, .. }) => RangeIncluded::No,
            Some(IdRanges {
                node_ids: Some(range),
                ..
            }) => {
                if range_included(range.clone(), node_ids) {
                    RangeIncluded::Yes(range.clone())
                } else {
                    RangeIncluded::No
                }
            }
        }
    }
}

/// Allows filtering elements and iterating over their dependencies.
/// It chooses an efficient method for navigating the PBF structure to achieve this in reasonable
/// time and with reasonable memory.
// wontfix(type-generic-bounds): bounds on struct match osmpbf API and document intent
pub struct IndexedReader<R: BlobReaderSource + Send> {
    reader: BlobReader<R>,
    index: Vec<BlobInfo>,
}

impl<R: BlobReaderSource + Send> IndexedReader<R> {
    /// Creates a new `IndexedReader`.
    ///
    /// # Example
    /// ```
    /// use pbfhogg::*;
    ///
    /// # fn foo() -> Result<()> {
    /// let f = std::fs::File::open("tests/test.osm.pbf")?;
    /// let buf_reader = std::io::BufReader::new(f);
    ///
    /// let reader = IndexedReader::new(buf_reader)?;
    ///
    /// # Ok(())
    /// # }
    /// # foo().unwrap();
    /// ```
    pub fn new(reader: R) -> Result<Self> {
        let reader = BlobReader::new_seekable(reader)?;
        Ok(Self {
            reader,
            index: vec![],
        })
    }

    /// Initializes the index of the PBF structure without decompressing the blobs.
    /// You do not need to call this method explicitly as the other methods already take care of
    /// it.
    pub fn create_index(&mut self) -> Result<()> {
        if !self.index.is_empty() {
            // Index is already present -> Do nothing
            return Ok(());
        }

        // Seek to the beginning of the reader.
        self.reader.seek(ByteOffset(0))?;

        while let Some(result) = self.reader.next_header_skip_blob() {
            let (header, offset) = result?;
            // The original code used `offset.unwrap()` here because IndexedReader
            // always constructs its BlobReader via `new_seekable`, which initializes
            // offset tracking. A seekable BlobReader always returns
            // `Some(ByteOffset)` from `next_header_skip_blob`. However, relying on
            // that invariant via unwrap is fragile -- if the BlobReader internals
            // ever change, this would panic at runtime. Instead, we propagate a
            // descriptive IO error so callers get a clear message rather than a panic.
            let offset = offset.ok_or_else(|| {
                crate::error::new_error(crate::error::ErrorKind::Io(std::io::Error::other(
                    // Starts with type name intentionally - it IS a proper noun here.
                    "IndexedReader requires a seekable BlobReader with offset tracking",
                )))
            })?;
            let blob_type = match header.blob_type() {
                BlobType::OsmHeader => SimpleBlobType::Header,
                BlobType::OsmData => SimpleBlobType::Primitive,
                BlobType::Unknown(_) => SimpleBlobType::Unknown,
            };

            self.index.push(BlobInfo {
                offset,
                blob_type,
                id_ranges: None,
            });
        }

        Ok(())
    }

    /// Check element IDs of this block. Record min and max for every node, way and relation.
    ///
    /// This is an ID-only consumer - it iterates all elements but only calls `.id()`.
    /// A lightweight "scan mode" that skips protobuf parsing of tags, coordinates, refs,
    /// and metadata was considered but deemed not worth the complexity: this runs once per
    /// `IndexedReader` session (during `create_index`), not in a hot loop. The full parse
    /// is already needed for the subsequent filtered iteration passes. See the doc comment
    /// on `PrimitiveBlock` in block.rs for the full rationale.
    #[allow(clippy::cognitive_complexity)]
    fn update_element_id_ranges(info: &mut BlobInfo, block: &PrimitiveBlock) {
        if info.id_ranges.is_some() {
            // Ranges are already present -> Do nothing
            return;
        }

        let mut min_node_id: Option<i64> = None;
        let mut max_node_id: Option<i64> = None;
        let mut min_way_id: Option<i64> = None;
        let mut max_way_id: Option<i64> = None;
        // Check each primitive group
        for group in block.groups() {
            let check_min_max = |id, min_id: &mut Option<i64>, max_id: &mut Option<i64>| {
                *min_id = Some(min_id.map_or(id, |x| x.min(id)));
                *max_id = Some(max_id.map_or(id, |x| x.max(id)));
            };

            for node in group.nodes() {
                check_min_max(node.id(), &mut min_node_id, &mut max_node_id);
            }
            for node in group.dense_nodes() {
                check_min_max(node.id(), &mut min_node_id, &mut max_node_id);
            }
            for way in group.ways() {
                check_min_max(way.id(), &mut min_way_id, &mut max_way_id);
            }
        }

        let to_range = |min_id, max_id| -> Option<RangeInclusive<i64>> {
            if let (Some(min), Some(max)) = (min_id, max_id) {
                Some(RangeInclusive::new(min, max))
            } else {
                None
            }
        };

        info.id_ranges = Some(IdRanges {
            node_ids: to_range(min_node_id, max_node_id),
            way_ids: to_range(min_way_id, max_way_id),
        });
    }

    /// Filter ways using a closure and return matching ways and their dependent nodes
    /// ([`Node`](crate::elements::Node)s and [`DenseNode`](crate::dense::DenseNode)s)
    /// in another closure.
    /// This method also creates a lightweight in-memory index that speeds up future invocations of
    /// this or any other method of `IndexedReader`.
    ///
    /// # Memory
    ///
    /// Collects all node IDs referenced by matching ways into a `BTreeSet<i64>`.
    /// For planet-scale files with broad way filters, this set can grow to
    /// billions of entries (~48 bytes each in `BTreeSet`), potentially exceeding
    /// available memory. For planet-scale extractions, prefer the `extract`
    /// command which uses [`IdSet`](crate::idset::IdSet)
    /// (1 bit per ID) and streaming passes.
    ///
    /// # Example
    /// ```
    /// use pbfhogg::*;
    ///
    /// # fn foo() -> Result<()> {
    /// let mut reader = IndexedReader::from_path("tests/test.osm.pbf")?;
    /// let mut ways = 0;
    /// let mut nodes = 0;
    ///
    /// // Filter all ways that are buildings and count their nodes.
    /// reader.read_ways_and_deps(
    ///     |way| {
    ///         // Filter ways. Return true if tags contain "building": "yes".
    ///         way.tags().any(|key_value| key_value == ("building", "yes"))
    ///     },
    ///     |element| {
    ///         // Increment counter
    ///         match element {
    ///             Element::Way(way) => ways += 1,
    ///             Element::Node(node) => nodes += 1,
    ///             Element::DenseNode(dense_node) => nodes += 1,
    ///             Element::Relation(_) | _ => (), // should not occur
    ///         }
    ///     },
    /// )?;
    ///
    /// println!("ways:  {ways}\nnodes: {nodes}");
    ///
    /// # assert_eq!(ways, 1);
    /// # assert_eq!(nodes, 3);
    /// # Ok(())
    /// # }
    /// # foo().unwrap();
    /// ```
    #[allow(clippy::too_many_lines)]
    pub fn read_ways_and_deps<F, E>(&mut self, mut filter: F, mut element_callback: E) -> Result<()>
    where
        F: for<'a> FnMut(&Way<'a>) -> bool,
        E: for<'a> FnMut(&Element<'a>),
    {
        self.create_index()?;

        let mut node_ids: BTreeSet<i64> = BTreeSet::new();

        // First pass:
        //   * Filter ways and store their dependencies as node IDs
        for info in &mut self.index {
            if info.blob_type == SimpleBlobType::Primitive
                && info.ways_available() != ElementsAvailable::No
            {
                let block = self
                    .reader
                    .blob_from_offset(info.offset)?
                    .to_primitiveblock()?;
                Self::update_element_id_ranges(info, &block);

                for group in block.groups() {
                    // filter ways and record node IDs
                    for way in group.ways() {
                        if filter(&way) {
                            let refs = way.refs();

                            node_ids.extend(refs);

                            // Return way
                            element_callback(&Element::Way(way));
                        }
                    }
                }
            }
        }

        // Second pass:
        //   * Iterate only over blobs that may include the node IDs we're searching for
        for info in &mut self.index {
            if let RangeIncluded::Yes(_) = info.node_range_included(&node_ids) {
                let block = self
                    .reader
                    .blob_from_offset(info.offset)?
                    .to_primitiveblock()?;
                for group in block.groups() {
                    for node in group.nodes() {
                        if node_ids.contains(&node.id()) {
                            element_callback(&Element::Node(node));
                        }
                    }
                    for node in group.dense_nodes() {
                        if node_ids.contains(&node.id()) {
                            element_callback(&Element::DenseNode(node));
                        }
                    }
                }
            }
        }

        Ok(())
    }

    /// Decodes the PBF structure sequentially and calls the given closure on each node.
    /// This method also creates a lightweight in-memory index that speeds up future invocations of
    /// this or any other method of `IndexedReader`.
    ///
    /// # Errors
    /// Returns the first Error encountered while parsing the PBF structure.
    ///
    /// # Example
    /// ```
    /// use pbfhogg::*;
    ///
    /// # fn foo() -> Result<()> {
    /// let mut reader = IndexedReader::from_path("tests/test.osm.pbf")?;
    /// let mut nodes = 0;
    ///
    /// reader.for_each_node(
    ///     |element| {
    ///         match element {
    ///             Element::Node(node) => nodes += 1,
    ///             Element::DenseNode(dense_node) => nodes += 1,
    ///             _ => {}
    ///         }
    ///     },
    /// )?;
    ///
    /// println!("nodes: {nodes}");
    ///
    /// # assert_eq!(nodes, 3);
    /// # Ok(())
    /// # }
    /// # foo().unwrap();
    /// ```
    pub fn for_each_node<F>(&mut self, mut f: F) -> Result<()>
    where
        F: for<'a> FnMut(Element<'a>),
    {
        self.create_index()?;

        for info in &mut self.index {
            // Skip header blobs and blobs where there are certainly no nodes available.
            if info.blob_type == SimpleBlobType::Primitive
                && info.nodes_available() != ElementsAvailable::No
            {
                let block = self
                    .reader
                    .blob_from_offset(info.offset)?
                    .to_primitiveblock()?;
                Self::update_element_id_ranges(info, &block);

                for group in block.groups() {
                    for node in group.nodes() {
                        f(Element::Node(node));
                    }
                    for dense_node in group.dense_nodes() {
                        f(Element::DenseNode(dense_node));
                    }
                }
            }
        }

        Ok(())
    }
}

impl IndexedReader<File> {
    /// Creates a new `IndexedReader` from a given path.
    ///
    /// # Example
    /// ```
    /// use pbfhogg::*;
    ///
    /// # fn foo() -> Result<()> {
    /// let reader = IndexedReader::from_path("tests/test.osm.pbf")?;
    ///
    /// # Ok(())
    /// # }
    /// # foo().unwrap();
    /// ```
    pub fn from_path<P: AsRef<Path>>(path: P) -> Result<Self> {
        let f = File::open(path)?;
        // BufReader is intentionally NOT used here. IndexedReader performs random
        // seeks to jump between blobs (e.g. first pass reads ways, second pass seeks
        // back to read dependent nodes). A large read-ahead buffer would be wasted
        // on every seek - the buffered data would be discarded immediately.
        //
        // This was investigated and confirmed: adding BufReader here showed no
        // improvement and slightly increased memory overhead. See commit a38c258.
        //
        // For sequential reads, see BlobReader::from_path in blob.rs which uses a
        // 256KB BufReader to amortize syscall overhead.
        Self::new(f)
    }
}

// Tests use `unwrap()` throughout because panicking is the correct failure mode
// for unit tests -- it immediately fails the test with a clear backtrace pointing
// to the exact call site. Propagating Results via `-> Result<()>` in tests would
// lose the backtrace and produce less actionable error messages. The crate-wide
// `unwrap_used = "deny"` lint is designed for production code where panics are
// unacceptable; test code is exempt via this module-level allow.
#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;

    #[test]
    fn test_range_included_set() {
        let mut set = BTreeSet::<i64>::new();
        set.extend(&[1, 2, 6]);

        assert!(!range_included(RangeInclusive::new(0, 0), &set));
        assert!(range_included(RangeInclusive::new(1, 1), &set));
        assert!(range_included(RangeInclusive::new(2, 2), &set));
        assert!(!range_included(RangeInclusive::new(3, 3), &set));
        assert!(!range_included(RangeInclusive::new(3, 5), &set));
        assert!(range_included(RangeInclusive::new(3, 6), &set));
        assert!(range_included(RangeInclusive::new(6, 6), &set));
        assert!(!range_included(RangeInclusive::new(7, 7), &set));
        assert!(range_included(RangeInclusive::new(0, 1), &set));
        assert!(range_included(RangeInclusive::new(6, 7), &set));
        assert!(range_included(RangeInclusive::new(2, 3), &set));
        assert!(range_included(RangeInclusive::new(5, 6), &set));
        assert!(range_included(RangeInclusive::new(5, 8), &set));
        assert!(range_included(RangeInclusive::new(0, 8), &set));
        assert!(range_included(RangeInclusive::new(0, 4), &set));
    }
}