gistools/writers/pmtiles/
mod.rs

1use crate::{
2    parsers::Writer,
3    readers::{
4        PMDirectory, PMEntry, PMHeader, PMTilePos, PMTileType, S2_PM_HEADER_SIZE_BYTES,
5        S2_PM_ROOT_SIZE, S2PMEntries, S2PMHeader,
6    },
7    util::CompressionFormat,
8    writers::TileWriter,
9};
10use alloc::{vec, vec::Vec};
11use s2_tilejson::Metadata;
12use s2json::Face;
13
14/// The result of an optimized directory computation
15#[derive(Debug, Clone, Default)]
16struct OptimizedDirectory {
17    /// The root directory bytes
18    pub root_bytes: Vec<u8>,
19    /// The leaf directories bytes
20    pub leaves_bytes: Vec<u8>,
21    /// The number of leaf directories
22    #[allow(dead_code)]
23    pub num_leaves: u64,
24}
25impl OptimizedDirectory {
26    /// Optimize the directory for storage
27    pub fn optimize_directories(
28        directory: &mut PMDirectory,
29        target_root_length: usize,
30    ) -> OptimizedDirectory {
31        directory.entries.sort_by(|a, b| a.tile_id.cmp(&b.tile_id));
32        let test_bytes = directory.serialize();
33        if test_bytes.len() < target_root_length {
34            OptimizedDirectory { root_bytes: test_bytes, leaves_bytes: Vec::new(), num_leaves: 0 }
35        } else {
36            let mut leaf_size = 4096;
37            loop {
38                let build = OptimizedDirectory::build_root_leaves(directory, leaf_size);
39                if build.root_bytes.len() < target_root_length {
40                    return build;
41                }
42                leaf_size *= 2;
43            }
44        }
45    }
46
47    /// Build the root and leaf directories
48    pub fn build_root_leaves(directory: &PMDirectory, leaf_size: usize) -> OptimizedDirectory {
49        let mut root_entries = PMDirectory::default();
50        let mut leaves_bytes = Vec::<u8>::new();
51        let mut num_leaves = 0;
52
53        let mut i = 0;
54        let entries = &directory.entries;
55        while i < entries.len() {
56            num_leaves += 1;
57            let mut end = i + leaf_size;
58            if i + leaf_size > entries.len() {
59                end = entries.len();
60            }
61            let new_dir_slice = PMDirectory::new(entries[i..end].to_vec());
62            let serialized = new_dir_slice.serialize();
63            let entry = PMEntry {
64                tile_id: entries[i].tile_id,
65                offset: leaves_bytes.len() as u64,
66                length: serialized.len() as u32,
67                run_length: 0,
68            };
69            root_entries.entries.push(entry);
70            leaves_bytes.extend(serialized);
71            i += leaf_size;
72        }
73
74        OptimizedDirectory { root_bytes: root_entries.serialize(), leaves_bytes, num_leaves }
75    }
76}
77
78/// The data writer
79pub trait DataWriter: core::fmt::Debug {
80    /// Write data at the specified offset
81    fn write_data(&mut self, data: &[u8], offset: u64);
82    /// Append data to the end of the storage
83    fn append_data(&mut self, data: &[u8]);
84    /// Assuming local writer, take ownership of the data when finished writing it
85    fn take(&self) -> Vec<u8>;
86}
87
88/// # (S2) PM Tiles Writer
89///
90/// ## Description
91/// A PM or S2PM Tile Writer to store tile and metadata in a cloud optimized format
92///
93/// Implements the [`TileWriter`] interface.
94///
95/// Takes a [`Writer`] input to write data to.
96///
97/// ## Usage
98///
99/// The methods you have access to:
100/// - [`PMTilesWriter::new`]: Create a new PMTilesWriter
101/// - [`PMTilesWriter::take`]: Take ownership of the data (assuming local buffered data and not a filesystem writer)
102/// - [`PMTilesWriter::write_tile`]: Write a tile to the PMTiles file given its tile ID.
103/// - [`PMTilesWriter::write_tile_wm`]: Write a Web Mercator tile to the folder location given its (zoom, x, y) coordinates.
104/// - [`PMTilesWriter::write_tile_s2`]: Write a S2 tile to the folder location given its (face, zoom, x, y) coordinates.
105/// - [`PMTilesWriter::commit`]: Write the metadata to the folder location.
106///
107/// ```rust
108/// use gistools::{util::CompressionFormat, parsers::BufferWriter, writers::{TileWriter, PMTilesWriter}};
109/// use s2_tilejson::Metadata;
110/// use s2json::Face;
111///
112/// let local_writer = BufferWriter::default();
113/// let mut pmtiles_writer = PMTilesWriter::new(local_writer, CompressionFormat::None);
114///
115/// // setup data
116/// let tmp_str = "hello world";
117/// // write data in tile
118/// pmtiles_writer.write_tile_s2(Face::Face0, 0, 0, 0, tmp_str.as_bytes().to_vec());
119/// pmtiles_writer.write_tile_s2(Face::Face3, 2, 1, 1, tmp_str.as_bytes().to_vec());
120/// // finish (we don't need to specify the compression again)
121/// pmtiles_writer.commit(Metadata::default(), None);
122///
123/// let pmtiles_data = pmtiles_writer.take();
124/// ```
125#[derive(Debug)]
126pub struct PMTilesWriter<W: Writer> {
127    tile_entries: PMDirectory,
128    s2tile_entries: S2PMEntries,
129    offset: u64,
130    addressed_tiles: u64,
131    clustered: bool,
132    compression: CompressionFormat,
133    writer: W,
134}
135impl<W: Writer> PMTilesWriter<W> {
136    /// given a compression scheme and a data writer, create an instance to start storing tiles
137    /// and metadata.
138    /// Compression will only describle how tiles are stored, nothing more.
139    pub fn new(writer: W, compression: CompressionFormat) -> Self {
140        let mut writer = PMTilesWriter {
141            tile_entries: PMDirectory::default(),
142            s2tile_entries: S2PMEntries::default(),
143            offset: 0,
144            addressed_tiles: 0,
145            clustered: false,
146            compression,
147            writer,
148        };
149        writer.writer.append(&vec![0u8; S2_PM_ROOT_SIZE]);
150        writer
151    }
152
153    /// take ownership of writer data (if local this actually has content)
154    pub fn take(&mut self) -> Vec<u8> {
155        self.writer.take()
156    }
157
158    /// Write a tile to the PMTiles file given its tile ID.
159    pub fn write_tile(&mut self, tile_id: u64, data: &[u8], face: Option<Face>) {
160        let length = data.len();
161        let tile_entries = match face {
162            None => &mut self.tile_entries,
163            Some(f) => self.s2tile_entries.get_mut(f),
164        };
165        if !tile_entries.is_empty() && tile_id < tile_entries.last().unwrap().tile_id {
166            self.clustered = false;
167        }
168
169        let offset = self.offset;
170        self.writer.append(data);
171        tile_entries.insert(PMEntry { tile_id, offset, length: length as u32, run_length: 1 });
172        self.offset += length as u64;
173
174        self.addressed_tiles += 1;
175    }
176
177    /// Finish writing by building the header with root and leaf directories
178    fn _commit(&mut self, metadata: &Metadata) {
179        if !self.tile_entries.is_empty() {
180            self.commit_wm(metadata);
181        } else {
182            self.commit_s2(metadata);
183        }
184
185        self.writer.flush();
186    }
187
188    /// Finish writing by building the header with root and leaf directories
189    fn commit_wm(&mut self, metadata: &Metadata) {
190        // build metadata
191        let meta_buffer = serde_json::to_vec(metadata).unwrap();
192
193        // optimize directories
194        let od: OptimizedDirectory = OptimizedDirectory::optimize_directories(
195            &mut self.tile_entries,
196            S2_PM_ROOT_SIZE - S2_PM_HEADER_SIZE_BYTES - meta_buffer.len(),
197        );
198        let OptimizedDirectory { root_bytes, leaves_bytes, .. } = od;
199
200        // build header data
201        let root_directory_offset = S2_PM_HEADER_SIZE_BYTES as u64;
202        let root_directory_length = root_bytes.len() as u64;
203        let metadata_offset = root_directory_offset + root_directory_length;
204        let metadata_length = meta_buffer.len() as u64;
205        let leaf_directory_offset = self.offset + S2_PM_ROOT_SIZE as u64;
206        let leaf_directory_length = leaves_bytes.len() as u64;
207        self.offset += leaves_bytes.len() as u64;
208
209        // write data
210        self.writer.append(&leaves_bytes);
211        // to make writing fasters
212        let min_zoom = PMTilePos::from_id(self.tile_entries.first().unwrap().tile_id).zoom;
213        let max_zoom = PMTilePos::from_id(self.tile_entries.last().unwrap().tile_id).zoom;
214
215        // build header
216        let header = PMHeader {
217            version: 3,
218            root_directory_offset,
219            root_directory_length,
220            metadata_offset,
221            metadata_length,
222            leaf_directory_offset,
223            leaf_directory_length,
224            data_offset: S2_PM_ROOT_SIZE as u64,
225            data_length: self.offset,
226            n_addressed_tiles: self.addressed_tiles,
227            n_tile_entries: self.tile_entries.len() as u64,
228            n_tile_contents: 0,
229            clustered: self.clustered,
230            internal_compression: CompressionFormat::None,
231            tile_compression: self.compression,
232            tile_type: PMTileType::Unknown,
233            min_zoom,
234            max_zoom,
235            ..Default::default()
236        };
237        let serialized_header = header.to_bytes().take();
238
239        // write header
240        self.writer.write(&serialized_header, 0);
241        self.writer.write(&root_bytes, root_directory_offset);
242        self.writer.write(&meta_buffer, metadata_offset);
243    }
244
245    /// Finish writing by building the header with root and leaf directories
246    fn commit_s2(&mut self, metadata: &Metadata) {
247        // build metadata
248        let meta_buffer = serde_json::to_vec(metadata).unwrap();
249
250        // optimize directories
251        let od = OptimizedDirectory::optimize_directories(
252            self.s2tile_entries.get_mut(Face::Face0),
253            S2_PM_ROOT_SIZE - S2_PM_HEADER_SIZE_BYTES - meta_buffer.len(),
254        );
255        let OptimizedDirectory { root_bytes, leaves_bytes, .. } = od;
256        let od1 = OptimizedDirectory::optimize_directories(
257            self.s2tile_entries.get_mut(Face::Face1),
258            S2_PM_ROOT_SIZE - S2_PM_HEADER_SIZE_BYTES - meta_buffer.len(),
259        );
260        let OptimizedDirectory { root_bytes: root_bytes1, leaves_bytes: leaves_bytes1, .. } = od1;
261        let od2 = OptimizedDirectory::optimize_directories(
262            self.s2tile_entries.get_mut(Face::Face2),
263            S2_PM_ROOT_SIZE - S2_PM_HEADER_SIZE_BYTES - meta_buffer.len(),
264        );
265        let OptimizedDirectory { root_bytes: root_bytes2, leaves_bytes: leaves_bytes2, .. } = od2;
266        let od3 = OptimizedDirectory::optimize_directories(
267            self.s2tile_entries.get_mut(Face::Face3),
268            S2_PM_ROOT_SIZE - S2_PM_HEADER_SIZE_BYTES - meta_buffer.len(),
269        );
270        let OptimizedDirectory { root_bytes: root_bytes3, leaves_bytes: leaves_bytes3, .. } = od3;
271        let od4 = OptimizedDirectory::optimize_directories(
272            self.s2tile_entries.get_mut(Face::Face4),
273            S2_PM_ROOT_SIZE - S2_PM_HEADER_SIZE_BYTES - meta_buffer.len(),
274        );
275        let OptimizedDirectory { root_bytes: root_bytes4, leaves_bytes: leaves_bytes4, .. } = od4;
276        let od5 = OptimizedDirectory::optimize_directories(
277            self.s2tile_entries.get_mut(Face::Face5),
278            S2_PM_ROOT_SIZE - S2_PM_HEADER_SIZE_BYTES - meta_buffer.len(),
279        );
280        let OptimizedDirectory { root_bytes: root_bytes5, leaves_bytes: leaves_bytes5, .. } = od5;
281
282        // build header data
283        // roots
284        let root_directory_offset = S2_PM_HEADER_SIZE_BYTES as u64;
285        let root_directory_length = root_bytes.len() as u64;
286        let root_directory_offset1 = root_directory_offset + root_directory_length;
287        let root_directory_length1 = root_bytes1.len() as u64;
288        let root_directory_offset2 = root_directory_offset1 + root_directory_length1;
289        let root_directory_length2 = root_bytes2.len() as u64;
290        let root_directory_offset3 = root_directory_offset2 + root_directory_length2;
291        let root_directory_length3 = root_bytes3.len() as u64;
292        let root_directory_offset4 = root_directory_offset3 + root_directory_length3;
293        let root_directory_length4 = root_bytes4.len() as u64;
294        let root_directory_offset5 = root_directory_offset4 + root_directory_length4;
295        let root_directory_length5 = root_bytes5.len() as u64;
296        // metadata
297        let metadata_offset = root_directory_offset5 + root_directory_length5;
298        let metadata_length = meta_buffer.len() as u64;
299        // leafs
300        let leaf_directory_offset = self.offset + S2_PM_ROOT_SIZE as u64;
301        let leaf_directory_length = leaves_bytes.len() as u64;
302        self.offset += leaf_directory_length;
303        self.writer.append(&leaves_bytes);
304        let leaf_directory_offset1 = self.offset + S2_PM_ROOT_SIZE as u64;
305        let leaf_directory_length1 = leaves_bytes1.len() as u64;
306        self.offset += leaf_directory_length1;
307        self.writer.append(&leaves_bytes1);
308        let leaf_directory_offset2 = self.offset + S2_PM_ROOT_SIZE as u64;
309        let leaf_directory_length2 = leaves_bytes2.len() as u64;
310        self.offset += leaf_directory_length2;
311        self.writer.append(&leaves_bytes2);
312        let leaf_directory_offset3 = self.offset + S2_PM_ROOT_SIZE as u64;
313        let leaf_directory_length3 = leaves_bytes3.len() as u64;
314        self.offset += leaf_directory_length3;
315        self.writer.append(&leaves_bytes3);
316        let leaf_directory_offset4 = self.offset + S2_PM_ROOT_SIZE as u64;
317        let leaf_directory_length4 = leaves_bytes4.len() as u64;
318        self.offset += leaf_directory_length4;
319        self.writer.append(&leaves_bytes4);
320        let leaf_directory_offset5 = self.offset + S2_PM_ROOT_SIZE as u64;
321        let leaf_directory_length5 = leaves_bytes5.len() as u64;
322        self.offset += leaf_directory_length5;
323        self.writer.append(&leaves_bytes5);
324
325        // write data
326        self.writer.append(&leaves_bytes);
327        // build header
328        let header = S2PMHeader {
329            is_s2: true,
330            version: 3,
331            root_directory_offset,
332            root_directory_length,
333            root_directory_offset1,
334            root_directory_length1,
335            root_directory_offset2,
336            root_directory_length2,
337            root_directory_offset3,
338            root_directory_length3,
339            root_directory_offset4,
340            root_directory_length4,
341            root_directory_offset5,
342            root_directory_length5,
343            metadata_offset,
344            metadata_length,
345            leaf_directory_offset,
346            leaf_directory_length,
347            leaf_directory_offset1,
348            leaf_directory_length1,
349            leaf_directory_offset2,
350            leaf_directory_length2,
351            leaf_directory_offset3,
352            leaf_directory_length3,
353            leaf_directory_offset4,
354            leaf_directory_length4,
355            leaf_directory_offset5,
356            leaf_directory_length5,
357            data_offset: S2_PM_ROOT_SIZE as u64,
358            data_length: self.offset,
359            n_addressed_tiles: self.addressed_tiles,
360            n_tile_entries: self.tile_entries.len() as u64,
361            n_tile_contents: 0,
362            clustered: self.clustered,
363            internal_compression: CompressionFormat::None,
364            tile_compression: self.compression,
365            tile_type: PMTileType::Unknown,
366            ..Default::default()
367        };
368        let serialized_header = header.to_bytes().take();
369
370        // write header
371        self.writer.write(&serialized_header, 0);
372        self.writer.write(&root_bytes, root_directory_offset);
373        self.writer.write(&root_bytes1, root_directory_offset1);
374        self.writer.write(&root_bytes2, root_directory_offset2);
375        self.writer.write(&root_bytes3, root_directory_offset3);
376        self.writer.write(&root_bytes4, root_directory_offset4);
377        self.writer.write(&root_bytes5, root_directory_offset5);
378        self.writer.write(&meta_buffer, metadata_offset);
379    }
380}
381
382impl<W: Writer> TileWriter for PMTilesWriter<W> {
383    fn write_tile_wm(&mut self, zoom: u8, x: u32, y: u32, data: Vec<u8>) {
384        let tile_id = PMTilePos::new(zoom, x as u64, y as u64).to_id();
385        self.write_tile(tile_id, &data, None);
386    }
387
388    fn write_tile_s2(&mut self, face: Face, zoom: u8, x: u32, y: u32, data: Vec<u8>) {
389        let tile_id = PMTilePos::new(zoom, x as u64, y as u64).to_id();
390        self.write_tile(tile_id, &data, Some(face));
391    }
392
393    fn commit(&mut self, metadata: Metadata, tile_compression: Option<CompressionFormat>) {
394        if let Some(tile_compression) = tile_compression {
395            self.compression = tile_compression;
396        }
397        self._commit(&metadata);
398    }
399}
gistools/writers/pmtiles/mod.rs

gistools/writers/pmtiles/
mod.rs
