kiddo 6.0.0

A high-performance, flexible, ergonomic k-d tree library. Ideal for geo- and astro- nearest-neighbour and k-nearest-neighbor queries
Documentation
use aligned_vec::{AVec, ConstAlign, CACHELINE_ALIGN};

use crate::kd_tree::ConstructionError;
use crate::traits::leaf_strategy::ConstructibleLeafStrategy;
use crate::{Axis, Content, StemStrategy};

pub(super) trait ConstructionIndex: Copy + Send {
    fn from_usize(value: usize) -> Self;
    fn as_usize(self) -> usize;
}

impl ConstructionIndex for u32 {
    #[inline(always)]
    fn from_usize(value: usize) -> Self {
        value as u32
    }

    #[inline(always)]
    fn as_usize(self) -> usize {
        self as usize
    }
}

impl ConstructionIndex for usize {
    #[inline(always)]
    fn from_usize(value: usize) -> Self {
        value
    }

    #[inline(always)]
    fn as_usize(self) -> usize {
        self
    }
}

#[inline(always)]
pub(super) const fn construction_index_fits_u32(item_count: usize) -> bool {
    item_count <= u32::MAX as usize
}

pub(super) struct ConstructionLeafScratch<A, T, const K: usize> {
    pub(super) points: [Vec<A>; K],
    pub(super) items: Vec<T>,
}

impl<A, T, const K: usize> ConstructionLeafScratch<A, T, K> {
    pub(super) fn with_capacity(capacity: usize) -> Self {
        Self {
            points: array_init::array_init(|_| Vec::with_capacity(capacity)),
            items: Vec::with_capacity(capacity),
        }
    }

    pub(super) fn clear_and_reserve(&mut self, leaf_len: usize) {
        for axis in &mut self.points {
            axis.clear();
            if axis.capacity() < leaf_len {
                axis.reserve(leaf_len);
            }
        }
        self.items.clear();
        if self.items.capacity() < leaf_len {
            self.items.reserve(leaf_len);
        }
    }
}

pub(in crate::kd_tree) fn validate_auto_generated_items<T>(
    item_count: usize,
) -> Result<(), ConstructionError>
where
    T: TryFrom<usize>,
{
    if let Some(max_src_idx) = item_count.checked_sub(1) {
        if T::try_from(max_src_idx).is_err() {
            return Err(ConstructionError::AutoGeneratedItemIndexOverflow {
                item_count,
                item_type: core::any::type_name::<T>(),
            });
        }
    }

    Ok(())
}

pub(super) trait SoftConstructionMode<A, T, SS, LS, I, X, FA, FI, const K: usize, const B: usize>
where
    A: Axis<Coord = A>,
    T: Content,
    SS: StemStrategy,
    LS: ConstructibleLeafStrategy<A, T, SS, K, B>,
    I: ConstructionIndex,
    FA: Fn(&X, usize) -> A,
    FI: FnMut(usize, &X) -> Result<T, ConstructionError>,
{
    #[allow(clippy::too_many_arguments)]
    fn populate(
        &self,
        stems: &mut AVec<A, ConstAlign<{ CACHELINE_ALIGN }>>,
        source: &[X],
        axis_at: &FA,
        sort_index: &mut [I],
        root_stem_ordering: SS,
        max_stem_level: i32,
        leaf_budget: usize,
        leaves: &mut LS,
        actual_max_stem_level: &mut i32,
        max_leaf_len: &mut usize,
        leaf_scratch: &mut ConstructionLeafScratch<A, T, K>,
        item_at: &mut FI,
    ) -> Result<(), ConstructionError>;
}