geo_index/
type.rs

1use std::fmt::Debug;
2
3use geo_traits::CoordTrait;
4use num_traits::{Bounded, Num, NumCast};
5
6use crate::kdtree::constants::KDBUSH_MAGIC;
7use crate::GeoIndexError;
8
9/// A trait for types that can be used for indexed coordinates.
10///
11/// This trait is sealed and cannot be implemented for external types. This is because we want to
12/// ensure FFI compatibility with other implementations, including the reference implementations in
13/// JavaScript ([rtree](https://github.com/mourner/flatbush),
14/// [kdtree](https://github.com/mourner/kdbush))
15pub trait IndexableNum:
16    private::Sealed + Num + NumCast + PartialOrd + Debug + Send + Sync + bytemuck::Pod + Bounded
17{
18    /// The type index to match the array order of `ARRAY_TYPES` in flatbush JS
19    const TYPE_INDEX: u8;
20    /// The number of bytes per element
21    const BYTES_PER_ELEMENT: usize;
22}
23
24impl IndexableNum for i8 {
25    const TYPE_INDEX: u8 = 0;
26    const BYTES_PER_ELEMENT: usize = 1;
27}
28
29impl IndexableNum for u8 {
30    const TYPE_INDEX: u8 = 1;
31    const BYTES_PER_ELEMENT: usize = 1;
32}
33
34impl IndexableNum for i16 {
35    const TYPE_INDEX: u8 = 3;
36    const BYTES_PER_ELEMENT: usize = 2;
37}
38
39impl IndexableNum for u16 {
40    const TYPE_INDEX: u8 = 4;
41    const BYTES_PER_ELEMENT: usize = 2;
42}
43
44impl IndexableNum for i32 {
45    const TYPE_INDEX: u8 = 5;
46    const BYTES_PER_ELEMENT: usize = 4;
47}
48
49impl IndexableNum for u32 {
50    const TYPE_INDEX: u8 = 6;
51    const BYTES_PER_ELEMENT: usize = 4;
52}
53
54impl IndexableNum for f32 {
55    const TYPE_INDEX: u8 = 7;
56    const BYTES_PER_ELEMENT: usize = 4;
57}
58
59impl IndexableNum for f64 {
60    const TYPE_INDEX: u8 = 8;
61    const BYTES_PER_ELEMENT: usize = 8;
62}
63
64/// For compatibility with JS, which contains a Uint8ClampedArray
65const U8_CLAMPED_TYPE_INDEX: u8 = 2;
66
67/// An enum over the allowed coordinate types in the spatial index.
68///
69/// This can be used to infer the coordinate type from an existing buffer.
70#[allow(missing_docs)]
71#[derive(Debug, Clone, Copy, PartialEq, Hash)]
72pub enum CoordType {
73    Int8,
74    UInt8,
75    Int16,
76    UInt16,
77    Int32,
78    UInt32,
79    Float32,
80    Float64,
81}
82
83impl CoordType {
84    /// Infer the CoordType from an existing buffer.
85    ///
86    /// This can be used to discern the generic type to use when constructing an `RTreeRef` or
87    /// `KDTreeRef`.
88    ///
89    /// ```
90    /// use geo_index::rtree::RTreeBuilder;
91    /// use geo_index::rtree::sort::HilbertSort;
92    /// use geo_index::CoordType;
93    ///
94    /// let mut builder = RTreeBuilder::<u32>::new(2);
95    /// builder.add(0, 0, 2, 2);
96    /// builder.add(1, 1, 3, 3);
97    /// let tree = builder.finish::<HilbertSort>();
98    ///
99    /// let coord_type = CoordType::from_buffer(&tree).unwrap();
100    /// assert!(matches!(coord_type, CoordType::UInt32));
101    /// ```
102    ///
103    /// This method works for both buffers representing RTree or KDTree trees.
104    pub fn from_buffer<T: AsRef<[u8]>>(data: &T) -> Result<Self, GeoIndexError> {
105        let data = data.as_ref();
106        let magic = data[0];
107        if magic != 0xfb && magic != KDBUSH_MAGIC {
108            return Err(GeoIndexError::General(
109                "Data not in Flatbush or Kdbush format.".to_string(),
110            ));
111        }
112
113        let version_and_type = data[1];
114        let type_ = version_and_type & 0x0f;
115        let result = match type_ {
116            i8::TYPE_INDEX => CoordType::Int8,
117            u8::TYPE_INDEX => CoordType::UInt8,
118            U8_CLAMPED_TYPE_INDEX => CoordType::UInt8,
119            i16::TYPE_INDEX => CoordType::Int16,
120            u16::TYPE_INDEX => CoordType::UInt16,
121            i32::TYPE_INDEX => CoordType::Int32,
122            u32::TYPE_INDEX => CoordType::UInt32,
123            f32::TYPE_INDEX => CoordType::Float32,
124            f64::TYPE_INDEX => CoordType::Float64,
125            t => return Err(GeoIndexError::General(format!("Unexpected type {}.", t))),
126        };
127        Ok(result)
128    }
129}
130
131// https://rust-lang.github.io/api-guidelines/future-proofing.html#sealed-traits-protect-against-downstream-implementations-c-sealed
132mod private {
133    pub trait Sealed {}
134
135    impl Sealed for i8 {}
136    impl Sealed for u8 {}
137    impl Sealed for i16 {}
138    impl Sealed for u16 {}
139    impl Sealed for i32 {}
140    impl Sealed for u32 {}
141    impl Sealed for f32 {}
142    impl Sealed for f64 {}
143}
144
145/// A single coordinate.
146///
147/// Used in the implementation of RectTrait for Node.
148pub struct Coord<N: IndexableNum> {
149    pub(crate) x: N,
150    pub(crate) y: N,
151}
152
153impl<N: IndexableNum> CoordTrait for Coord<N> {
154    type T = N;
155
156    fn dim(&self) -> geo_traits::Dimensions {
157        geo_traits::Dimensions::Xy
158    }
159
160    fn x(&self) -> Self::T {
161        self.x
162    }
163
164    fn y(&self) -> Self::T {
165        self.y
166    }
167
168    fn nth_or_panic(&self, n: usize) -> Self::T {
169        match n {
170            0 => self.x,
171            1 => self.y,
172            _ => panic!("Invalid index of coord"),
173        }
174    }
175}
geo_index/type.rs

geo_index/
type.rs
