geo_index/kdtree/
index.rs

1use std::marker::PhantomData;
2
3use bytemuck::cast_slice;
4
5use crate::error::{GeoIndexError, Result};
6use crate::indices::Indices;
7use crate::kdtree::constants::{KDBUSH_HEADER_SIZE, KDBUSH_MAGIC, KDBUSH_VERSION};
8use crate::r#type::IndexableNum;
9
10/// Common metadata to describe a KDTree
11///
12/// You can use the metadata to infer the total byte size of a tree given the provided criteria.
13/// See [`data_buffer_length`][Self::data_buffer_length].
14#[derive(Debug, Clone, Copy, PartialEq)]
15pub struct KDTreeMetadata<N: IndexableNum> {
16    node_size: u16,
17    num_items: u32,
18    phantom: PhantomData<N>,
19    pub(crate) indices_byte_size: usize,
20    pub(crate) pad_coords_byte_size: usize,
21    pub(crate) coords_byte_size: usize,
22}
23
24impl<N: IndexableNum> KDTreeMetadata<N> {
25    /// Construct a new [`KDTreeMetadata`] from a number of items and node size.
26    pub fn new(num_items: u32, node_size: u16) -> Self {
27        assert!((2..=65535).contains(&node_size));
28
29        let coords_byte_size = (num_items as usize) * 2 * N::BYTES_PER_ELEMENT;
30        let indices_bytes_per_element = if num_items < 65536 { 2 } else { 4 };
31        let indices_byte_size = (num_items as usize) * indices_bytes_per_element;
32        let pad_coords_byte_size = (8 - (indices_byte_size % 8)) % 8;
33
34        Self {
35            node_size,
36            num_items,
37            phantom: PhantomData,
38            indices_byte_size,
39            pad_coords_byte_size,
40            coords_byte_size,
41        }
42    }
43
44    /// Construct a new [`KDTreeMetadata`] from an existing byte slice conforming to the "kdbush
45    /// ABI", such as what [`KDTreeBuilder`] generates.
46    pub fn from_slice(data: &[u8]) -> Result<Self> {
47        if data[0] != KDBUSH_MAGIC {
48            return Err(GeoIndexError::General(
49                "Data not in Kdbush format.".to_string(),
50            ));
51        }
52
53        let version_and_type = data[1];
54        let version = version_and_type >> 4;
55        if version != KDBUSH_VERSION {
56            return Err(GeoIndexError::General(
57                format!("Got v{} data when expected v{}.", version, KDBUSH_VERSION).to_string(),
58            ));
59        }
60
61        let type_ = version_and_type & 0x0f;
62        if type_ != N::TYPE_INDEX {
63            return Err(GeoIndexError::General(
64                format!(
65                    "Got type {} data when expected type {}.",
66                    type_,
67                    N::TYPE_INDEX
68                )
69                .to_string(),
70            ));
71        }
72
73        let node_size: u16 = cast_slice(&data[2..4])[0];
74        let num_items: u32 = cast_slice(&data[4..8])[0];
75
76        let slf = Self::new(num_items, node_size);
77        if slf.data_buffer_length() != data.len() {
78            return Err(GeoIndexError::General(format!(
79                "Expected {} bytes but received byte slice with {} bytes",
80                slf.data_buffer_length(),
81                data.len()
82            )));
83        }
84
85        Ok(slf)
86    }
87
88    /// The maximum number of items per node.
89    pub fn node_size(&self) -> u16 {
90        self.node_size
91    }
92
93    /// The number of items indexed in the tree.
94    pub fn num_items(&self) -> u32 {
95        self.num_items
96    }
97
98    /// The number of bytes that a KDTree with this metadata would have.
99    ///
100    /// ```
101    /// use geo_index::kdtree::KDTreeMetadata;
102    ///
103    /// let metadata = KDTreeMetadata::<f64>::new(25000, 16);
104    /// assert_eq!(metadata.data_buffer_length(), 450_008);
105    /// ```
106    pub fn data_buffer_length(&self) -> usize {
107        KDBUSH_HEADER_SIZE
108            + self.coords_byte_size
109            + self.indices_byte_size
110            + self.pad_coords_byte_size
111    }
112
113    /// Access the slice of coordinates from the data buffer this metadata represents.
114    pub fn coords_slice<'a>(&self, data: &'a [u8]) -> &'a [N] {
115        let coords_byte_start =
116            KDBUSH_HEADER_SIZE + self.indices_byte_size + self.pad_coords_byte_size;
117        let coords_byte_end = KDBUSH_HEADER_SIZE
118            + self.indices_byte_size
119            + self.pad_coords_byte_size
120            + self.coords_byte_size;
121        cast_slice(&data[coords_byte_start..coords_byte_end])
122    }
123
124    /// Access the slice of indices from the data buffer this metadata represents.
125    pub fn indices_slice<'a>(&self, data: &'a [u8]) -> Indices<'a> {
126        let indices_buf = &data[KDBUSH_HEADER_SIZE..KDBUSH_HEADER_SIZE + self.indices_byte_size];
127
128        if self.num_items < 65536 {
129            Indices::U16(cast_slice(indices_buf))
130        } else {
131            Indices::U32(cast_slice(indices_buf))
132        }
133    }
134}
135
136/// An owned KDTree buffer.
137///
138/// Usually this will be created from scratch via [`KDTreeBuilder`][crate::kdtree::KDTreeBuilder].
139#[derive(Debug, Clone, PartialEq)]
140pub struct KDTree<N: IndexableNum> {
141    pub(crate) buffer: Vec<u8>,
142    pub(crate) metadata: KDTreeMetadata<N>,
143}
144
145impl<N: IndexableNum> KDTree<N> {
146    /// Consume this KDTree, returning the underlying buffer.
147    pub fn into_inner(self) -> Vec<u8> {
148        self.buffer
149    }
150}
151
152impl<N: IndexableNum> AsRef<[u8]> for KDTree<N> {
153    fn as_ref(&self) -> &[u8] {
154        &self.buffer
155    }
156}
157
158/// A reference on an external KDTree buffer.
159#[derive(Debug, Clone, PartialEq)]
160pub struct KDTreeRef<'a, N: IndexableNum> {
161    pub(crate) coords: &'a [N],
162    pub(crate) indices: Indices<'a>,
163    pub(crate) metadata: KDTreeMetadata<N>,
164}
165
166impl<'a, N: IndexableNum> KDTreeRef<'a, N> {
167    /// Construct a new KDTreeRef from an external byte slice.
168    ///
169    /// This byte slice must conform to the "kdbush ABI", that is, the ABI originally implemented
170    /// by the JavaScript [`kdbush` library](https://github.com/mourner/kdbush). You can extract
171    /// such a buffer either via [`KDTree::into_inner`] or from the `.data` attribute of the
172    /// JavaScript `KDBush` object.
173    pub fn try_new<T: AsRef<[u8]>>(data: &'a T) -> Result<Self> {
174        let data = data.as_ref();
175        let metadata = KDTreeMetadata::from_slice(data)?;
176        let coords = metadata.coords_slice(data);
177        let indices = metadata.indices_slice(data);
178
179        Ok(Self {
180            coords,
181            indices,
182            metadata,
183        })
184    }
185
186    /// Construct a new KDTreeRef without doing any validation
187    ///
188    /// # Safety
189    ///
190    /// `metadata` must be valid for this data buffer.
191    pub unsafe fn new_unchecked<T: AsRef<[u8]>>(
192        data: &'a T,
193        metadata: KDTreeMetadata<N>,
194    ) -> Result<Self> {
195        let data = data.as_ref();
196        let coords = metadata.coords_slice(data);
197        let indices = metadata.indices_slice(data);
198
199        Ok(Self {
200            coords,
201            indices,
202            metadata,
203        })
204    }
205}
geo_index/kdtree/index.rs

geo_index/kdtree/
index.rs
