gistools/readers/pmtiles/
reader.rs

1use super::{
2    PMDirectory, PMTilePos, S2_PM_HEADER_SIZE_BYTES, S2_PM_ROOT_SIZE, S2PMEntries, S2PMHeader,
3    find_tile,
4};
5use crate::{
6    data_structures::Cache,
7    parsers::{Buffer, Reader},
8    util::decompress_data,
9};
10use alloc::{string::String, vec::Vec};
11use s2_tilejson::{Metadata, UnknownMetadata};
12use s2json::Face;
13
14/// # (S2) PMTiles Reader
15///
16/// ## Description
17/// A V3.0 PMTiles reader for reading standard WebMercator Tile data and V1.0 S2 Tile data.
18///
19/// A Modified implementation of the PMTiles library. It is backwards compatible but
20/// offers support for the S2 Projection.
21///
22/// You can learn more about the [S2PMTiles Specification here](https://github.com/Open-S2/s2-pmtiles/blob/master/s2-pmtiles-spec/1.0.0/README.md).
23///
24/// ## Usage
25///
26/// PMTilesReader utilizes any struct that implements the [`Reader`] trait.
27/// Options are [`crate::parsers::BufferReader`], [`crate::parsers::FileReader`], [`crate::parsers::MMapReader`], and [`crate::parsers::FetchReader`].
28///
29/// The methods you have access to:
30/// - [`PMTilesReader::new`]: Create a new PMTilesReader
31/// - [`PMTilesReader::get_header`]: Get the PMTiles header
32/// - [`PMTilesReader::get_s2_metadata`]: Get the S2 PMTiles metadata
33/// - [`PMTilesReader::get_metadata`]: Get the PMTiles metadata
34/// - [`PMTilesReader::get_tile_s2`]: Get an S2 Tile
35/// - [`PMTilesReader::get_tile_wm`]: Get an WM Tile
36/// - [`PMTilesReader::get_tile`]: Get a Tile irregardless of the projection type
37///
38/// ```rust
39/// use gistools::{parsers::FileReader, readers::PMTilesReader};
40/// use std::path::PathBuf;
41///
42/// let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
43///     .join("tests/readers/pmtiles/fixtures/test_fixture_1.pmtiles");
44/// let file_reader = FileReader::new(path).unwrap();
45/// let mut reader = PMTilesReader::new(file_reader, None);
46///
47/// smol::block_on(async {
48///
49/// // pull out the header
50/// let header = reader.get_header().await;
51///
52/// // get the metadata
53/// let metadata = reader.get_metadata();
54///
55/// // S2 specific functions
56/// let tile = reader.get_tile_s2(0.into(), 0, 0, 0).await;
57///
58/// // WM functions
59/// let tile = reader.get_tile_wm(0, 0, 0).await.unwrap();
60///
61/// });
62/// ```
63///
64/// ## Links
65/// - <https://github.com/Open-S2/s2-pmtiles>
66/// - <https://github.com/Open-S2/s2-pmtiles/blob/master/s2-pmtiles-spec/1.0.0/README.md>
67/// - <https://github.com/protomaps/PMTiles>
68/// - <https://github.com/protomaps/PMTiles/blob/main/spec/v3/spec.md>
69#[derive(Debug)]
70pub struct PMTilesReader<R: Reader> {
71    header: Option<S2PMHeader>,
72    root_dir: PMDirectory,
73    root_dir_s2: S2PMEntries,
74    metadata: Metadata,
75    dir_cache: Cache<u64, PMDirectory>,
76    reader: R,
77}
78impl<R: Reader> PMTilesReader<R> {
79    /// Given an input path, read in the header and root directory
80    pub fn new(reader: R, max_size: Option<usize>) -> Self {
81        let max_size = max_size.unwrap_or(20);
82        Self {
83            header: None,
84            root_dir: PMDirectory::default(),
85            root_dir_s2: S2PMEntries::default(),
86            metadata: Metadata::default(),
87            dir_cache: Cache::new(max_size, None),
88            reader,
89        }
90    }
91
92    /// fetch the s2 metadata as needed
93    pub async fn get_header(&mut self) -> S2PMHeader {
94        if self.header.is_some() {
95            return self.header.unwrap();
96        }
97
98        let data = self.get_range(0, S2_PM_ROOT_SIZE as u64).await;
99        let header_data = data[0..S2_PM_HEADER_SIZE_BYTES].to_vec();
100        // header
101        let mut header = S2PMHeader::from_bytes(&mut header_data.into());
102
103        // json metadata
104        let json_offset = header.metadata_offset as usize;
105        let json_length = header.metadata_length as usize;
106        let json_metadata = decompress_data(
107            &data[json_offset..(json_offset + json_length)],
108            header.internal_compression,
109        )
110        .unwrap();
111        let meta: UnknownMetadata = serde_json::from_str(&String::from_utf8_lossy(&json_metadata))
112            .unwrap_or_else(|e| panic!("ERROR: {}", e));
113        self.metadata = meta.to_metadata();
114
115        // root directory data
116        let root_dir_offset = header.root_directory_offset as usize;
117        let root_dir_length = header.root_directory_length as usize;
118        let root_dir_data = decompress_data(
119            &data[root_dir_offset..(root_dir_offset + root_dir_length)],
120            header.internal_compression,
121        )
122        .unwrap();
123        self.root_dir = PMDirectory::from_buffer(&mut root_dir_data.into());
124
125        if header.is_s2 {
126            self.get_s2_metadata(&data, &mut header);
127        }
128
129        self.header = Some(header);
130
131        header
132    }
133
134    /// If S2, we need to build the other face's root directories
135    pub fn get_s2_metadata(&mut self, data: &[u8], header: &mut S2PMHeader) {
136        // move the root directory to the s2 root
137        self.root_dir_s2.face_0 = self.root_dir.clone();
138        // add the 5 other faces
139        for face in [Face::Face1, Face::Face2, Face::Face3, Face::Face4, Face::Face5] {
140            let root_offset = header.get_root_offset(face) as usize;
141            let root_length = header.get_root_length(face) as usize;
142            let face_dir_data = decompress_data(
143                &data[root_offset..(root_offset + root_length)],
144                header.internal_compression,
145            )
146            .unwrap();
147            self.root_dir_s2.set_dir(face, PMDirectory::from_buffer(&mut face_dir_data.into()));
148        }
149    }
150
151    /// get the metadata
152    pub fn get_metadata(&mut self) -> &Metadata {
153        &self.metadata
154    }
155
156    /// get an S2 tile
157    pub async fn get_tile_s2(&mut self, face: Face, zoom: u8, x: u64, y: u64) -> Option<Vec<u8>> {
158        self.get_tile(Some(face), zoom, x, y).await
159    }
160
161    /// get an WM tile
162    pub async fn get_tile_wm(&mut self, zoom: u8, x: u64, y: u64) -> Option<Vec<u8>> {
163        self.get_tile(None, zoom, x, y).await
164    }
165
166    /// get a tile, wheather WM or S2
167    pub async fn get_tile(
168        &mut self,
169        face: Option<Face>,
170        zoom: u8,
171        x: u64,
172        y: u64,
173    ) -> Option<Vec<u8>> {
174        let header = self.get_header().await;
175        let tile_id = PMTilePos::new(zoom, x, y).to_id();
176        // if zoom < header.min_zoom || zoom > header.max_zoom { return None; }
177
178        let mut d_o = header.root_directory_offset;
179        let mut d_l = header.root_directory_length;
180
181        for _ in 0..4 {
182            let directory = self.get_directory(d_o, d_l, face).await;
183            if directory.is_empty() {
184                return None;
185            }
186            let entry = find_tile(&directory.entries, tile_id);
187            match entry {
188                None => {
189                    return None;
190                }
191                Some(entry) => {
192                    if entry.run_length > 0 {
193                        let entry_data = self
194                            .get_range(header.data_offset + entry.offset, entry.length as u64)
195                            .await;
196                        return Some(
197                            decompress_data(&entry_data, header.internal_compression).unwrap(),
198                        );
199                    } else {
200                        d_o = header.leaf_directory_offset + entry.offset;
201                        d_l = entry.length as u64;
202                    }
203                }
204            }
205        }
206
207        panic!("Maximum directory depth exceeded");
208    }
209
210    /// Get a full directory
211    async fn get_directory(&mut self, offset: u64, length: u64, face: Option<Face>) -> PMDirectory {
212        let dir = match face {
213            None => &self.root_dir,
214            Some(f) => self.root_dir_s2.get(f),
215        };
216        let internal_compression = self.header.unwrap().internal_compression;
217        let root_directory_offset = self.header.unwrap().root_directory_offset;
218        // if root_directory_offset, return roon
219        if offset == root_directory_offset {
220            return dir.clone();
221        }
222        // check cache
223        if let Some(cache) = self.dir_cache.get(&offset) {
224            cache.clone()
225        } else {
226            // get from archive
227            let resp = self.get_range(offset, length).await;
228            let data = decompress_data(&resp, internal_compression).unwrap();
229            let mut buffer: Buffer = Buffer::new(data);
230            let directory = PMDirectory::from_buffer(&mut buffer);
231            if directory.is_empty() {
232                panic!("Empty directory is invalid");
233            }
234            // save in cache
235            self.dir_cache.set(offset, directory.clone());
236
237            directory
238        }
239    }
240
241    /// Get a range of bytes given an offset and length
242    async fn get_range(&mut self, offset: u64, length: u64) -> Vec<u8> {
243        let len = self.reader.len();
244        if len != 0 {
245            // This is not a FetchReader
246            let end = u64::min(len, offset + length);
247            self.reader.slice(Some(offset), Some(end))
248        } else {
249            self.reader.get_slice(offset, Some(length)).await
250        }
251    }
252}
gistools/readers/pmtiles/reader.rs

gistools/readers/pmtiles/
reader.rs
