pbfhogg 0.4.0

Fast OpenStreetMap PBF reader and writer for Rust. Read, write, and merge .osm.pbf files with pipelined parallel decoding.
Documentation 
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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230

//! Wire-format way-blob reframe for the dense / sparse pass 2 way arm.
//!
//! Lifted from `external/stage4.rs::reframe_way_blob_with_locations`. Same
//! structure (StringTable byte-pass, group iteration, way reframe) but the
//! coordinate source is a `NodeIndex` lookup instead of a stage-3
//! pre-encoded `coord_payload`.
//!
//! Hot path: decompress -> walk wire -> per-way `NodeIndex::get` per ref ->
//! zigzag-delta-encode lat / lon -> append fields 9 / 10 to original way
//! bytes -> compress. No `BlockBuilder`, no `StringTable::add`, no Info
//! decode / encode, no ref redelta, no tag re-intern.
//!
//! If the input way already declares fields 9 and 10 (input is itself a
//! `LocationsOnWays` PBF), those fields are stripped before append - we
//! own the coordinates from `NodeIndex`, the input's are stale at best.

use protohoggr::{encode_bytes_field, encode_int64_field, encode_tag, encode_varint, zigzag_encode_64, Cursor, PackedSint64Iter, WIRE_LEN, WIRE_VARINT};

use super::NodeIndex;

/// Per-blob scratch buffers reused across way-blob reframes within a worker.
#[derive(Default)]
pub(super) struct WayReframeScratch {
    pub group_ranges: Vec<(usize, usize)>,
    pub scalar_fields: Vec<u8>,
    pub group_out: Vec<u8>,
    pub reframed_way: Vec<u8>,
    pub packed_lats: Vec<u8>,
    pub packed_lons: Vec<u8>,
    pub refs: Vec<i64>,
}

/// Per-blob counters that the caller folds into `Stats` after the reframe.
#[derive(Default)]
pub(super) struct WayReframeStats {
    pub way_count: u64,
    pub min_way_id: i64,
    pub max_way_id: i64,
    pub missing_locations: u64,
}

/// Reframe one decompressed way-blob `PrimitiveBlock` by appending packed
/// `lat` / `lon` fields (9 / 10) to each way using `node_index` for
/// coordinate lookups. Writes the resulting `PrimitiveBlock` bytes into
/// `output` (caller compresses + writes).
///
/// Errors are returned as `String` to match the existing par_iter result
/// channel in `process_slot_batch`.
#[allow(clippy::too_many_lines)]
pub(super) fn reframe_way_blob_with_locations(
    decompressed: &[u8],
    node_index: &NodeIndex,
    output: &mut Vec<u8>,
    scratch: &mut WayReframeScratch,
) -> std::result::Result<WayReframeStats, String> {
    let (st_offset, st_len) = parse_block_top(decompressed, scratch)?;
    let stringtable_bytes = &decompressed[st_offset..st_offset + st_len];

    output.clear();
    encode_bytes_field(output, 1, stringtable_bytes);

    let mut stats = WayReframeStats {
        min_way_id: i64::MAX,
        max_way_id: i64::MIN,
        ..WayReframeStats::default()
    };

    for i in 0..scratch.group_ranges.len() {
        let (gr_offset, gr_len) = scratch.group_ranges[i];
        let group_bytes = &decompressed[gr_offset..gr_offset + gr_len];
        process_group(group_bytes, node_index, output, scratch, &mut stats)?;
    }

    output.extend_from_slice(&scratch.scalar_fields);

    if stats.way_count == 0 {
        stats.min_way_id = 0;
        stats.max_way_id = 0;
    }

    Ok(stats)
}

/// Walk the `PrimitiveBlock` top-level fields. Captures the StringTable
/// byte range (no parse), group offsets, and scalar fields (17-20)
/// verbatim into `scratch`. Returns the StringTable's `(offset, len)` in
/// `decompressed`.
fn parse_block_top(
    decompressed: &[u8],
    scratch: &mut WayReframeScratch,
) -> std::result::Result<(usize, usize), String> {
    scratch.group_ranges.clear();
    scratch.scalar_fields.clear();
    let mut stringtable_range: Option<(usize, usize)> = None;

    let mut cursor = Cursor::new(decompressed);
    while let Some((field, wire_type)) = cursor.read_tag().map_err(|e| format!("reframe block tag: {e}"))? {
        match (field, wire_type) {
            (1, WIRE_LEN) => {
                let data = cursor.read_len_delimited().map_err(|e| format!("reframe st: {e}"))?;
                let offset = data.as_ptr() as usize - decompressed.as_ptr() as usize;
                stringtable_range = Some((offset, data.len()));
            }
            (2, WIRE_LEN) => {
                let data = cursor.read_len_delimited().map_err(|e| format!("reframe group: {e}"))?;
                let offset = data.as_ptr() as usize - decompressed.as_ptr() as usize;
                scratch.group_ranges.push((offset, data.len()));
            }
            (17..=20, WIRE_VARINT) => {
                let raw = cursor.read_raw_field(wire_type).map_err(|e| format!("reframe scalar: {e}"))?;
                encode_tag(&mut scratch.scalar_fields, field, wire_type);
                scratch.scalar_fields.extend_from_slice(raw);
            }
            _ => cursor.skip_field(wire_type).map_err(|e| format!("reframe skip: {e}"))?,
        }
    }

    stringtable_range.ok_or_else(|| "reframe: no StringTable in PrimitiveBlock".to_string())
}

/// Walk one `PrimitiveGroup`. Way submessages get reframed (locations
/// appended); non-way fields are copied verbatim. Builds the
/// `PrimitiveGroup` bytes in `scratch.group_out`, then emits as field 2
/// of the `PrimitiveBlock` in `output`.
fn process_group(
    group_bytes: &[u8],
    node_index: &NodeIndex,
    output: &mut Vec<u8>,
    scratch: &mut WayReframeScratch,
    stats: &mut WayReframeStats,
) -> std::result::Result<(), String> {
    scratch.group_out.clear();

    let mut gr_cursor = Cursor::new(group_bytes);
    while let Some((field, wire_type)) = gr_cursor.read_tag().map_err(|e| format!("reframe gtag: {e}"))? {
        if field == 3 && wire_type == WIRE_LEN {
            let way_bytes = gr_cursor.read_len_delimited().map_err(|e| format!("reframe way: {e}"))?;
            splice_way_locations(way_bytes, node_index, scratch, stats)?;
            encode_bytes_field(&mut scratch.group_out, 3, &scratch.reframed_way);
        } else {
            let raw = gr_cursor.read_raw_field(wire_type).map_err(|e| format!("reframe gskip: {e}"))?;
            encode_tag(&mut scratch.group_out, field, wire_type);
            scratch.group_out.extend_from_slice(raw);
        }
    }

    encode_bytes_field(output, 2, &scratch.group_out);
    Ok(())
}

/// Reframe a single Way submessage. Copies every existing field verbatim
/// EXCEPT fields 9 (lat) and 10 (lon) - those are stripped if present
/// (input was already a LocationsOnWays PBF) since we own the
/// coordinates. Then appends fresh fields 9 / 10 from `node_index`
/// lookups.
fn splice_way_locations(
    way_bytes: &[u8],
    node_index: &NodeIndex,
    scratch: &mut WayReframeScratch,
    stats: &mut WayReframeStats,
) -> std::result::Result<(), String> {
    scratch.reframed_way.clear();
    scratch.refs.clear();
    scratch.packed_lats.clear();
    scratch.packed_lons.clear();

    let mut way_id: i64 = 0;
    let mut have_id = false;

    let mut way_cursor = Cursor::new(way_bytes);
    while let Some((wf, wt)) = way_cursor.read_tag().map_err(|e| format!("reframe wtag: {e}"))? {
        match (wf, wt) {
            (1, WIRE_VARINT) => {
                way_id = way_cursor.read_varint_i64().map_err(|e| format!("reframe wid: {e}"))?;
                have_id = true;
                encode_int64_field(&mut scratch.reframed_way, 1, way_id);
            }
            (8, WIRE_LEN) => {
                let refs_data = way_cursor.read_len_delimited().map_err(|e| format!("reframe wrefs: {e}"))?;
                let mut cum: i64 = 0;
                for delta in PackedSint64Iter::new(refs_data) {
                    cum += delta;
                    scratch.refs.push(cum);
                }
                encode_tag(&mut scratch.reframed_way, 8, WIRE_LEN);
                encode_varint(&mut scratch.reframed_way, refs_data.len() as u64);
                scratch.reframed_way.extend_from_slice(refs_data);
            }
            (9 | 10, WIRE_LEN) => {
                let _ = way_cursor.read_raw_field(wt).map_err(|e| format!("reframe wstrip: {e}"))?;
            }
            _ => {
                let raw = way_cursor.read_raw_field(wt).map_err(|e| format!("reframe wskip: {e}"))?;
                encode_tag(&mut scratch.reframed_way, wf, wt);
                scratch.reframed_way.extend_from_slice(raw);
            }
        }
    }

    if !have_id {
        return Err("reframe: way without field 1 (id)".to_string());
    }

    if way_id < stats.min_way_id { stats.min_way_id = way_id; }
    if way_id > stats.max_way_id { stats.max_way_id = way_id; }
    stats.way_count += 1;

    let mut last_lat: i64 = 0;
    let mut last_lon: i64 = 0;
    for &node_id in &scratch.refs {
        let (lat, lon) = match node_index.get(node_id) {
            Some(loc) => (i64::from(loc.0), i64::from(loc.1)),
            None => {
                stats.missing_locations += 1;
                (0, 0)
            }
        };
        encode_varint(&mut scratch.packed_lats, zigzag_encode_64(lat - last_lat));
        encode_varint(&mut scratch.packed_lons, zigzag_encode_64(lon - last_lon));
        last_lat = lat;
        last_lon = lon;
    }

    if !scratch.refs.is_empty() {
        encode_bytes_field(&mut scratch.reframed_way, 9, &scratch.packed_lats);
        encode_bytes_field(&mut scratch.reframed_way, 10, &scratch.packed_lons);
    }

    Ok(())
}