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    /// Convert to f64 for distance calculations
24    fn to_f64(self) -> Option<f64> {
25        NumCast::from(self)
26    }
27
28    /// Convert from f64 for distance calculations
29    fn from_f64(value: f64) -> Option<Self> {
30        NumCast::from(value)
31    }
32
33    /// Get the square root of this value
34    fn sqrt(self) -> Option<Self> {
35        self.to_f64()
36            .and_then(|value| {
37                if value >= 0.0 {
38                    Some(value.sqrt())
39                } else {
40                    None
41                }
42            })
43            .and_then(NumCast::from)
44    }
45}
46
47impl IndexableNum for i8 {
48    const TYPE_INDEX: u8 = 0;
49    const BYTES_PER_ELEMENT: usize = 1;
50}
51
52impl IndexableNum for u8 {
53    const TYPE_INDEX: u8 = 1;
54    const BYTES_PER_ELEMENT: usize = 1;
55}
56
57impl IndexableNum for i16 {
58    const TYPE_INDEX: u8 = 3;
59    const BYTES_PER_ELEMENT: usize = 2;
60}
61
62impl IndexableNum for u16 {
63    const TYPE_INDEX: u8 = 4;
64    const BYTES_PER_ELEMENT: usize = 2;
65}
66
67impl IndexableNum for i32 {
68    const TYPE_INDEX: u8 = 5;
69    const BYTES_PER_ELEMENT: usize = 4;
70}
71
72impl IndexableNum for u32 {
73    const TYPE_INDEX: u8 = 6;
74    const BYTES_PER_ELEMENT: usize = 4;
75}
76
77impl IndexableNum for f32 {
78    const TYPE_INDEX: u8 = 7;
79    const BYTES_PER_ELEMENT: usize = 4;
80}
81
82impl IndexableNum for f64 {
83    const TYPE_INDEX: u8 = 8;
84    const BYTES_PER_ELEMENT: usize = 8;
85}
86
87/// For compatibility with JS, which contains a Uint8ClampedArray
88const U8_CLAMPED_TYPE_INDEX: u8 = 2;
89
90/// An enum over the allowed coordinate types in the spatial index.
91///
92/// This can be used to infer the coordinate type from an existing buffer.
93#[allow(missing_docs)]
94#[derive(Debug, Clone, Copy, PartialEq, Hash)]
95pub enum CoordType {
96    Int8,
97    UInt8,
98    Int16,
99    UInt16,
100    Int32,
101    UInt32,
102    Float32,
103    Float64,
104}
105
106impl CoordType {
107    /// Infer the CoordType from an existing buffer.
108    ///
109    /// This can be used to discern the generic type to use when constructing an `RTreeRef` or
110    /// `KDTreeRef`.
111    ///
112    /// ```
113    /// use geo_index::rtree::RTreeBuilder;
114    /// use geo_index::rtree::sort::HilbertSort;
115    /// use geo_index::CoordType;
116    ///
117    /// let mut builder = RTreeBuilder::<u32>::new(2);
118    /// builder.add(0, 0, 2, 2);
119    /// builder.add(1, 1, 3, 3);
120    /// let tree = builder.finish::<HilbertSort>();
121    ///
122    /// let coord_type = CoordType::from_buffer(&tree).unwrap();
123    /// assert!(matches!(coord_type, CoordType::UInt32));
124    /// ```
125    ///
126    /// This method works for both buffers representing RTree or KDTree trees.
127    pub fn from_buffer<T: AsRef<[u8]>>(data: &T) -> Result<Self, GeoIndexError> {
128        let data = data.as_ref();
129        let magic = data[0];
130        if magic != 0xfb && magic != KDBUSH_MAGIC {
131            return Err(GeoIndexError::General(
132                "Data not in Flatbush or Kdbush format.".to_string(),
133            ));
134        }
135
136        let version_and_type = data[1];
137        let type_ = version_and_type & 0x0f;
138        let result = match type_ {
139            i8::TYPE_INDEX => CoordType::Int8,
140            u8::TYPE_INDEX => CoordType::UInt8,
141            U8_CLAMPED_TYPE_INDEX => CoordType::UInt8,
142            i16::TYPE_INDEX => CoordType::Int16,
143            u16::TYPE_INDEX => CoordType::UInt16,
144            i32::TYPE_INDEX => CoordType::Int32,
145            u32::TYPE_INDEX => CoordType::UInt32,
146            f32::TYPE_INDEX => CoordType::Float32,
147            f64::TYPE_INDEX => CoordType::Float64,
148            t => return Err(GeoIndexError::General(format!("Unexpected type {}.", t))),
149        };
150        Ok(result)
151    }
152}
153
154// https://rust-lang.github.io/api-guidelines/future-proofing.html#sealed-traits-protect-against-downstream-implementations-c-sealed
155mod private {
156    pub trait Sealed {}
157
158    impl Sealed for i8 {}
159    impl Sealed for u8 {}
160    impl Sealed for i16 {}
161    impl Sealed for u16 {}
162    impl Sealed for i32 {}
163    impl Sealed for u32 {}
164    impl Sealed for f32 {}
165    impl Sealed for f64 {}
166}
167
168/// A single coordinate.
169///
170/// Used in the implementation of RectTrait for Node.
171pub struct Coord<N: IndexableNum> {
172    pub(crate) x: N,
173    pub(crate) y: N,
174}
175
176impl<N: IndexableNum> CoordTrait for Coord<N> {
177    type T = N;
178
179    fn dim(&self) -> geo_traits::Dimensions {
180        geo_traits::Dimensions::Xy
181    }
182
183    fn x(&self) -> Self::T {
184        self.x
185    }
186
187    fn y(&self) -> Self::T {
188        self.y
189    }
190
191    fn nth_or_panic(&self, n: usize) -> Self::T {
192        match n {
193            0 => self.x,
194            1 => self.y,
195            _ => panic!("Invalid index of coord"),
196        }
197    }
198}
geo_index/type.rs

geo_index/
type.rs
