osm_io/osm/pbf/
file_block.rs

1use std::fs::File;
2use std::io::{Cursor, Read, Seek, SeekFrom, Write};
3
4use anyhow::{anyhow, Context};
5use flate2::bufread::ZlibDecoder;
6use flate2::Compression;
7use flate2::write::ZlibEncoder;
8use prost::Message;
9
10use crate::osmpbf;
11use crate::osm::model::bounding_box::BoundingBox;
12use crate::osm::model::element::Element;
13use crate::osm::pbf::blob_desc::BlobDesc;
14use crate::osm::pbf::compression_type::CompressionType;
15use crate::osm::pbf::file_block_metadata::FileBlockMetadata;
16use crate::osm::pbf::osm_data::OsmData;
17use crate::osm::pbf::osm_header::OsmHeader;
18use crate::osmpbf::BlobHeader;
19use crate::osmpbf::blob::Data;
20
21/// A header or data file block in *.osm.pbf file
22#[derive(Debug)]
23pub enum FileBlock {
24    Header {
25        metadata: FileBlockMetadata,
26        header: OsmHeader,
27    },
28    Data {
29        metadata: FileBlockMetadata,
30        data: OsmData,
31    },
32}
33
34impl FileBlock {
35    pub(crate) fn new(index: usize, blob_type: String, data: Vec<u8>) -> Result<FileBlock, anyhow::Error> {
36        let blob_type_str = blob_type.as_str();
37        match blob_type_str {
38            "OSMHeader" => {
39                Ok(
40                    FileBlock::Header {
41                        metadata: FileBlockMetadata::new(blob_type, index),
42                        header: OsmHeader::from_bytes(data)?,
43                    }
44                )
45            }
46            "OSMData" => {
47                Ok(
48                    FileBlock::Data {
49                        metadata: FileBlockMetadata::new(blob_type, index),
50                        data: OsmData::new(data)?,
51                    }
52                )
53            }
54            _ => {
55                Err(anyhow!("Failed to decode file block"))
56            }
57        }
58    }
59
60    #[allow(dead_code)]
61    pub(crate) fn index(&self) -> usize {
62        match self {
63            FileBlock::Header { metadata, header: _header } => {
64                metadata.index()
65            }
66            FileBlock::Data { metadata, data: _data } => {
67                metadata.index()
68            }
69        }
70    }
71
72    pub(crate) fn from_elements(index: usize, elements: Vec<Element>) -> FileBlock {
73        FileBlock::Data {
74            metadata: FileBlockMetadata::new("OSMData".to_string(), index),
75            data: OsmData::from_elements(elements, None),
76        }
77    }
78
79    #[allow(dead_code)]
80    pub(crate) fn compute_bounding_box(&self) -> Option<BoundingBox> {
81        match self {
82            FileBlock::Header { metadata: _, header } => {
83                header.info().bounding_box().clone()
84            }
85            FileBlock::Data { metadata: _, data } => {
86                data.compute_bounding_box()
87            }
88        }
89    }
90
91    pub(crate) fn from_header(osm_header: OsmHeader) -> FileBlock {
92        FileBlock::Header {
93            metadata: FileBlockMetadata::new("OSMHeader".to_string(), 0),
94            header: osm_header.clone(),
95        }
96    }
97
98    fn zlib_decode(data: Vec<u8>, raw_size: usize) -> Result<Vec<u8>, anyhow::Error> {
99        let mut decoder = ZlibDecoder::new(data.as_slice());
100        let mut decoded = vec![0_u8; raw_size];
101        decoder.read_exact(&mut decoded)?;
102        Ok(decoded)
103    }
104
105    fn zlib_encode(buf: Vec<u8>, compression_level: Compression) -> Result<Vec<u8>, anyhow::Error> {
106        let mut encoder = ZlibEncoder::new(Vec::new(), compression_level);
107        encoder.write_all(buf.as_slice())?;
108        encoder.flush()?;
109        let encoded = encoder.finish()?;
110        Ok(encoded)
111    }
112
113    pub(crate) fn read_blob_data(blob: osmpbf::Blob) -> Result<Vec<u8>, anyhow::Error> {
114        match blob.data {
115            None => {
116                Err(
117                    anyhow!("Input file too short")
118                )
119            }
120            Some(data) => {
121                match data {
122                    Data::Raw(_) => {
123                        Err(
124                            // TODO:
125                            anyhow!("Raw data type not implemented")
126                        )
127                    }
128                    Data::ZlibData(zlib_data) => {
129                        // for now ignore that the uncompressed size is optional
130                        FileBlock::zlib_decode(zlib_data, blob.raw_size.unwrap() as usize)
131                    }
132                    Data::LzmaData(_) => {
133                        Err(
134                            // TODO:
135                            anyhow!("Lzma data type not implemented")
136                        )
137                    }
138                    Data::ObsoleteBzip2Data(_) => {
139                        Err(
140                            anyhow!("Obsolete Bzip data type not implemented")
141                        )
142                    }
143                    Data::Lz4Data(_) => {
144                        Err(
145                            // TODO:
146                            anyhow!("Lz4 data type not implemented")
147                        )
148                    }
149                    Data::ZstdData(_) => {
150                        Err(
151                            anyhow!("Zstd data type not implemented")
152                        )
153                    }
154                }
155            }
156        }
157    }
158
159    pub(crate) fn from_blob_desc(blob_desc: &BlobDesc) -> Result<FileBlock, anyhow::Error> {
160        let mut file = File::open(blob_desc.path()).with_context(
161            || anyhow!("Failed to open {:?} for reading", blob_desc.path())
162        )?;
163        file.seek(SeekFrom::Start(blob_desc.start())).with_context(
164            || anyhow!("Failed seek to {} in {:?} ", blob_desc.start(), blob_desc.path())
165        )?;
166        let mut blob_buffer = vec![0; blob_desc.length() as usize];
167        file.read_exact(&mut blob_buffer).ok().with_context(
168            || anyhow!("Failed to read {} bytes from {:?} ", blob_desc.length(), blob_desc.path())
169        )?;
170        Self::deserialize(blob_desc, &mut blob_buffer)
171    }
172
173    pub(crate) fn serialize(file_block: &FileBlock, compression: CompressionType) -> Result<(Vec<u8>, Vec<u8>), anyhow::Error> {
174        let (blob_type, compression_level, block_data) = match file_block {
175            FileBlock::Header { metadata: _, header } => {
176                ("OSMHeader".to_string(), Compression::none(), header.serialize()?)
177            }
178            FileBlock::Data { metadata: _, data } => {
179                ("OSMData".to_string(), Compression::default(), data.serialize()?)
180            }
181        };
182
183        let mut raw_size = None;
184        let mut data = None;
185        if !block_data.is_empty() {
186            raw_size = Some(block_data.len() as i32);
187            data = match compression {
188                CompressionType::Uncompressed => {
189                    Some(Data::Raw(block_data))
190                }
191                CompressionType::Zlib => {
192                    let encoded = Self::zlib_encode(block_data, compression_level)?;
193                    Some(Data::ZlibData(encoded))
194                }
195            };
196        }
197
198        let blob = osmpbf::Blob {
199            raw_size,
200            data,
201        };
202        let body = blob.encode_to_vec();
203
204        let blob_header = BlobHeader {
205            r#type: blob_type,
206            indexdata: None,
207            datasize: body.len() as i32,
208        };
209
210
211        let header = blob_header.encode_to_vec();
212
213        Ok((header, body))
214    }
215
216    fn deserialize(blob_desc: &BlobDesc, blob_buffer: &mut Vec<u8>) -> Result<FileBlock, anyhow::Error> {
217        // use BlobDesc rather than BlobHeader to skip reading again the blob header
218        let protobuf_blob = osmpbf::Blob::decode(&mut Cursor::new(blob_buffer)).with_context(
219            || anyhow!("Failed to decode a message from blob {} from {:?}", blob_desc.index(), blob_desc.path())
220        )?;
221        let data = FileBlock::read_blob_data(protobuf_blob)?;
222        FileBlock::new(blob_desc.index(), blob_desc.t(), data)
223    }
224
225    #[allow(dead_code)]
226    pub(crate) fn metadata(&self) -> &FileBlockMetadata {
227        match self {
228            FileBlock::Header { metadata, header: _ } => {
229                metadata
230            }
231            FileBlock::Data { metadata, data: _ } => {
232                metadata
233            }
234        }
235    }
236
237    pub(crate) fn as_osm_header(&self) -> Result<&OsmHeader, anyhow::Error> {
238        match self {
239            FileBlock::Header { header, .. } => {
240                Ok(header)
241            }
242            FileBlock::Data { .. } => {
243                Err(anyhow!("Not an OSMHeader"))
244            }
245        }
246    }
247    pub(crate) fn is_osm_header(&self) -> bool {
248        match self {
249            FileBlock::Header { header: _, .. } => {
250                true
251            }
252            FileBlock::Data { .. } => {
253                false
254            }
255        }
256    }
257
258    pub(crate) fn is_osm_data(&self) -> bool {
259        !self.is_osm_header()
260    }
261
262    #[allow(dead_code)]
263    pub(crate) fn as_osm_data(&self) -> Result<&OsmData, anyhow::Error> {
264        match self {
265            FileBlock::Header { .. } => {
266                Err(anyhow!("Not an OSMData"))
267            }
268            FileBlock::Data { data, .. } => {
269                Ok(data)
270            }
271        }
272    }
273
274    #[allow(dead_code)]
275    pub(crate) fn elements(&self) -> &Vec<Element> {
276        self.as_osm_data().unwrap().elements()
277    }
278
279    pub(crate) fn take_elements(&mut self) -> Vec<Element> {
280        match self {
281            FileBlock::Header { .. } => {
282                panic!("Not a Data variant")
283            }
284            FileBlock::Data { data, .. } => {
285                data.take_elements()
286            }
287        }
288    }
289}
290
291impl Default for FileBlock {
292    fn default() -> Self {
293        FileBlock::Data { metadata: Default::default(), data: Default::default() }
294    }
295}
osm_io/osm/pbf/file_block.rs

osm_io/osm/pbf/
file_block.rs
