geo_index/rtree/
index.rs

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

geo_index/rtree/
index.rs
