Skip to main content

pbf_craft/readers/
iter_reader.rs

1use std::fs::File;
2use std::io::{BufReader, Read};
3use std::path::Path;
4
5use super::raw_reader::PbfReader;
6use super::raw_reader::ReaderProgress;
7use super::traits::BlobData;
8use crate::models::{Element, ElementType};
9
10/// A reader that provides an iterable interface for reading PBF data.
11///
12/// The `IterableReader` struct allows for sequential reading of PBF data by iterating over blobs
13/// and elements. It is generic over a type `R` that implements the `Read` and `Send` traits, which
14/// provide the necessary methods for reading PBF data from a source.
15///
16/// # Type Parameters
17///
18/// * `R` - A type that implements the `Read` and `Send` traits, providing methods for reading PBF data.
19///
20/// # Example
21///
22/// ```rust
23/// use pbf_craft::models::{Element, ElementType};
24/// use pbf_craft::readers::IterableReader;
25///
26/// let mut reader = IterableReader::from_path("resources/andorra-latest.osm.pbf").unwrap();
27/// for element in reader {
28///    // Process the element
29/// }
30/// ```
31pub struct IterableReader<R: Read + Send> {
32    pbf_reader: PbfReader<R>,
33    current_blob: Option<BlobData>,
34    current_element_type: ElementType,
35    current_element_index: usize,
36    /// A read error that occurred between blobs. `Iterator` cannot return `Err`, so the error
37    /// is stored here and surfaced (with context) on the next call to `next()`.
38    read_error: Option<anyhow::Error>,
39}
40
41impl<R: Read + Send> IterableReader<R> {
42    /// Creates a new `IterableReader` from a raw pbf reader, reading the first blob eagerly so
43    /// a malformed stream fails at construction time.
44    pub fn new(mut pbf_reader: PbfReader<R>) -> anyhow::Result<Self> {
45        Ok(Self {
46            current_blob: pbf_reader.read_next_blob()?,
47            current_element_type: ElementType::Node,
48            current_element_index: 0,
49            read_error: None,
50            pbf_reader,
51        })
52    }
53
54    /// Reports the reader's consumption progress. Delegates to the underlying `PbfReader`,
55    /// so `total_bytes` is known when created via `from_path`.
56    pub fn progress(&self) -> ReaderProgress {
57        self.pbf_reader.progress()
58    }
59
60    fn next_element(&mut self) -> Option<Element> {
61        loop {
62            if let Some(blob) = &self.current_blob {
63                if ElementType::Node == self.current_element_type {
64                    if self.current_element_index < blob.nodes.len() {
65                        let node = blob.nodes.get(self.current_element_index).unwrap();
66                        self.current_element_index += 1;
67                        return Some(Element::Node(node.clone()));
68                    } else {
69                        self.current_element_type = ElementType::Way;
70                        self.current_element_index = 0;
71                    }
72                }
73                if ElementType::Way == self.current_element_type {
74                    if self.current_element_index < blob.ways.len() {
75                        let way = blob.ways.get(self.current_element_index).unwrap();
76                        self.current_element_index += 1;
77                        return Some(Element::Way(way.clone()));
78                    } else {
79                        self.current_element_type = ElementType::Relation;
80                        self.current_element_index = 0;
81                    }
82                }
83                if ElementType::Relation == self.current_element_type {
84                    if self.current_element_index < blob.relations.len() {
85                        let relation = blob.relations.get(self.current_element_index).unwrap();
86                        self.current_element_index += 1;
87                        return Some(Element::Relation(relation.clone()));
88                    } else {
89                        match self.pbf_reader.read_next_blob() {
90                            Ok(next) => self.current_blob = next,
91                            Err(err) => {
92                                self.read_error = Some(err);
93                                self.current_blob = None;
94                            }
95                        }
96                        self.current_element_type = ElementType::Node;
97                        self.current_element_index = 0;
98                    }
99                }
100            } else {
101                return None;
102            }
103        }
104    }
105}
106
107impl<R: Read + Send> Iterator for IterableReader<R> {
108    type Item = Element;
109
110    fn next(&mut self) -> Option<Self::Item> {
111        if let Some(err) = &self.read_error {
112            panic!("PBF read error during iteration: {}", err);
113        }
114        self.next_element()
115    }
116}
117
118impl IterableReader<BufReader<File>> {
119    /// Creates a new `IterableReader` from a file path.
120    pub fn from_path<P: AsRef<Path>>(path: P) -> anyhow::Result<Self> {
121        let pbf_reader = PbfReader::from_path(path)?;
122        Self::new(pbf_reader)
123    }
124}