feat(wip): rkyv v0.8

Cargo.toml

@@ -95,6 +95,11 @@ optional = true
 default-features = false
 features = ["alloc", "copy_unsafe", "size_64"]
 
+[dependencies.rkyv_08]
+package = "rkyv"
+version = "0.8.9"
+optional = true
+
 [dependencies.serde]
 version = "1"
 optional = true
@@ -127,6 +132,7 @@ las = ["dep:las"]
 serde = ["dep:serde", "serde/derive", "dep:serde_derive", "dep:serde_with", "fixed/serde", "aligned-vec/serde"]
 simd = []
 rkyv = ["dep:rkyv"]
+rkyv_08 = ["dep:rkyv_08"]
 test_utils = ["dep:rand", "dep:rand_chacha", "dep:rayon"]
 tracing = ["dep:tracing", "dep:tracing-subscriber"]
 
@@ -215,6 +221,16 @@ name = "immutable-rkyv-deserialize"
 path = "examples/immutable-rkyv-deserialize.rs"
 required-features = ["rkyv"]
 
+[[example]]
+name = "immutable-rkyv_08-serialize"
+path = "examples/immutable-rkyv_08-serialize.rs"
+required-features = ["rkyv_08"]
+
+[[example]]
+name = "immutable-rkyv_08-deserialize"
+path = "examples/immutable-rkyv_08-deserialize.rs"
+required-features = ["rkyv_08"]
+
 [[example]]
 name = "pointcloud-las"
 path = "examples/pointcloud-las.rs"

examples/immutable-rkyv_08-deserialize.rs

@@ -0,0 +1,72 @@
+use elapsed::ElapsedDuration;
+use memmap::MmapOptions;
+use std::error::Error;
+use std::fs::File;
+use std::time::Instant;
+#[cfg(feature = "tracing")]
+use tracing::Level;
+#[cfg(feature = "tracing")]
+use tracing_subscriber::fmt;
+
+use rkyv_08::{from_bytes_unchecked, rancor::Error as RkyvError};
+
+// use kiddo::immutable::float::kdtree::ArchivedImmutableKdTree;
+use kiddo::immutable::float::kdtree::ImmutableKdTree;
+use kiddo::SquaredEuclidean;
+
+type Tree = ImmutableKdTree<f64, u32, 3, 256>;
+
+fn main() -> Result<(), Box<dyn Error>>
+where
+{
+    #[cfg(feature = "tracing")]
+    let subscriber = fmt().with_max_level(Level::TRACE).without_time().finish();
+    #[cfg(feature = "tracing")]
+    tracing::subscriber::set_global_default(subscriber)?;
+
+    let query = [0.123f64, 0.456f64, 0.789f64];
+
+    let start = Instant::now();
+
+    // memmap the file into a buffer
+    let buf =
+        unsafe { MmapOptions::new().map(&File::open("./examples/immutable-test-tree-r08.rkyv")?)? };
+
+    // TODO: unsatisfied trait bounds when trying to call nearest_one on archived tree
+
+    // You can use the safe API for fast zero-copy deserialization
+    // let archived_tree = rkyv_08::access::<ArchivedImmutableKdTree<f64, u32, 3, 256>, RkyvError>(&buf[..]).unwrap();
+
+    // zero-copy deserialization into archived type
+    // let archived_tree =
+    //     unsafe { rkyv_08::access_unchecked::<ArchivedImmutableKdTree<f64, u32, 3, 256>>(&buf) };
+
+    // perform a query
+    // let nearest_neighbour = archived_tree.nearest_one::<SquaredEuclidean>(&query);
+
+    // println!(
+    //     "total elapsed: {}\n\n",
+    //     ElapsedDuration::new(start.elapsed())
+    // );
+    // println!(
+    //     "Nearest item to query (archived): {:?}",
+    //     nearest_neighbour.item
+    // );
+
+    // full deserialization
+    let tree = unsafe { from_bytes_unchecked::<Tree, RkyvError>(&buf) }?;
+
+    // perform a query
+    let nearest_neighbour = tree.nearest_one::<SquaredEuclidean>(&query);
+
+    println!(
+        "Nearest item to query (deserialized): {:?}",
+        nearest_neighbour.item
+    );
+    println!(
+        "total elapsed: {}\n\n",
+        ElapsedDuration::new(start.elapsed())
+    );
+
+    Ok(())
+}

examples/immutable-rkyv_08-serialize.rs

@@ -0,0 +1,63 @@
+use elapsed::ElapsedDuration;
+// use memmap::MmapOptions;
+use rand::Rng;
+use rand_chacha::rand_core::SeedableRng;
+use rkyv_08::{rancor::Error as RkyvError, to_bytes};
+use std::error::Error;
+use std::fs::File;
+use std::io::Write;
+use std::time::Instant;
+#[cfg(feature = "tracing")]
+use tracing::Level;
+#[cfg(feature = "tracing")]
+use tracing_subscriber::fmt;
+use ubyte::ToByteUnit;
+
+use kiddo::float::distance::SquaredEuclidean;
+use kiddo::immutable::float::kdtree::ImmutableKdTree;
+
+const NUM_ITEMS: usize = 50_000_000;
+
+type Tree = ImmutableKdTree<f64, u32, 3, 256>;
+
+fn main() -> Result<(), Box<dyn Error>> {
+    #[cfg(feature = "tracing")]
+    let subscriber = fmt().with_max_level(Level::TRACE).without_time().finish();
+    #[cfg(feature = "tracing")]
+    tracing::subscriber::set_global_default(subscriber)?;
+
+    let query = [0.123f64, 0.456f64, 0.789f64];
+
+    // build and serialize a large ImmutableKdTree
+    let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(1);
+    let content_to_add: Vec<[f64; 3]> = (0..NUM_ITEMS).map(|_| rng.gen::<[f64; 3]>()).collect();
+
+    let start = Instant::now();
+    let tree: Tree = ImmutableKdTree::new_from_slice(&content_to_add);
+    println!(
+        "Populated ImmutableKdTree with {} items ({})",
+        tree.size(),
+        ElapsedDuration::new(start.elapsed())
+    );
+
+    let nearest_neighbour = tree.nearest_one::<SquaredEuclidean>(&query);
+
+    println!("Nearest item to query: {:?}", nearest_neighbour.item);
+
+    let start = Instant::now();
+
+    let buf = to_bytes::<RkyvError>(&tree)?;
+
+    let mut file = File::create("./examples/immutable-test-tree-r08.rkyv")?;
+    file.write_all(&buf)
+        .expect("Could not write serialized rkyv to file");
+
+    let file_size = file.metadata().unwrap().len().bytes();
+    println!(
+        "Serialized k-d tree to rkyv file 'immutable-test-tree-r08.rkyv' ({}). File size: {:.2}",
+        ElapsedDuration::new(start.elapsed()),
+        file_size
+    );
+
+    Ok(())
+}

src/immutable/common/generate_immutable_nearest_one.rs

@@ -68,7 +68,7 @@ macro_rules! generate_immutable_nearest_one {
                 use cmov::Cmov;
                 use $crate::modified_van_emde_boas::modified_van_emde_boas_get_child_idx_v2_branchless;
 
-                if level > self.max_stem_level || self.stems.is_empty() {
+                if level > Into::<i32>::into(self.max_stem_level) || self.stems.is_empty() {
                     self.search_leaf_for_nearest_one::<D>(query, nearest, leaf_idx as usize);
                     return;
                 }
@@ -140,7 +140,7 @@ macro_rules! generate_immutable_nearest_one {
                 where
                     D: DistanceMetric<A, K>,
             {
-                if level > self.max_stem_level as usize || self.stems.is_empty() {
+                if level as i32 > Into::<i32>::into(self.max_stem_level) || self.stems.is_empty() {
                     self.search_leaf_for_nearest_one::<D>(query, nearest, leaf_idx as usize);
                     return;
                 }

src/immutable/float/kdtree.rs

@@ -11,12 +11,9 @@
 //! As with the vanilla tree, [`f64`] or [`f32`] are supported currently for co-ordinate
 //! values, or [`f16`](https://docs.rs/half/latest/half/struct.f16.html) if the `f16` feature is enabled
 
-pub use crate::float::kdtree::Axis;
-use crate::float_leaf_slice::leaf_slice::{LeafSlice, LeafSliceFloat, LeafSliceFloatChunk};
-#[cfg(feature = "modified_van_emde_boas")]
-use crate::modified_van_emde_boas::modified_van_emde_boas_get_child_idx_v2_branchless;
-use crate::traits::Content;
-#[cfg(feature = "rkyv")]
+use std::{cmp::PartialEq, fmt::Debug};
+
+#[cfg(feature = "rkyv_08")]
 use aligned_vec::CACHELINE_ALIGN;
 use aligned_vec::{avec, AVec, ConstAlign};
 use array_init::array_init;
@@ -27,8 +24,14 @@ use ordered_float::OrderedFloat;
 use rkyv::vec::ArchivedVec;
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
-use std::cmp::PartialEq;
-use std::fmt::Debug;
+
+pub use crate::float::kdtree::Axis;
+use crate::float_leaf_slice::leaf_slice::{LeafSlice, LeafSliceFloat, LeafSliceFloatChunk};
+#[cfg(feature = "rkyv_08")]
+use crate::immutable::float::rkyv_aligned_vec::EncodeAVec;
+#[cfg(feature = "modified_van_emde_boas")]
+use crate::modified_van_emde_boas::modified_van_emde_boas_get_child_idx_v2_branchless;
+use crate::traits::Content;
 
 /// Immutable floating point k-d tree
 ///
@@ -47,8 +50,14 @@ use std::fmt::Debug;
 ///
 /// A convenient type alias exists for ImmutableKdTree with some sensible defaults set: [`kiddo::ImmutableKdTree`](`crate::ImmutableKdTree`).
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(
+    feature = "rkyv_08",
+    derive(rkyv_08::Archive, rkyv_08::Serialize, rkyv_08::Deserialize)
+)]
+#[cfg_attr(feature = "rkyv_08", rkyv(crate=rkyv_08))]
 #[derive(Clone, Debug, PartialEq)]
 pub struct ImmutableKdTree<A: Copy + Default, T: Copy + Default, const K: usize, const B: usize> {
+    #[cfg_attr(feature = "rkyv_08", rkyv(with = EncodeAVec<A>))]
     pub(crate) stems: AVec<A>,
 
     #[cfg_attr(feature = "serde", serde(with = "crate::custom_serde::array_of_vecs"))]
@@ -130,6 +139,37 @@ where
     }
 }
 
+#[cfg(feature = "rkyv")]
+impl<A: Axis, T: Content, const K: usize, const B: usize> From<ImmutableKdTreeRK<A, T, K, B>>
+    for ImmutableKdTree<A, T, K, B>
+where
+    A: Axis + LeafSliceFloat<T> + LeafSliceFloatChunk<T, K>,
+    T: Content,
+    usize: Cast<T>,
+{
+    /// Creates an [`ImmutableKdTree`]  from an [`ImmutableKdTreeRK`]
+    fn from(orig: ImmutableKdTreeRK<A, T, K, B>) -> Self {
+        let ImmutableKdTreeRK {
+            stems,
+            leaf_points,
+            leaf_items,
+            leaf_extents,
+            max_stem_level,
+        } = orig;
+
+        let (ptr, length, capacity) = stems.into_raw_parts();
+        let stems = unsafe { AVec::from_raw_parts(ptr, CACHELINE_ALIGN, length, capacity) };
+
+        ImmutableKdTree {
+            stems,
+            leaf_points,
+            leaf_items,
+            leaf_extents,
+            max_stem_level,
+        }
+    }
+}
+
 /// rkyv zero-copy deserializable version of an `ImmutableKdTree`.
 ///
 /// Convert an `ImmutableKdTreeRK` into this in order to perform queries.
@@ -214,6 +254,37 @@ where
     }
 }
 
+#[cfg(feature = "rkyv_08")]
+impl<A, T, const K: usize, const B: usize> ArchivedImmutableKdTree<A, T, K, B>
+where
+    A: Axis + LeafSliceFloat<T> + LeafSliceFloatChunk<T, K> + rkyv_08::Archive<Archived = A>,
+    T: Content + rkyv_08::Archive<Archived = T>,
+    usize: Cast<T>,
+{
+    /// Returns the current number of elements stored in the tree
+    #[inline]
+    pub fn size(&self) -> usize {
+        self.leaf_items.len()
+    }
+
+    /// Returns a LeafSlice for a given leaf index
+    #[inline]
+    pub(crate) fn get_leaf_slice(&self, leaf_idx: usize) -> LeafSlice<A, T, K> {
+        let extents = unsafe { self.leaf_extents.get_unchecked(leaf_idx) };
+        let start = Into::<u32>::into(extents.0) as usize;
+        let end = Into::<u32>::into(extents.1) as usize;
+
+        // Artificially extend size to be at least chunk length for faster processing
+        // TODO: why does this slow things down?
+        // let end = end.max(start + 32).min(self.leaf_items.len() as u32);
+
+        LeafSlice::new(
+            array_init::array_init(|i| &self.leaf_points[i][start..end]),
+            &self.leaf_items[start..end],
+        )
+    }
+}
+
 impl<A: Axis, T: Content, const K: usize, const B: usize> From<&[[A; K]]>
     for ImmutableKdTree<A, T, K, B>
 where

src/immutable/float/mod.rs

@@ -24,3 +24,6 @@
 pub mod kdtree;
 #[doc(hidden)]
 pub mod query;
+
+#[cfg(feature = "rkyv_08")]
+mod rkyv_aligned_vec;

src/immutable/float/query/nearest_one.rs

@@ -72,6 +72,26 @@ where
     );
 }
 
+#[cfg(feature = "rkyv_08")]
+use crate::immutable::float::kdtree::ArchivedImmutableKdTree;
+#[cfg(feature = "rkyv_08")]
+impl<A, T, const K: usize, const B: usize> ArchivedImmutableKdTree<A, T, K, B>
+where
+    A: Axis + LeafSliceFloat<T> + LeafSliceFloatChunk<T, K> + rkyv_08::Archive<Archived = A>,
+    T: Content + rkyv_08::Archive<Archived = T>,
+    usize: Cast<T>,
+{
+    generate_immutable_float_nearest_one!(
+        "use std::fs::File;
+    use memmap::MmapOptions;
+
+    use kiddo::immutable::float::kdtree::AlignedArchivedImmutableKdTree;
+
+    let mmap = unsafe { MmapOptions::new().map(&File::open(\"./examples/immutable-doctest-tree.rkyv\").unwrap()).unwrap() };
+    let tree = unsafe { access_unchecked::<Archived<ImmutableKdTree<f64, u32, 3, 256>>>(&buf) };"
+    );
+}
+
 #[cfg(test)]
 mod tests {
     use crate::float::distance::SquaredEuclidean;