osm_pbf_proto/
primitives.rs

1use std::ops::Deref;
2
3use bitflags::bitflags;
4use bytes::Bytes;
5use protobuf::SpecialFields;
6
7use crate::osmformat::{
8    ChangeSet, DenseInfo, Info, Node, PrimitiveBlock, PrimitiveGroup, Relation, Way,
9};
10
11bitflags! {
12    pub struct PrimitiveType: u32 {
13        const NODE = 1;
14        const WAY = 2;
15        const RELATION = 4;
16        const CHANGE_SET = 8;
17
18        const DEFAULT = Self::NODE.bits() | Self::WAY.bits() | Self::RELATION.bits();
19    }
20}
21
22#[derive(Copy, Clone)]
23pub struct PrimitiveRef<'l, T: ?Sized> {
24    value: &'l T,
25    block: &'l PrimitiveBlock,
26}
27
28impl<T: ?Sized> Deref for PrimitiveRef<'_, T> {
29    type Target = T;
30    fn deref(&self) -> &T {
31        self.value
32    }
33}
34
35pub type PrimitiveGroupRef<'l> = PrimitiveRef<'l, PrimitiveGroup>;
36pub type WayRef<'l> = PrimitiveRef<'l, Way>;
37pub type RelationRef<'l> = PrimitiveRef<'l, Relation>;
38pub type ChangeSetRef<'l> = PrimitiveRef<'l, ChangeSet>;
39
40#[derive(PartialEq, Clone, Debug)]
41pub struct NodeRef<'l> {
42    pub id: i64,
43    pub nano_lat: i64,
44    pub nano_lon: i64,
45    pub index: usize,
46    data: NodeData<'l>,
47    block: &'l PrimitiveBlock,
48}
49
50impl<'l> NodeRef<'l> {
51    #[inline]
52    fn from_node(index: usize, node: &'l Node, block: &'l PrimitiveBlock) -> Self {
53        Self {
54            id: node.id(),
55            nano_lat: block.lat_offset() + node.lat() * block.granularity() as i64,
56            nano_lon: block.lon_offset() + node.lon() * block.granularity() as i64,
57            index,
58            data: NodeData::Node {
59                keys: &node.keys,
60                vals: &node.vals,
61                info: &node.info,
62            },
63            block,
64        }
65    }
66
67    #[inline]
68    fn from_dense_node(
69        index: usize,
70        dense_state: &DenseState,
71        kv_pairs: &'l [i32],
72        info: &'l DenseInfo,
73        block: &'l PrimitiveBlock,
74    ) -> Self {
75        Self {
76            id: dense_state.id,
77            nano_lat: block.lat_offset() + dense_state.lat * block.granularity() as i64,
78            nano_lon: block.lon_offset() + dense_state.lon * block.granularity() as i64,
79            index,
80            data: NodeData::DenseNode { kv_pairs, info },
81            block,
82        }
83    }
84
85    #[inline]
86    pub const fn id(&self) -> i64 {
87        self.id
88    }
89
90    /// Latitude in degrees.
91    #[inline]
92    pub fn lat(&self) -> f64 {
93        self.nano_lat as f64 * 1e-9
94    }
95
96    /// Longitude in degrees.
97    #[inline]
98    pub fn lon(&self) -> f64 {
99        self.nano_lon as f64 * 1e-9
100    }
101
102    pub fn info(&self) -> Info {
103        match self.data {
104            NodeData::Node { info, .. } => info.clone(),
105            NodeData::DenseNode { info, .. } => Info {
106                version: info.version.get(self.index).copied(),
107                timestamp: info.timestamp.get(self.index).copied(),
108                changeset: info.changeset.get(self.index).copied(),
109                uid: info.uid.get(self.index).copied(),
110                user_sid: info.user_sid.get(self.index).copied().map(|v| v as u32),
111                visible: info.visible.get(self.index).copied(),
112                special_fields: SpecialFields::new(),
113            },
114        }
115    }
116
117    pub fn tags(&self) -> Tags<'_> {
118        match self.data {
119            NodeData::Node { keys, vals, .. } => Tags {
120                kv: TagsData::Normal(keys.iter(), vals.iter()),
121                s: &self.block.stringtable.s,
122            },
123            NodeData::DenseNode { kv_pairs, .. } => Tags {
124                kv: TagsData::Dense(kv_pairs.iter()),
125                s: &self.block.stringtable.s,
126            },
127        }
128    }
129}
130
131#[derive(PartialEq, Clone, Debug)]
132enum NodeData<'l> {
133    Node {
134        keys: &'l [u32],
135        vals: &'l [u32],
136        info: &'l Info,
137    },
138    DenseNode {
139        kv_pairs: &'l [i32],
140        info: &'l DenseInfo,
141    },
142}
143
144#[derive(Default)]
145struct DenseState {
146    id: i64,
147    lat: i64,
148    lon: i64,
149    kv_pos: usize,
150}
151
152impl WayRef<'_> {
153    #[inline]
154    pub fn tags(&self) -> Tags<'_> {
155        Tags {
156            kv: TagsData::Normal(self.value.keys.iter(), self.value.vals.iter()),
157            s: &self.block.stringtable.s,
158        }
159    }
160}
161
162impl RelationRef<'_> {
163    #[inline]
164    pub fn tags(&self) -> Tags<'_> {
165        Tags {
166            kv: TagsData::Normal(self.value.keys.iter(), self.value.vals.iter()),
167            s: &self.block.stringtable.s,
168        }
169    }
170}
171
172#[derive(Clone, Debug)]
173enum TagsData<'l> {
174    Normal(std::slice::Iter<'l, u32>, std::slice::Iter<'l, u32>),
175    Dense(std::slice::Iter<'l, i32>),
176}
177
178#[derive(Clone, Debug)]
179pub struct Tags<'l> {
180    kv: TagsData<'l>,
181    s: &'l [Bytes],
182}
183
184impl<'l> Tags<'l> {
185    pub fn get(&self, key: &str) -> Option<&'l str> {
186        let key = key.as_bytes();
187        let s = self.s;
188        match self.kv.clone() {
189            TagsData::Normal(mut keys, mut vals) => loop {
190                let key_index = keys.next().copied()? as usize;
191                if let Some(k) = s.get(key_index) {
192                    if k == key {
193                        let value_index = vals.next().copied()? as usize;
194                        return s.get(value_index).and_then(|b| std::str::from_utf8(b).ok());
195                    }
196                }
197            },
198            TagsData::Dense(mut kv_pairs) => loop {
199                let key_index = kv_pairs.next().copied()? as usize;
200                let value_index = kv_pairs.next().copied()? as usize;
201                if let Some(k) = s.get(key_index) {
202                    if k == key {
203                        return s.get(value_index).and_then(|b| std::str::from_utf8(b).ok());
204                    }
205                }
206            },
207        }
208    }
209}
210
211impl<'l> Iterator for Tags<'l> {
212    type Item = (&'l str, &'l str);
213    #[inline]
214    fn next(&mut self) -> Option<Self::Item> {
215        let s = self.s;
216        loop {
217            let key_index;
218            let value_index;
219            match &mut self.kv {
220                TagsData::Normal(keys, vals) => {
221                    key_index = keys.next().copied()? as usize;
222                    value_index = vals.next().copied()? as usize;
223                }
224                TagsData::Dense(kv_pairs) => {
225                    key_index = kv_pairs.next().copied()? as usize;
226                    value_index = kv_pairs.next().copied()? as usize;
227                }
228            }
229            // get string as utf8 (ognore invalid)
230            let Some(key) = s.get(key_index).and_then(|b| std::str::from_utf8(b).ok()) else {
231                continue;
232            };
233            let Some(value) = s.get(value_index).and_then(|b| std::str::from_utf8(b).ok()) else {
234                continue;
235            };
236            return Some((key, value));
237        }
238    }
239
240    #[inline]
241    fn size_hint(&self) -> (usize, Option<usize>) {
242        match &self.kv {
243            TagsData::Normal(keys, vals) => {
244                let (keys_lower, keys_upper) = keys.size_hint();
245                let (values_lower, values_upper) = vals.size_hint();
246                (
247                    keys_lower.min(values_lower),
248                    keys_upper.zip(values_upper).map(|(k, v)| k.min(v)),
249                )
250            }
251            TagsData::Dense(kv_pairs) => {
252                let (lower, upper) = kv_pairs.size_hint();
253                (lower / 2, upper.map(|l| l / 2))
254            }
255        }
256    }
257}
258
259impl<'l> std::iter::FusedIterator for Tags<'l> {}
260
261#[non_exhaustive]
262pub enum Primitive<'l> {
263    Node(NodeRef<'l>),
264    Way(WayRef<'l>),
265    Relation(RelationRef<'l>),
266    ChangeSet(ChangeSetRef<'l>),
267}
268
269pub struct PrimitivesIter<'l> {
270    block: &'l PrimitiveBlock,
271    groups: &'l [PrimitiveGroup],
272    filter: PrimitiveType,
273    group_pos: usize,
274    prim_pos: usize,
275    dense_state: DenseState,
276}
277
278impl PrimitiveBlock {
279    #[inline]
280    pub fn primitives(&self) -> PrimitivesIter<'_> {
281        PrimitivesIter {
282            block: self,
283            groups: &self.primitivegroup,
284            filter: PrimitiveType::DEFAULT,
285            group_pos: 0,
286            prim_pos: 0,
287            dense_state: DenseState::default(),
288        }
289    }
290    #[inline]
291    pub fn primitivegroup(&self, i: usize) -> Option<PrimitiveGroupRef<'_>> {
292        Some(PrimitiveGroupRef {
293            value: self.primitivegroup.get(i)?,
294            block: self,
295        })
296    }
297}
298
299impl<'l> PrimitiveGroupRef<'l> {
300    #[inline]
301    pub fn primitives(self) -> PrimitivesIter<'l> {
302        PrimitivesIter {
303            block: self.block,
304            groups: std::slice::from_ref(self.value),
305            filter: PrimitiveType::DEFAULT,
306            group_pos: 0,
307            prim_pos: 0,
308            dense_state: DenseState::default(),
309        }
310    }
311}
312
313impl<'l> PrimitivesIter<'l> {
314    #[inline]
315    pub fn filter_types(mut self, types: PrimitiveType) -> Self {
316        self.filter = types;
317        self
318    }
319}
320
321impl<'l> IntoIterator for &'l PrimitiveBlock {
322    type Item = Primitive<'l>;
323    type IntoIter = PrimitivesIter<'l>;
324    #[inline]
325    fn into_iter(self) -> PrimitivesIter<'l> {
326        self.primitives()
327    }
328}
329
330impl<'l> IntoIterator for PrimitiveGroupRef<'l> {
331    type Item = Primitive<'l>;
332    type IntoIter = PrimitivesIter<'l>;
333    #[inline]
334    fn into_iter(self) -> PrimitivesIter<'l> {
335        self.primitives()
336    }
337}
338
339impl<'l> Iterator for PrimitivesIter<'l> {
340    type Item = Primitive<'l>;
341    fn next(&mut self) -> Option<Self::Item> {
342        loop {
343            let group = self.groups.get(self.group_pos)?;
344            if self.filter.contains(PrimitiveType::NODE) && !group.nodes.is_empty() {
345                let index = self.prim_pos;
346                if let Some(n) = group.nodes.get(index) {
347                    self.prim_pos = index + 1;
348                    return Some(Primitive::Node(NodeRef::from_node(index, n, self.block)));
349                }
350            } else if self.filter.contains(PrimitiveType::NODE) && group.dense.is_some() {
351                let dense = &group.dense;
352                let index = self.prim_pos;
353                if let (Some(id), Some(lat), Some(lon)) = (
354                    dense.id.get(index).copied(),
355                    dense.lat.get(index).copied(),
356                    dense.lon.get(index).copied(),
357                ) {
358                    self.prim_pos = index + 1;
359                    self.dense_state.id += id;
360                    self.dense_state.lat += lat;
361                    self.dense_state.lon += lon;
362
363                    // find range for key-value pairs
364                    let kv_from = self.dense_state.kv_pos;
365                    while let Some(k) = dense.keys_vals.get(self.dense_state.kv_pos).copied() {
366                        if k == 0 {
367                            self.dense_state.kv_pos += 1;
368                        } else {
369                            self.dense_state.kv_pos += 2;
370                        }
371                    }
372                    let key_values = &dense.keys_vals[kv_from..self.dense_state.kv_pos];
373
374                    let n = NodeRef::from_dense_node(
375                        index,
376                        &self.dense_state,
377                        key_values,
378                        &dense.denseinfo,
379                        self.block,
380                    );
381                    return Some(Primitive::Node(n));
382                }
383                // reset dense state for next group
384                self.dense_state = DenseState::default();
385            } else if self.filter.contains(PrimitiveType::WAY) && !group.ways.is_empty() {
386                if let Some(w) = group.ways.get(self.prim_pos) {
387                    self.prim_pos += 1;
388                    return Some(Primitive::Way(PrimitiveRef {
389                        value: w,
390                        block: self.block,
391                    }));
392                }
393            } else if self.filter.contains(PrimitiveType::RELATION) && !group.relations.is_empty() {
394                if let Some(r) = group.relations.get(self.prim_pos) {
395                    self.prim_pos += 1;
396                    return Some(Primitive::Relation(PrimitiveRef {
397                        value: r,
398                        block: self.block,
399                    }));
400                }
401            } else if self.filter.contains(PrimitiveType::CHANGE_SET)
402                && !group.changesets.is_empty()
403            {
404                if let Some(c) = group.changesets.get(self.prim_pos) {
405                    self.prim_pos += 1;
406                    return Some(Primitive::ChangeSet(PrimitiveRef {
407                        value: c,
408                        block: self.block,
409                    }));
410                }
411            }
412            self.group_pos += 1;
413        }
414    }
415}
osm_pbf_proto/primitives.rs

osm_pbf_proto/
primitives.rs
