diff --git a/Cargo.toml b/Cargo.toml index cbbae5a0..aa57636f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "osmpbfreader" -version = "0.19.1" +version = "0.20.0" authors = ["Guillaume Pinot "] description = "Read OpenStreetMap PBF files in rust." documentation = "https://docs.rs/osmpbfreader" diff --git a/examples/count.rs b/examples/count.rs index b3a025eb..b696bb3a 100644 --- a/examples/count.rs +++ b/examples/count.rs @@ -12,7 +12,7 @@ extern crate osmpbfreader; fn count bool>(filter: F, filename: &std::ffi::OsStr) { let r = std::fs::File::open(std::path::Path::new(filename)).unwrap(); - let mut pbf = osmpbfreader::OsmPbfReader::new(r); + let mut pbf = osmpbfreader::ParReader::new(r); let mut nb_nodes = 0; let mut sum_lon = 0.; let mut sum_lat = 0.; @@ -20,7 +20,7 @@ fn count bool>(filter: F, filename: &std::ffi::OsS let mut nb_way_nodes = 0; let mut nb_rels = 0; let mut nb_rel_refs = 0; - for obj in pbf.par_iter().map(Result::unwrap) { + for obj in pbf.iter().map(Result::unwrap) { if !filter(obj.tags()) { continue; } diff --git a/examples/count_with_deps.rs b/examples/count_with_deps.rs index 93bf26d8..7b47a9db 100644 --- a/examples/count_with_deps.rs +++ b/examples/count_with_deps.rs @@ -12,7 +12,7 @@ extern crate osmpbfreader; fn count bool>(filter: F, filename: &std::ffi::OsStr) { let r = std::fs::File::open(std::path::Path::new(filename)).unwrap(); - let mut pbf = osmpbfreader::OsmPbfReader::new(r); + let mut pbf = osmpbfreader::Reader::new(r); let objs = pbf.get_objs_and_deps(|obj| filter(obj.tags())).unwrap(); let mut nb_nodes = 0; let mut sum_lon = 0.; diff --git a/examples/relation.rs b/examples/relation.rs index ff772cc6..1804b71e 100644 --- a/examples/relation.rs +++ b/examples/relation.rs @@ -15,7 +15,7 @@ fn main() { let filename = std::env::args_os().nth(1).unwrap(); let path = std::path::Path::new(&filename); let r = std::fs::File::open(path).unwrap(); - let mut pbf = osmpbfreader::OsmPbfReader::new(r); + let mut pbf = osmpbfreader::Reader::new(r); let objects = pbf.get_objs_and_deps(wanted).unwrap(); println!( "The relation Paris is composed of {:?} items", diff --git a/examples/tutorial.rs b/examples/tutorial.rs index ab409492..44132eb9 100644 --- a/examples/tutorial.rs +++ b/examples/tutorial.rs @@ -11,7 +11,7 @@ fn main() { let filename = std::env::args_os().nth(1).unwrap(); let path = std::path::Path::new(&filename); let r = std::fs::File::open(path).unwrap(); - let mut pbf = osmpbfreader::OsmPbfReader::new(r); + let mut pbf = osmpbfreader::Reader::new(r); let mut nb = 0; for _obj in pbf.iter().map(Result::unwrap) { nb += 1; diff --git a/src/lib.rs b/src/lib.rs index 403c5cb5..8da4b2e7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -18,7 +18,7 @@ //! availlable. //! //! ```rust -//! let mut pbf = osmpbfreader::OsmPbfReader::new(std::io::Cursor::new([])); +//! let mut pbf = osmpbfreader::Reader::new(std::io::Cursor::new([])); //! let objs = pbf.get_objs_and_deps(|obj| { //! obj.is_way() && obj.tags().contains_key("highway") //! }) @@ -35,7 +35,7 @@ //! //! ```rust //! use std::process::exit; -//! let mut pbf = osmpbfreader::OsmPbfReader::new(std::io::empty()); +//! let mut pbf = osmpbfreader::Reader::new(std::io::empty()); //! for obj in pbf.iter() { //! // error handling: //! let obj = obj.unwrap_or_else(|e| {println!("{:?}", e); exit(1)}); @@ -49,8 +49,8 @@ //! //! ```rust //! use std::process::exit; -//! let mut pbf = osmpbfreader::OsmPbfReader::new(std::io::empty()); -//! for obj in pbf.par_iter() { +//! let mut pbf = osmpbfreader::ParReader::new(std::io::empty()); +//! for obj in pbf.iter() { //! // error handling: //! let obj = obj.unwrap_or_else(|e| {println!("{:?}", e); exit(1)}); //! @@ -67,7 +67,7 @@ //! //! ```rust //! use osmpbfreader::{primitive_block_from_blob, groups}; -//! let mut pbf = osmpbfreader::OsmPbfReader::new(std::io::empty()); +//! let mut pbf = osmpbfreader::Reader::new(std::io::empty()); //! for block in pbf.blobs().map(|b| primitive_block_from_blob(&b.unwrap())) { //! let block = block.unwrap(); //! for group in block.primitivegroup.iter() { @@ -98,10 +98,12 @@ pub use error::Error; pub use error::Result; pub use objects::*; -pub use reader::{primitive_block_from_blob, OsmPbfReader, StoreObjs}; +pub use par_reader::ParReader; +pub use reader::{primitive_block_from_blob, Reader}; pub mod blobs; pub mod objects; +pub mod par_reader; pub mod reader; #[allow(missing_docs)] @@ -110,6 +112,7 @@ pub mod blocks; pub mod error; #[allow(missing_docs)] pub mod groups; +mod store_objs; mod pbf { include!(concat!(env!("OUT_DIR"), "/mod.rs")); diff --git a/src/par_reader.rs b/src/par_reader.rs new file mode 100644 index 00000000..df2f1f3f --- /dev/null +++ b/src/par_reader.rs @@ -0,0 +1,313 @@ +// Copyright (c) 2014-2017 Guillaume Pinot +// +// This work is free. You can redistribute it and/or modify it under +// the terms of the Do What The Fuck You Want To Public License, +// Version 2, as published by Sam Hocevar. See the COPYING file for +// more details. + +//! Tools for reading a pbf file. + +use crate::blobs; +use crate::error::{Error, Result}; +use crate::fileformat::{Blob, BlobHeader}; +use crate::objects::{OsmId, OsmObj}; +use crate::osmformat::PrimitiveBlock; +use crate::store_objs::StoreObjs; +use par_map::{self, ParMap}; +use protobuf::Message; +use pub_iterator_type::pub_iterator_type; +use std::collections::btree_map::BTreeMap; +use std::collections::BTreeSet; +use std::convert::From; +use std::io::{self, Read}; +use std::iter; + +/// The object to manage a pbf file. +pub struct ParReader { + buf: Vec, + r: R, + finished: bool, +} + +impl ParReader { + /// Creates an OsmPbfReader from a Read object. + pub fn new(r: R) -> ParReader { + ParReader { + buf: vec![], + r, + finished: false, + } + } + + /// Returns a parallel iterator on the OsmObj of the pbf file. + /// + /// Several threads decode in parallel the file. The memory and + /// CPU usage are guaranteed to be bounded even if the caller stop + /// consuming items. + /// + /// # Example + /// + /// ``` + /// let mut pbf = osmpbfreader::ParReader::new(std::io::empty()); + /// for obj in pbf.iter().map(Result::unwrap) { + /// println!("{:?}", obj); + /// } + /// ``` + pub fn iter(&mut self) -> ParIter<'_, R> { + ParIter(self.blobs().par_flat_map(blobs::result_blob_into_iter)) + } + + /// Returns a parallel iterator on the Node of the pbf file. + /// + /// Several threads decode in parallel the file. The memory and + /// CPU usage are guaranteed to be bounded even if the caller stop + /// consuming items. + pub fn iter_nodes(&mut self) -> NodeParIter<'_, R> { + NodeParIter(self.blobs().par_flat_map(blobs::result_blob_into_node_iter)) + } + + /// Returns a parallel iterator on the Way of the pbf file. + /// + /// Several threads decode in parallel the file. The memory and + /// CPU usage are guaranteed to be bounded even if the caller stop + /// consuming items. + pub fn iter_ways(&mut self) -> WayParIter<'_, R> { + WayParIter(self.blobs().par_flat_map(blobs::result_blob_into_way_iter)) + } + + /// Returns a parallel iterator on the Relation of the pbf file. + /// + /// Several threads decode in parallel the file. The memory and + /// CPU usage are guaranteed to be bounded even if the caller stop + /// consuming items. + pub fn iter_relations(&mut self) -> RelationParIter<'_, R> { + RelationParIter( + self.blobs() + .par_flat_map(blobs::result_blob_into_relation_iter), + ) + } + + /// Rewinds the pbf file to the begining. + /// + /// Useful if you want to read several consecutive times the same + /// file. + /// + /// # Example + /// + /// ``` + /// let mut cursor = std::io::Cursor::new([0, 0, 0]); + /// cursor.set_position(2); + /// let mut pbf = osmpbfreader::ParReader::new(cursor); + /// pbf.rewind().unwrap(); + /// assert_eq!(pbf.into_inner().position(), 0); + /// ``` + pub fn rewind(&mut self) -> Result<()> + where + R: io::Seek, + { + self.r.seek(io::SeekFrom::Start(0))?; + self.finished = false; + Ok(()) + } + + /// Same as `get_objs_and_deps` but generic. + pub fn get_objs_and_deps_store(&mut self, mut pred: F, objects: &mut T) -> Result<()> + where + R: io::Seek, + F: FnMut(&OsmObj) -> bool, + T: StoreObjs, + { + let mut finished = false; + let mut deps = BTreeSet::new(); + let mut first_pass = true; + while !finished { + self.rewind()?; + finished = true; + for obj in self.iter() { + let obj = obj?; + if (!first_pass || !pred(&obj)) && !deps.contains(&obj.id()) { + continue; + } + finished = match obj { + OsmObj::Relation(ref rel) => rel + .refs + .iter() + .filter(|r| !objects.contains_key(&r.member)) + .fold(finished, |accu, r| !deps.insert(r.member) && accu), + OsmObj::Way(ref way) => way + .nodes + .iter() + .filter(|n| !objects.contains_key(&(**n).into())) + .fold(finished, |accu, n| !deps.insert((*n).into()) && accu), + OsmObj::Node(_) => finished, + }; + deps.remove(&obj.id()); + objects.insert(obj.id(), obj); + } + first_pass = false; + } + Ok(()) + } + + /// This function give you the ability to find all the objects + /// validating a predicate and all their dependencies. The file + /// will be decoded in parallel. + /// + /// # Example + /// + /// If you want to extract all the administrative boundaries + /// and all their dependencies you can do something like that: + /// + /// ``` + /// fn is_admin(obj: &osmpbfreader::OsmObj) -> bool { + /// // get relations with tags[boundary] == administrative + /// obj.is_relation() && obj.tags().contains("boundary", "administrative") + /// } + /// + /// let mut pbf = osmpbfreader::ParReader::new(std::io::Cursor::new([])); + /// let objs = pbf.get_objs_and_deps(is_admin).unwrap(); + /// for (id, obj) in &objs { + /// println!("{:?}: {:?}", id, obj); + /// } + /// ``` + pub fn get_objs_and_deps(&mut self, pred: F) -> Result> + where + R: io::Seek, + F: FnMut(&OsmObj) -> bool, + { + let mut objects = BTreeMap::new(); + match self.get_objs_and_deps_store(pred, &mut objects) { + Ok(_) => Ok(objects), + Err(e) => Err(e), + } + } + /// Extract the Read object. + /// + /// Consumes the object. + pub fn into_inner(self) -> R { + self.r + } + /// Returns an iterator on the blobs of the pbf file. + pub fn blobs(&mut self) -> Blobs<'_, R> { + Blobs { opr: self } + } + /// Returns an iterator on the blocks of the pbf file. + pub fn primitive_blocks(&mut self) -> PrimitiveBlocks<'_, R> { + fn and_then_primitive_block(blob_res: Result) -> Result { + blob_res.and_then(|b| primitive_block_from_blob(&b)) + } + PrimitiveBlocks(self.blobs().map(and_then_primitive_block)) + } + + fn push(&mut self, sz: u64) -> Result<()> { + self.buf.clear(); + self.r.by_ref().take(sz).read_to_end(&mut self.buf)?; + assert_eq!(sz, self.buf.len() as u64); + Ok(()) + } + fn try_blob(&mut self, sz: u64) -> Result> { + self.push(sz)?; + let header: BlobHeader = Message::parse_from_bytes(&self.buf)?; + let sz = header.datasize() as u64; + self.push(sz)?; + let blob: Blob = Message::parse_from_bytes(&self.buf)?; + if header.type_() == "OSMData" { + Ok(Some(blob)) + } else if header.type_() == "OSMHeader" { + Ok(None) + } else { + println!("Unknown type: {}", header.type_()); + Ok(None) + } + } + fn next_blob(&mut self) -> Option> { + use byteorder::{BigEndian, ReadBytesExt}; + use std::io::ErrorKind; + if self.finished { + return None; + } + let sz = match self.r.read_u32::() { + Ok(sz) if sz > 64 * 1024 => return Some(Err(Error::InvalidData)), + Ok(sz) => sz, + Err(ref e) if e.kind() == ErrorKind::UnexpectedEof => { + self.finished = true; + return None; + } + Err(e) => { + self.finished = true; + return Some(Err(From::from(e))); + } + } as u64; + match self.try_blob(sz) { + Ok(Some(p)) => Some(Ok(p)), + Ok(None) => self.next_blob(), + Err(e) => { + self.finished = true; + Some(Err(e)) + } + } + } +} + +/// Iterator on the blobs of a file. +pub struct Blobs<'a, R: 'a> { + opr: &'a mut ParReader, +} +impl Iterator for Blobs<'_, R> { + type Item = Result; + fn next(&mut self) -> Option { + self.opr.next_blob() + } +} + +pub_iterator_type! { + #[doc="Iterator on the blocks of a file."] + PrimitiveBlocks['a, R] = iter::Map, fn(Result) -> Result> + where R: Read + 'a +} + +/// Returns an iterator on the blocks of a blob. +pub fn primitive_block_from_blob(blob: &Blob) -> Result { + if blob.has_raw() { + Message::parse_from_bytes(blob.raw()).map_err(From::from) + } else if blob.has_zlib_data() { + use flate2::read::ZlibDecoder; + let r = io::Cursor::new(blob.zlib_data()); + let mut zr = ZlibDecoder::new(r); + Message::parse_from_reader(&mut zr).map_err(From::from) + } else { + Err(Error::UnsupportedData) + } +} + +pub_iterator_type! { + #[doc="Parallel iterator on the `OsmObj` of the pbf file."] + ParIter['a, R] = par_map::FlatMap, + blobs::OsmObjs, + fn(Result) -> blobs::OsmObjs> + where R: io::Read + 'a +} + +pub_iterator_type! { + #[doc="Parallel iterator on the `OsmObj` of the pbf file."] + NodeParIter['a, R] = par_map::FlatMap, + blobs::OsmObjs, + fn(Result) -> blobs::OsmObjs> + where R: io::Read + 'a +} + +pub_iterator_type! { + #[doc="Parallel iterator on the `OsmObj` of the pbf file."] + WayParIter['a, R] = par_map::FlatMap, + blobs::OsmObjs, + fn(Result) -> blobs::OsmObjs> + where R: io::Read + 'a +} + +pub_iterator_type! { + #[doc="Parallel iterator on the `OsmObj` of the pbf file."] + RelationParIter['a, R] = par_map::FlatMap, + blobs::OsmObjs, + fn(Result) -> blobs::OsmObjs> + where R: io::Read + 'a +} diff --git a/src/reader.rs b/src/reader.rs index 2af97975..634e1193 100644 --- a/src/reader.rs +++ b/src/reader.rs @@ -12,7 +12,7 @@ use crate::error::{Error, Result}; use crate::fileformat::{Blob, BlobHeader}; use crate::objects::{OsmId, OsmObj}; use crate::osmformat::PrimitiveBlock; -use par_map::{self, ParMap}; +use crate::store_objs::StoreObjs; use protobuf::Message; use pub_iterator_type::pub_iterator_type; use std::collections::btree_map::BTreeMap; @@ -21,35 +21,17 @@ use std::convert::From; use std::io::{self, Read}; use std::iter; -/// Trait to allow generic objects (not just BTreeMap) in some methods. -pub trait StoreObjs { - /// Insert given object at given key index. - fn insert(&mut self, key: OsmId, value: OsmObj); - /// Check if object contains the given key. - fn contains_key(&self, key: &OsmId) -> bool; -} - -impl StoreObjs for BTreeMap { - fn insert(&mut self, key: OsmId, value: OsmObj) { - self.insert(key, value); - } - - fn contains_key(&self, key: &OsmId) -> bool { - self.contains_key(key) - } -} - /// The object to manage a pbf file. -pub struct OsmPbfReader { +pub struct Reader { buf: Vec, r: R, finished: bool, } -impl OsmPbfReader { +impl Reader { /// Creates an OsmPbfReader from a Read object. - pub fn new(r: R) -> OsmPbfReader { - OsmPbfReader { + pub fn new(r: R) -> Reader { + Reader { buf: vec![], r, finished: false, @@ -61,7 +43,7 @@ impl OsmPbfReader { /// # Example /// /// ``` - /// let mut pbf = osmpbfreader::OsmPbfReader::new(std::io::empty()); + /// let mut pbf = osmpbfreader::Reader::new(std::io::empty()); /// for obj in pbf.iter().map(Result::unwrap) { /// println!("{:?}", obj); /// } @@ -70,70 +52,22 @@ impl OsmPbfReader { Iter(self.blobs().flat_map(blobs::result_blob_into_iter)) } - /// Returns a parallel iterator on the OsmObj of the pbf file. - /// - /// Several threads decode in parallel the file. The memory and - /// CPU usage are guaranteed to be bounded even if the caller stop - /// consuming items. - /// - /// # Example - /// - /// ``` - /// let mut pbf = osmpbfreader::OsmPbfReader::new(std::io::empty()); - /// for obj in pbf.par_iter().map(Result::unwrap) { - /// println!("{:?}", obj); - /// } - /// ``` - pub fn par_iter(&mut self) -> ParIter<'_, R> { - ParIter(self.blobs().par_flat_map(blobs::result_blob_into_iter)) - } - /// Returns an iterator on the Node of the pbf file. pub fn iter_nodes(&mut self) -> NodeIter<'_, R> { NodeIter(self.blobs().flat_map(blobs::result_blob_into_node_iter)) } - /// Returns a parallel iterator on the Node of the pbf file. - /// - /// Several threads decode in parallel the file. The memory and - /// CPU usage are guaranteed to be bounded even if the caller stop - /// consuming items. - pub fn par_iter_nodes(&mut self) -> NodeParIter<'_, R> { - NodeParIter(self.blobs().par_flat_map(blobs::result_blob_into_node_iter)) - } - /// Returns an iterator on the Way of the pbf file. pub fn iter_ways(&mut self) -> WayIter<'_, R> { WayIter(self.blobs().flat_map(blobs::result_blob_into_way_iter)) } - /// Returns a parallel iterator on the Way of the pbf file. - /// - /// Several threads decode in parallel the file. The memory and - /// CPU usage are guaranteed to be bounded even if the caller stop - /// consuming items. - pub fn par_iter_ways(&mut self) -> WayParIter<'_, R> { - WayParIter(self.blobs().par_flat_map(blobs::result_blob_into_way_iter)) - } - /// Returns an iterator on the Relation of the pbf file. pub fn iter_relations(&mut self) -> RelationIter<'_, R> { RelationIter(self.blobs().flat_map(blobs::result_blob_into_relation_iter)) } - /// Returns a parallel iterator on the Relation of the pbf file. - /// - /// Several threads decode in parallel the file. The memory and - /// CPU usage are guaranteed to be bounded even if the caller stop - /// consuming items. - pub fn par_iter_relations(&mut self) -> RelationParIter<'_, R> { - RelationParIter( - self.blobs() - .par_flat_map(blobs::result_blob_into_relation_iter), - ) - } - - /// Rewinds the pbf file to the begining. + /// Rewinds the pbf file to the beginning. /// /// Useful if you want to read several consecutive times the same /// file. @@ -143,7 +77,7 @@ impl OsmPbfReader { /// ``` /// let mut cursor = std::io::Cursor::new([0, 0, 0]); /// cursor.set_position(2); - /// let mut pbf = osmpbfreader::OsmPbfReader::new(cursor); + /// let mut pbf = osmpbfreader::Reader::new(cursor); /// pbf.rewind().unwrap(); /// assert_eq!(pbf.into_inner().position(), 0); /// ``` @@ -169,7 +103,7 @@ impl OsmPbfReader { while !finished { self.rewind()?; finished = true; - for obj in self.par_iter() { + for obj in self.iter() { let obj = obj?; if (!first_pass || !pred(&obj)) && !deps.contains(&obj.id()) { continue; @@ -196,8 +130,7 @@ impl OsmPbfReader { } /// This function give you the ability to find all the objects - /// validating a predicate and all their dependencies. The file - /// will be decoded in parallel. + /// validating a predicate and all their dependencies. /// /// # Example /// @@ -210,7 +143,7 @@ impl OsmPbfReader { /// obj.is_relation() && obj.tags().contains("boundary", "administrative") /// } /// - /// let mut pbf = osmpbfreader::OsmPbfReader::new(std::io::Cursor::new([])); + /// let mut pbf = osmpbfreader::Reader::new(std::io::Cursor::new([])); /// let objs = pbf.get_objs_and_deps(is_admin).unwrap(); /// for (id, obj) in &objs { /// println!("{:?}: {:?}", id, obj); @@ -297,7 +230,7 @@ impl OsmPbfReader { /// Iterator on the blobs of a file. pub struct Blobs<'a, R: 'a> { - opr: &'a mut OsmPbfReader, + opr: &'a mut Reader, } impl Iterator for Blobs<'_, R> { type Item = Result; @@ -332,52 +265,20 @@ pub_iterator_type! { where R: io::Read + 'a } -pub_iterator_type! { - #[doc="Parallel iterator on the `OsmObj` of the pbf file."] - ParIter['a, R] = par_map::FlatMap, - blobs::OsmObjs, - fn(Result) -> blobs::OsmObjs> - where R: io::Read + 'a -} - pub_iterator_type! { #[doc="Iterator on the `OsmObj` of the pbf file."] NodeIter['a, R] = iter::FlatMap, blobs::OsmObjs, fn(Result) -> blobs::OsmObjs> where R: io::Read + 'a } -pub_iterator_type! { - #[doc="Parallel iterator on the `OsmObj` of the pbf file."] - NodeParIter['a, R] = par_map::FlatMap, - blobs::OsmObjs, - fn(Result) -> blobs::OsmObjs> - where R: io::Read + 'a -} - pub_iterator_type! { #[doc="Iterator on the `OsmObj` of the pbf file."] WayIter['a, R] = iter::FlatMap, blobs::OsmObjs, fn(Result) -> blobs::OsmObjs> where R: io::Read + 'a } -pub_iterator_type! { - #[doc="Parallel iterator on the `OsmObj` of the pbf file."] - WayParIter['a, R] = par_map::FlatMap, - blobs::OsmObjs, - fn(Result) -> blobs::OsmObjs> - where R: io::Read + 'a -} - pub_iterator_type! { #[doc="Iterator on the `OsmObj` of the pbf file."] RelationIter['a, R] = iter::FlatMap, blobs::OsmObjs, fn(Result) -> blobs::OsmObjs> where R: io::Read + 'a } - -pub_iterator_type! { - #[doc="Parallel iterator on the `OsmObj` of the pbf file."] - RelationParIter['a, R] = par_map::FlatMap, - blobs::OsmObjs, - fn(Result) -> blobs::OsmObjs> - where R: io::Read + 'a -} diff --git a/src/store_objs.rs b/src/store_objs.rs new file mode 100644 index 00000000..61572ad3 --- /dev/null +++ b/src/store_objs.rs @@ -0,0 +1,20 @@ +use crate::objects::{OsmId, OsmObj}; +use std::collections::BTreeMap; + +/// Trait to allow generic objects (not just BTreeMap) in some methods. +pub trait StoreObjs { + /// Insert given object at given key index. + fn insert(&mut self, key: OsmId, value: OsmObj); + /// Check if object contains the given key. + fn contains_key(&self, key: &OsmId) -> bool; +} + +impl StoreObjs for BTreeMap { + fn insert(&mut self, key: OsmId, value: OsmObj) { + self.insert(key, value); + } + + fn contains_key(&self, key: &OsmId) -> bool { + self.contains_key(key) + } +}