Infer coord type from buffer

benches/rtree.rs

@@ -1,9 +1,9 @@
 use bytemuck::cast_slice;
 use criterion::{criterion_group, criterion_main, Criterion};
-use geo_index::r#type::IndexableNum;
 use geo_index::rtree::sort::{HilbertSort, STRSort};
 use geo_index::rtree::util::f64_box_to_f32;
 use geo_index::rtree::{OwnedRTree, RTreeBuilder, RTreeIndex};
+use geo_index::IndexableNum;
 use rstar::primitives::{GeomWithData, Rectangle};
 use rstar::{RTree, AABB};
 use std::fs::read;

src/kdtree/index.rs

@@ -42,6 +42,7 @@ pub struct KDTreeRef<'a, N: IndexableNum> {
 impl<'a, N: IndexableNum> KDTreeRef<'a, N> {
     pub fn try_new<T: AsRef<[u8]>>(data: &'a T) -> Result<Self, GeoIndexError> {
         let data = data.as_ref();
+        // TODO: validate length of slice?
 
         if data[0] != KDBUSH_MAGIC {
             return Err(GeoIndexError::General(
@@ -57,6 +58,18 @@ impl<'a, N: IndexableNum> KDTreeRef<'a, N> {
             ));
         }
 
+        let type_ = version_and_type & 0x0f;
+        if type_ != N::TYPE_INDEX {
+            return Err(GeoIndexError::General(
+                format!(
+                    "Got type {} data when expected type {}.",
+                    type_,
+                    N::TYPE_INDEX
+                )
+                .to_string(),
+            ));
+        }
+
         let node_size: u16 = cast_slice(&data[2..4])[0];
         let num_items: u32 = cast_slice(&data[4..8])[0];
         let node_size = node_size as usize;

src/kdtree/mod.rs

@@ -1,5 +1,5 @@
 mod builder;
-mod constants;
+pub(crate) mod constants;
 mod index;
 mod r#trait;
 

src/lib.rs

@@ -2,9 +2,10 @@ mod error;
 pub mod indices;
 pub mod kdtree;
 pub mod rtree;
-pub mod r#type;
+mod r#type;
 
 pub use error::GeoIndexError;
+pub use r#type::{CoordType, IndexableNum};
 
 #[cfg(test)]
 pub(crate) mod test;

src/rtree/index.rs

@@ -72,6 +72,18 @@ impl<'a, N: IndexableNum> RTreeRef<'a, N> {
             ));
         }
 
+        let type_ = version_and_type & 0x0f;
+        if type_ != N::TYPE_INDEX {
+            return Err(GeoIndexError::General(
+                format!(
+                    "Got type {} data when expected type {}.",
+                    type_,
+                    N::TYPE_INDEX
+                )
+                .to_string(),
+            ));
+        }
+
         let node_size: u16 = cast_slice(&data[2..4])[0];
         let num_items: u32 = cast_slice(&data[4..8])[0];
         let node_size = node_size as usize;

src/type.rs

@@ -2,6 +2,9 @@ use std::fmt::Debug;
 
 use num_traits::{Bounded, Num, NumCast, ToPrimitive};
 
+use crate::kdtree::constants::KDBUSH_MAGIC;
+use crate::GeoIndexError;
+
 /// A trait for types that can be used for indexed coordinates.
 ///
 /// This trait is sealed and cannot be implemented for external types. This is because we want to
@@ -65,6 +68,50 @@ impl IndexableNum for f64 {
     const TYPE_INDEX: u8 = 8;
     const BYTES_PER_ELEMENT: usize = 8;
 }
+
+/// An enum over the allowed coordinate types in the spatial index.
+pub enum CoordType {
+    Int8,
+    UInt8,
+    Int16,
+    UInt16,
+    Int32,
+    UInt32,
+    Float32,
+    Float64,
+}
+
+impl CoordType {
+    /// Infer the CoordType from an existing buffer.
+    ///
+    /// This can be used to discern the generic type to use when constructing an OwnedRTree or
+    /// OwnedKDTree.
+    pub fn from_buffer<T: AsRef<[u8]>>(data: &T) -> Result<Self, GeoIndexError> {
+        let data = data.as_ref();
+        let magic = data[0];
+        if magic != 0xfb && magic != KDBUSH_MAGIC {
+            return Err(GeoIndexError::General(
+                "Data not in Flatbush or Kdbush format.".to_string(),
+            ));
+        }
+
+        let version_and_type = data[1];
+        let type_ = version_and_type & 0x0f;
+        let result = match type_ {
+            i8::TYPE_INDEX => CoordType::Int8,
+            u8::TYPE_INDEX => CoordType::UInt8,
+            i16::TYPE_INDEX => CoordType::Int16,
+            u16::TYPE_INDEX => CoordType::UInt16,
+            i32::TYPE_INDEX => CoordType::Int32,
+            u32::TYPE_INDEX => CoordType::UInt32,
+            f32::TYPE_INDEX => CoordType::Float32,
+            f64::TYPE_INDEX => CoordType::Float64,
+            t => return Err(GeoIndexError::General(format!("Unexpected type {}.", t))),
+        };
+        Ok(result)
+    }
+}
+
 // https://rust-lang.github.io/api-guidelines/future-proofing.html#sealed-traits-protect-against-downstream-implementations-c-sealed
 mod private {
     pub trait Sealed {}