1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

use std::vec::IntoIter;

use crate::osm::model::element::Element;
use crate::osm::pbf::file_block::FileBlock;
use crate::osm::pbf::file_block_iterator::FileBlockIterator;

/// Iterate over elements in *.osm.pbf file
pub struct ElementIterator {
    file_block_iterator: FileBlockIterator,
    element_iterator: Option<IntoIter<Element>>,
}

impl ElementIterator {
    pub(crate) fn new(mut file_block_iterator: FileBlockIterator) -> ElementIterator {
        // skip the header
        file_block_iterator.next();
        let element_iterator = Self::create_element_iterator(&mut file_block_iterator);
        ElementIterator {
            file_block_iterator,
            element_iterator,
        }
    }

    fn create_element_iterator(file_block_iterator: &mut FileBlockIterator) -> Option<IntoIter<Element>> {
        if let Some(current_block) = file_block_iterator.next() {
            if let FileBlock::Data { metadata: _, mut data } = current_block {
                Some(data.take_elements().into_iter())
            } else {
                None
            }
        } else {
            None
        }
    }
}

impl Iterator for ElementIterator {
    type Item = Element;

    fn next(&mut self) -> Option<Self::Item> {
        match &mut self.element_iterator {
            None => {
                None
            }
            Some(element_iterator) => {
                let element = element_iterator.next();
                match &element {
                    None => {
                        self.element_iterator = ElementIterator::create_element_iterator(&mut self.file_block_iterator);
                        self.next()
                    }
                    Some(_) => {
                        element
                    }
                }
            }
        }
    }
}