robert/another-boids-in-rust
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Vendor dependencies for 0.3.0 release

Browse Source
This commit is contained in:
Robert Garrett
2025-09-27 10:29:08 -05:00
parent 0c8d39d483
commit 82ab7f317b
26803 changed files with 16134934 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

840
vendor/kd-tree/src/lib.rs vendored Normal file
View File

@@ -0,0 +1,840 @@
//! k-dimensional tree.
//!
//! # Usage
//! ```
//! // construct kd-tree
//! let kdtree = kd_tree::KdTree::build_by_ordered_float(vec![
//! [1.0, 2.0, 3.0],
//! [3.0, 1.0, 2.0],
//! [2.0, 3.0, 1.0],
//! ]);
//!
//! // search the nearest neighbor
//! let found = kdtree.nearest(&[3.1, 0.9, 2.1]).unwrap();
//! assert_eq!(found.item, &[3.0, 1.0, 2.0]);
//!
//! // search k-nearest neighbors
//! let found = kdtree.nearests(&[1.5, 2.5, 1.8], 2);
//! assert_eq!(found[0].item, &[2.0, 3.0, 1.0]);
//! assert_eq!(found[1].item, &[1.0, 2.0, 3.0]);
//!
//! // search points within a sphere
//! let found = kdtree.within_radius(&[2.0, 1.5, 2.5], 1.5);
//! assert_eq!(found.len(), 2);
//! assert!(found.iter().any(|&&p| p == [1.0, 2.0, 3.0]));
//! assert!(found.iter().any(|&&p| p == [3.0, 1.0, 2.0]));
//! ```
mod nalgebra;
mod nearest;
mod nearests;
mod sort;
mod tests;
mod within;
use nearest::*;
use nearests::*;
use sort::*;
use std::cmp::Ordering;
use std::marker::PhantomData;
use typenum::Unsigned;
use within::*;
/// A trait to represent k-dimensional point.
///
/// # Example
/// ```
/// struct MyItem {
/// point: [f64; 3],
/// id: usize,
/// }
/// impl kd_tree::KdPoint for MyItem {
/// type Scalar = f64;
/// type Dim = typenum::U3;
/// fn at(&self, k: usize) -> f64 { self.point[k] }
/// }
/// let kdtree: kd_tree::KdTree<MyItem> = kd_tree::KdTree::build_by_ordered_float(vec![
/// MyItem { point: [1.0, 2.0, 3.0], id: 111 },
/// MyItem { point: [3.0, 1.0, 2.0], id: 222 },
/// MyItem { point: [2.0, 3.0, 1.0], id: 333 },
/// ]);
/// assert_eq!(kdtree.nearest(&[3.1, 0.1, 2.2]).unwrap().item.id, 222);
/// ```
pub trait KdPoint {
type Scalar: num_traits::NumAssign + Copy + PartialOrd;
type Dim: Unsigned;
fn dim() -> usize {
<Self::Dim as Unsigned>::to_usize()
}
fn at(&self, i: usize) -> Self::Scalar;
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ItemAndDistance<'a, T, Scalar> {
pub item: &'a T,
pub squared_distance: Scalar,
}
/// A slice of kd-tree.
/// This type implements [`std::ops::Deref`] to `[T]`.
/// This is an unsized type, meaning that it must always be used as a reference.
/// For an owned version of this type, see [`KdTree`].
#[derive(Debug, PartialEq, Eq)]
pub struct KdSliceN<T, N: Unsigned>(PhantomData<N>, [T]);
pub type KdSlice<T> = KdSliceN<T, <T as KdPoint>::Dim>;
impl<T, N: Unsigned> std::ops::Deref for KdSliceN<T, N> {
type Target = [T];
fn deref(&self) -> &[T] {
&self.1
}
}
impl<T: Clone, N: Unsigned> std::borrow::ToOwned for KdSliceN<T, N> {
type Owned = KdTreeN<T, N>;
fn to_owned(&self) -> Self::Owned {
KdTreeN(PhantomData, self.1.to_vec())
}
}
impl<T, N: Unsigned> KdSliceN<T, N> {
pub fn items(&self) -> &[T] {
&self.1
}
unsafe fn new_unchecked(items: &[T]) -> &Self {
&*(items as *const _ as *const Self)
}
/// # Example
/// ```
/// struct Item {
/// point: [i32; 3],
/// id: usize,
/// }
/// let mut items: Vec<Item> = vec![
/// Item { point: [1, 2, 3], id: 111 },
/// Item { point: [3, 1, 2], id: 222 },
/// Item { point: [2, 3, 1], id: 333 },
/// ];
/// let kdtree = kd_tree::KdSlice3::sort_by(&mut items, |item1, item2, k| item1.point[k].cmp(&item2.point[k]));
/// assert_eq!(kdtree.nearest_by(&[3, 1, 2], |item, k| item.point[k]).unwrap().item.id, 222);
/// ```
pub fn sort_by<F>(items: &mut [T], compare: F) -> &Self
where
F: Fn(&T, &T, usize) -> Ordering + Copy,
{
kd_sort_by(items, N::to_usize(), compare);
unsafe { Self::new_unchecked(items) }
}
/// # Example
/// ```
/// struct Item {
/// point: [f64; 3],
/// id: usize,
/// }
/// let mut items: Vec<Item> = vec![
/// Item { point: [1.0, 2.0, 3.0], id: 111 },
/// Item { point: [3.0, 1.0, 2.0], id: 222 },
/// Item { point: [2.0, 3.0, 1.0], id: 333 },
/// ];
/// use ordered_float::OrderedFloat;
/// let kdtree = kd_tree::KdSlice3::sort_by_key(&mut items, |item, k| OrderedFloat(item.point[k]));
/// assert_eq!(kdtree.nearest_by(&[3.1, 0.9, 2.1], |item, k| item.point[k]).unwrap().item.id, 222);
/// ```
pub fn sort_by_key<Key: Ord, F>(items: &mut [T], kd_key: F) -> &Self
where
F: Fn(&T, usize) -> Key + Copy,
{
Self::sort_by(items, |item1, item2, k| {
kd_key(item1, k).cmp(&kd_key(item2, k))
})
}
/// # Example
/// ```
/// use kd_tree::KdSlice;
/// let mut items: Vec<[f64; 3]> = vec![[1.0, 2.0, 3.0], [3.0, 1.0, 2.0], [2.0, 3.0, 1.0]];
/// let kdtree: &KdSlice<[f64; 3]> = KdSlice::sort_by_ordered_float(&mut items);
/// assert_eq!(kdtree.nearest(&[3.1, 0.9, 2.1]).unwrap().item, &[3.0, 1.0, 2.0]);
/// ```
pub fn sort_by_ordered_float(points: &mut [T]) -> &Self
where
T: KdPoint<Dim = N>,
T::Scalar: ordered_float::FloatCore,
{
Self::sort_by_key(points, |item, k| ordered_float::OrderedFloat(item.at(k)))
}
/// # Example
/// ```
/// use kd_tree::KdSlice;
/// let mut items: Vec<[i32; 3]> = vec![[1, 2, 3], [3, 1, 2], [2, 3, 1]];
/// let kdtree: &KdSlice<[i32; 3]> = KdSlice::sort(&mut items);
/// assert_eq!(kdtree.nearest(&[3, 1, 2]).unwrap().item, &[3, 1, 2]);
/// ```
pub fn sort(points: &mut [T]) -> &Self
where
T: KdPoint<Dim = N>,
T::Scalar: Ord,
{
Self::sort_by_key(points, |item, k| item.at(k))
}
/// Returns the nearest item from the input point. Returns `None` if `self.is_empty()`.
/// # Example
/// ```
/// struct Item {
/// point: [f64; 3],
/// id: usize,
/// }
/// let mut items: Vec<Item> = vec![
/// Item { point: [1.0, 2.0, 3.0], id: 111 },
/// Item { point: [3.0, 1.0, 2.0], id: 222 },
/// Item { point: [2.0, 3.0, 1.0], id: 333 },
/// ];
/// use ordered_float::OrderedFloat;
/// let kdtree = kd_tree::KdSlice3::sort_by_key(&mut items, |item, k| OrderedFloat(item.point[k]));
/// assert_eq!(kdtree.nearest_by(&[3.1, 0.9, 2.1], |item, k| item.point[k]).unwrap().item.id, 222);
/// ```
pub fn nearest_by<Q: KdPoint<Dim = N>>(
&self,
query: &Q,
coord: impl Fn(&T, usize) -> Q::Scalar + Copy,
) -> Option<ItemAndDistance<T, Q::Scalar>> {
if self.is_empty() {
None
} else {
Some(kd_nearest_by(self.items(), query, coord))
}
}
/// Returns the nearest item from the input point. Returns `None` if `self.is_empty()`.
/// # Example
/// ```
/// let mut items: Vec<[i32; 3]> = vec![[1, 2, 3], [3, 1, 2], [2, 3, 1]];
/// let kdtree = kd_tree::KdSlice::sort(&mut items);
/// assert_eq!(kdtree.nearest(&[3, 1, 2]).unwrap().item, &[3, 1, 2]);
/// ```
pub fn nearest(
&self,
query: &impl KdPoint<Scalar = T::Scalar, Dim = N>,
) -> Option<ItemAndDistance<T, T::Scalar>>
where
T: KdPoint<Dim = N>,
{
if self.is_empty() {
None
} else {
Some(kd_nearest(self.items(), query))
}
}
/*
/// # Example
/// ```
/// let kdtree = kd_tree::KdTree3::build(vec![[1, 2, 3], [3, 1, 2], [2, 3, 1]]);
/// let key = [3, 1, 2];
/// assert_eq!(kdtree.nearest_with(|p, k| key[k] - p[k]).item, &[3, 1, 2]);
/// ```
pub fn nearest_with<Scalar>(
&self,
kd_difference: impl Fn(&T, usize) -> Scalar + Copy,
) -> ItemAndDistance<T, Scalar>
where
Scalar: num_traits::NumAssign + Copy + PartialOrd,
{
kd_nearest_with(self.items(), N::to_usize(), kd_difference)
}
*/
/// Returns the nearest item from the input point. Returns `None` if `self.is_empty()`.
/// # Example
/// ```
/// struct Item {
/// point: [f64; 3],
/// id: usize,
/// }
/// let mut items: Vec<Item> = vec![
/// Item { point: [1.0, 2.0, 3.0], id: 111 },
/// Item { point: [3.0, 1.0, 2.0], id: 222 },
/// Item { point: [2.0, 3.0, 1.0], id: 333 },
/// ];
/// use ordered_float::OrderedFloat;
/// let kdtree = kd_tree::KdSlice3::sort_by_key(&mut items, |item, k| OrderedFloat(item.point[k]));
/// let nearests = kdtree.nearests_by(&[2.5, 2.0, 1.4], 2, |item, k| item.point[k]);
/// assert_eq!(nearests.len(), 2);
/// assert_eq!(nearests[0].item.id, 333);
/// assert_eq!(nearests[1].item.id, 222);
/// ```
pub fn nearests_by<Q: KdPoint<Dim = N>>(
&self,
query: &Q,
num: usize,
coord: impl Fn(&T, usize) -> Q::Scalar + Copy,
) -> Vec<ItemAndDistance<T, Q::Scalar>> {
kd_nearests_by(self.items(), query, num, coord)
}
/// Returns kNN(k nearest neighbors) from the input point.
/// # Example
/// ```
/// let mut items: Vec<[i32; 3]> = vec![[1, 2, 3], [3, 1, 2], [2, 3, 1], [3, 2, 2]];
/// let kdtree = kd_tree::KdSlice::sort(&mut items);
/// let nearests = kdtree.nearests(&[3, 1, 2], 2);
/// assert_eq!(nearests.len(), 2);
/// assert_eq!(nearests[0].item, &[3, 1, 2]);
/// assert_eq!(nearests[1].item, &[3, 2, 2]);
/// ```
pub fn nearests(
&self,
query: &impl KdPoint<Scalar = T::Scalar, Dim = N>,
num: usize,
) -> Vec<ItemAndDistance<T, T::Scalar>>
where
T: KdPoint<Dim = N>,
{
kd_nearests(self.items(), query, num)
}
pub fn within_by_cmp(&self, compare: impl Fn(&T, usize) -> Ordering + Copy) -> Vec<&T> {
kd_within_by_cmp(self, N::to_usize(), compare)
}
pub fn within_by<Q: KdPoint<Dim = N>>(
&self,
query: &[Q; 2],
coord: impl Fn(&T, usize) -> Q::Scalar + Copy,
) -> Vec<&T> {
assert!((0..Q::dim()).all(|k| query[0].at(k) <= query[1].at(k)));
self.within_by_cmp(|item, k| {
let a = coord(item, k);
if a < query[0].at(k) {
Ordering::Less
} else if a > query[1].at(k) {
Ordering::Greater
} else {
Ordering::Equal
}
})
}
/// search points within a rectangular region
/// # Example
/// ```
/// let mut items: Vec<[i32; 2]> = vec![[0, 0], [1, 0], [0, 1], [1, 1]];
/// let kdtree = kd_tree::KdSlice::sort(&mut items);
/// let within = kdtree.within(&[[1, 0], [2, 1]]);
/// assert_eq!(within.len(), 2);
/// assert!(within.contains(&&[1, 0]));
/// assert!(within.contains(&&[1, 1]));
/// ```
pub fn within(&self, query: &[impl KdPoint<Scalar = T::Scalar, Dim = N>; 2]) -> Vec<&T>
where
T: KdPoint<Dim = N>,
{
self.within_by(query, |item, k| item.at(k))
}
pub fn within_radius_by<Q: KdPoint<Dim = N>>(
&self,
query: &Q,
radius: Q::Scalar,
coord: impl Fn(&T, usize) -> Q::Scalar + Copy,
) -> Vec<&T> {
let mut results = self.within_by_cmp(|item, k| {
let coord = coord(item, k);
if coord < query.at(k) - radius {
Ordering::Less
} else if coord > query.at(k) + radius {
Ordering::Greater
} else {
Ordering::Equal
}
});
results.retain(|item| {
let mut distance = <Q::Scalar as num_traits::Zero>::zero();
for k in 0..N::to_usize() {
let diff = coord(item, k) - query.at(k);
distance += diff * diff;
}
distance < radius * radius
});
results
}
/// search points within k-dimensional sphere
pub fn within_radius(
&self,
query: &impl KdPoint<Scalar = T::Scalar, Dim = N>,
radius: T::Scalar,
) -> Vec<&T>
where
T: KdPoint<Dim = N>,
{
self.within_radius_by(query, radius, |item, k| item.at(k))
}
}
#[cfg(feature = "rayon")]
impl<T: Send, N: Unsigned> KdSliceN<T, N> {
/// Same as [`Self::sort_by`], but using multiple threads.
pub fn par_sort_by<F>(items: &mut [T], compare: F) -> &Self
where
F: Fn(&T, &T, usize) -> Ordering + Copy + Send,
{
kd_par_sort_by(items, N::to_usize(), compare);
unsafe { Self::new_unchecked(items) }
}
/// Same as [`Self::sort_by_key`], but using multiple threads.
pub fn par_sort_by_key<Key: Ord, F>(items: &mut [T], kd_key: F) -> &Self
where
F: Fn(&T, usize) -> Key + Copy + Send,
{
Self::par_sort_by(items, move |item1, item2, k| {
kd_key(item1, k).cmp(&kd_key(item2, k))
})
}
/// Same as [`Self::sort_by_ordered_float`], but using multiple threads.
pub fn par_sort_by_ordered_float(points: &mut [T]) -> &Self
where
T: KdPoint<Dim = N>,
T::Scalar: ordered_float::FloatCore,
{
Self::par_sort_by_key(points, |item, k| ordered_float::OrderedFloat(item.at(k)))
}
/// Same as [`Self::sort`], but using multiple threads.
pub fn par_sort(points: &mut [T]) -> &Self
where
T: KdPoint<Dim = N>,
T::Scalar: Ord,
{
Self::par_sort_by_key(points, |item, k| item.at(k))
}
}
/// An owned kd-tree.
/// This type implements [`std::ops::Deref`] to [`KdSlice`].
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct KdTreeN<T, N: Unsigned>(PhantomData<N>, Vec<T>);
pub type KdTree<T> = KdTreeN<T, <T as KdPoint>::Dim>;
impl<T, N: Unsigned> std::ops::Deref for KdTreeN<T, N> {
type Target = KdSliceN<T, N>;
fn deref(&self) -> &Self::Target {
unsafe { KdSliceN::new_unchecked(&self.1) }
}
}
impl<T, N: Unsigned> AsRef<KdSliceN<T, N>> for KdTreeN<T, N> {
fn as_ref(&self) -> &KdSliceN<T, N> {
self
}
}
impl<T, N: Unsigned> std::borrow::Borrow<KdSliceN<T, N>> for KdTreeN<T, N> {
fn borrow(&self) -> &KdSliceN<T, N> {
self
}
}
impl<T, N: Unsigned> From<KdTreeN<T, N>> for Vec<T> {
fn from(src: KdTreeN<T, N>) -> Self {
src.1
}
}
impl<T, N: Unsigned> KdTreeN<T, N> {
pub fn into_vec(self) -> Vec<T> {
self.1
}
/// # Example
/// ```
/// struct Item {
/// point: [i32; 3],
/// id: usize,
/// }
/// let kdtree = kd_tree::KdTree3::build_by(
/// vec![
/// Item { point: [1, 2, 3], id: 111 },
/// Item { point: [3, 1, 2], id: 222 },
/// Item { point: [2, 3, 1], id: 333 },
/// ],
/// |item1, item2, k| item1.point[k].cmp(&item2.point[k])
/// );
/// assert_eq!(kdtree.nearest_by(&[3, 1, 2], |item, k| item.point[k]).unwrap().item.id, 222);
/// ```
pub fn build_by<F>(mut items: Vec<T>, compare: F) -> Self
where
F: Fn(&T, &T, usize) -> Ordering + Copy,
{
kd_sort_by(&mut items, N::to_usize(), compare);
Self(PhantomData, items)
}
/// # Example
/// ```
/// struct Item {
/// point: [f64; 3],
/// id: usize,
/// }
/// let kdtree = kd_tree::KdTree3::build_by_key(
/// vec![
/// Item { point: [1.0, 2.0, 3.0], id: 111 },
/// Item { point: [3.0, 1.0, 2.0], id: 222 },
/// Item { point: [2.0, 3.0, 1.0], id: 333 },
/// ],
/// |item, k| ordered_float::OrderedFloat(item.point[k])
/// );
/// assert_eq!(kdtree.nearest_by(&[3.1, 0.9, 2.1], |item, k| item.point[k]).unwrap().item.id, 222);
/// ```
pub fn build_by_key<Key, F>(items: Vec<T>, kd_key: F) -> Self
where
Key: Ord,
F: Fn(&T, usize) -> Key + Copy,
{
Self::build_by(items, |item1, item2, k| {
kd_key(item1, k).cmp(&kd_key(item2, k))
})
}
/// # Example
/// ```
/// use kd_tree::KdTree;
/// let kdtree: KdTree<[f64; 3]> = KdTree::build_by_ordered_float(vec![
/// [1.0, 2.0, 3.0], [3.0, 1.0, 2.0], [2.0, 3.0, 1.0]
/// ]);
/// assert_eq!(kdtree.nearest(&[3.1, 0.9, 2.1]).unwrap().item, &[3.0, 1.0, 2.0]);
/// ```
pub fn build_by_ordered_float(points: Vec<T>) -> Self
where
T: KdPoint<Dim = N>,
T::Scalar: ordered_float::FloatCore,
{
Self::build_by_key(points, |item, k| ordered_float::OrderedFloat(item.at(k)))
}
/// # Example
/// ```
/// use kd_tree::KdTree;
/// let kdtree: KdTree<[i32; 3]> = KdTree::build(vec![[1, 2, 3], [3, 1, 2], [2, 3, 1]]);
/// assert_eq!(kdtree.nearest(&[3, 1, 2]).unwrap().item, &[3, 1, 2]);
/// ```
pub fn build(points: Vec<T>) -> Self
where
T: KdPoint<Dim = N>,
T::Scalar: Ord,
{
Self::build_by_key(points, |item, k| item.at(k))
}
}
#[cfg(feature = "serde")]
mod impl_serde {
use super::{KdTreeN, PhantomData, Unsigned};
impl<T: serde::Serialize, N: Unsigned> serde::Serialize for KdTreeN<T, N> {
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
self.1.serialize(serializer)
}
}
impl<'de, T: serde::Deserialize<'de>, N: Unsigned> serde::Deserialize<'de> for KdTreeN<T, N> {
fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
Vec::<T>::deserialize(deserializer).map(|items| Self(PhantomData, items))
}
}
}
#[cfg(feature = "rayon")]
impl<T: Send, N: Unsigned> KdTreeN<T, N> {
/// Same as [`Self::build_by`], but using multiple threads.
pub fn par_build_by<F>(mut items: Vec<T>, compare: F) -> Self
where
F: Fn(&T, &T, usize) -> Ordering + Copy + Send,
{
kd_par_sort_by(&mut items, N::to_usize(), compare);
Self(PhantomData, items)
}
/// Same as [`Self::build_by_key`], but using multiple threads.
pub fn par_build_by_key<Key, F>(items: Vec<T>, kd_key: F) -> Self
where
Key: Ord,
F: Fn(&T, usize) -> Key + Copy + Send,
{
Self::par_build_by(items, move |item1, item2, k| {
kd_key(item1, k).cmp(&kd_key(item2, k))
})
}
/// Same as [`Self::build_by_ordered_float`], but using multiple threads.
pub fn par_build_by_ordered_float(points: Vec<T>) -> Self
where
T: KdPoint<Dim = N>,
T::Scalar: ordered_float::FloatCore,
{
Self::par_build_by_key(points, |item, k| ordered_float::OrderedFloat(item.at(k)))
}
/// Same as [`Self::build`], but using multiple threads.
pub fn par_build(points: Vec<T>) -> Self
where
T: KdPoint<Dim = N>,
T::Scalar: Ord,
{
Self::par_build_by_key(points, |item, k| item.at(k))
}
}
/// This type refers a slice of items, `[T]`, and contains kd-tree of indices to the items, `KdTree<usize, N>`.
/// Unlike [`KdSliceN::sort`], [`KdIndexTreeN::build`] doesn't sort input items.
/// ```
/// let items = vec![[1, 2, 3], [3, 1, 2], [2, 3, 1]];
/// let kdtree = kd_tree::KdIndexTree::build(&items);
/// assert_eq!(kdtree.nearest(&[3, 1, 2]).unwrap().item, &1); // nearest() returns an index of items.
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KdIndexTreeN<'a, T, N: Unsigned> {
source: &'a [T],
kdtree: KdTreeN<usize, N>,
}
pub type KdIndexTree<'a, T> = KdIndexTreeN<'a, T, <T as KdPoint>::Dim>;
impl<'a, T, N: Unsigned> KdIndexTreeN<'a, T, N> {
pub fn source(&self) -> &'a [T] {
self.source
}
pub fn indices(&self) -> &KdSliceN<usize, N> {
&self.kdtree
}
pub fn item(&self, i: usize) -> &'a T {
&self.source[i]
}
pub fn build_by<F>(source: &'a [T], compare: F) -> Self
where
F: Fn(&T, &T, usize) -> Ordering + Copy,
{
Self {
source,
kdtree: KdTreeN::build_by((0..source.len()).collect(), |i1, i2, k| {
compare(&source[*i1], &source[*i2], k)
}),
}
}
pub fn build_by_key<Key, F>(source: &'a [T], kd_key: F) -> Self
where
Key: Ord,
F: Fn(&T, usize) -> Key + Copy,
{
Self::build_by(source, |item1, item2, k| {
kd_key(item1, k).cmp(&kd_key(item2, k))
})
}
pub fn build_by_ordered_float(points: &'a [T]) -> Self
where
T: KdPoint<Dim = N>,
T::Scalar: ordered_float::FloatCore,
{
Self::build_by_key(points, |item, k| ordered_float::OrderedFloat(item.at(k)))
}
pub fn build(points: &'a [T]) -> Self
where
T: KdPoint<Dim = N>,
T::Scalar: Ord,
{
Self::build_by_key(points, |item, k| item.at(k))
}
pub fn nearest_by<Q: KdPoint<Dim = N>>(
&self,
query: &Q,
coord: impl Fn(&T, usize) -> Q::Scalar + Copy,
) -> Option<ItemAndDistance<usize, Q::Scalar>> {
self.kdtree
.nearest_by(query, |&index, k| coord(&self.source[index], k))
}
/// # Example
/// ```
/// let mut items: Vec<[i32; 3]> = vec![[1, 2, 3], [3, 1, 2], [2, 3, 1]];
/// let kdtree = kd_tree::KdIndexTree3::build(&items);
/// assert_eq!(kdtree.nearest(&[3, 1, 2]).unwrap().item, &1);
/// ```
pub fn nearest(
&self,
query: &impl KdPoint<Scalar = T::Scalar, Dim = N>,
) -> Option<ItemAndDistance<usize, T::Scalar>>
where
T: KdPoint<Dim = N>,
{
self.nearest_by(query, |item, k| item.at(k))
}
pub fn nearests_by<Q: KdPoint<Dim = N>>(
&self,
query: &Q,
num: usize,
coord: impl Fn(&T, usize) -> Q::Scalar + Copy,
) -> Vec<ItemAndDistance<usize, Q::Scalar>> {
self.kdtree
.nearests_by(query, num, |&index, k| coord(&self.source[index], k))
}
/// Returns kNN(k nearest neighbors) from the input point.
/// # Example
/// ```
/// let mut items: Vec<[i32; 3]> = vec![[1, 2, 3], [3, 1, 2], [2, 3, 1], [3, 2, 2]];
/// let kdtree = kd_tree::KdIndexTree::build(&mut items);
/// let nearests = kdtree.nearests(&[3, 1, 2], 2);
/// assert_eq!(nearests.len(), 2);
/// assert_eq!(nearests[0].item, &1);
/// assert_eq!(nearests[1].item, &3);
/// ```
pub fn nearests(
&self,
query: &impl KdPoint<Scalar = T::Scalar, Dim = N>,
num: usize,
) -> Vec<ItemAndDistance<usize, T::Scalar>>
where
T: KdPoint<Dim = N>,
{
self.nearests_by(query, num, |item, k| item.at(k))
}
pub fn within_by_cmp(&self, compare: impl Fn(&T, usize) -> Ordering + Copy) -> Vec<&usize> {
self.kdtree
.within_by_cmp(|&index, k| compare(&self.source[index], k))
}
pub fn within_by<Q: KdPoint<Dim = N>>(
&self,
query: &[Q; 2],
coord: impl Fn(&T, usize) -> Q::Scalar + Copy,
) -> Vec<&usize> {
self.kdtree
.within_by(query, |&index, k| coord(&self.source[index], k))
}
pub fn within(&self, query: &[impl KdPoint<Scalar = T::Scalar, Dim = N>; 2]) -> Vec<&usize>
where
T: KdPoint<Dim = N>,
{
self.within_by(query, |item, k| item.at(k))
}
pub fn within_radius_by<Q: KdPoint<Dim = N>>(
&self,
query: &Q,
radius: Q::Scalar,
coord: impl Fn(&T, usize) -> Q::Scalar + Copy,
) -> Vec<&usize> {
self.kdtree
.within_radius_by(query, radius, |&index, k| coord(&self.source[index], k))
}
pub fn within_radius(
&self,
query: &impl KdPoint<Scalar = T::Scalar, Dim = N>,
radius: T::Scalar,
) -> Vec<&usize>
where
T: KdPoint<Dim = N>,
{
self.within_radius_by(query, radius, |item, k| item.at(k))
}
}
#[cfg(feature = "rayon")]
impl<'a, T: Sync, N: Unsigned> KdIndexTreeN<'a, T, N> {
pub fn par_build_by<F>(source: &'a [T], compare: F) -> Self
where
F: Fn(&T, &T, usize) -> Ordering + Copy + Send,
{
Self {
source,
kdtree: KdTreeN::par_build_by((0..source.len()).collect(), move |i1, i2, k| {
compare(&source[*i1], &source[*i2], k)
}),
}
}
pub fn par_build_by_key<Key, F>(source: &'a [T], kd_key: F) -> Self
where
Key: Ord,
F: Fn(&T, usize) -> Key + Copy + Send,
{
Self::par_build_by(source, move |item1, item2, k| {
kd_key(item1, k).cmp(&kd_key(item2, k))
})
}
pub fn par_build_by_ordered_float(points: &'a [T]) -> Self
where
T: KdPoint<Dim = N>,
T::Scalar: ordered_float::FloatCore,
{
Self::par_build_by_key(points, |item, k| ordered_float::OrderedFloat(item.at(k)))
}
pub fn par_build(points: &'a [T]) -> Self
where
T: KdPoint<Dim = N>,
T::Scalar: Ord,
{
Self::par_build_by_key(points, |item, k| item.at(k))
}
}
macro_rules! define_kdtree_aliases {
($($dim:literal),*) => {
$(
paste::paste! {
pub type [<KdSlice $dim>]<T> = KdSliceN<T, typenum::[<U $dim>]>;
pub type [<KdTree $dim>]<T> = KdTreeN<T, typenum::[<U $dim>]>;
pub type [<KdIndexTree $dim>]<'a, T> = KdIndexTreeN<'a, T, typenum::[<U $dim>]>;
}
)*
};
}
define_kdtree_aliases!(1, 2, 3, 4, 5, 6, 7, 8);
macro_rules! impl_kd_points {
($($len:literal),*) => {
$(
paste::paste!{
impl<T: num_traits::NumAssign + Copy + PartialOrd> KdPoint for [T; $len] {
type Scalar = T;
type Dim = typenum::[<U $len>];
fn at(&self, i: usize) -> T { self[i] }
}
}
)*
};
}
impl_kd_points!(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16);
impl<P: KdPoint, T> KdPoint for (P, T) {
type Scalar = P::Scalar;
type Dim = P::Dim;
fn at(&self, k: usize) -> Self::Scalar {
self.0.at(k)
}
}
/// kd-tree of key-value pairs.
/// ```
/// let kdmap: kd_tree::KdMap<[isize; 3], &'static str> = kd_tree::KdMap::build(vec![
/// ([1, 2, 3], "foo"),
/// ([2, 3, 1], "bar"),
/// ([3, 1, 2], "buzz"),
/// ]);
/// assert_eq!(kdmap.nearest(&[3, 1, 2]).unwrap().item.1, "buzz");
/// ```
pub type KdMap<P, T> = KdTree<(P, T)>;
/// kd-tree slice of key-value pairs.
/// ```
/// let mut items: Vec<([isize; 3], &'static str)> = vec![
/// ([1, 2, 3], "foo"),
/// ([2, 3, 1], "bar"),
/// ([3, 1, 2], "buzz"),
/// ];
/// let kdmap = kd_tree::KdMapSlice::sort(&mut items);
/// assert_eq!(kdmap.nearest(&[3, 1, 2]).unwrap().item.1, "buzz");
/// ```
pub type KdMapSlice<P, T> = KdSlice<(P, T)>;

44
vendor/kd-tree/src/nalgebra.rs vendored Normal file
View File

@@ -0,0 +1,44 @@
#![cfg(feature = "nalgebra")]
use super::KdPoint;
macro_rules! impl_kdpoint_for_nalgebra_point {
($($dim:literal),*) => {
$(
paste::paste! {
impl<Scalar> KdPoint for nalgebra::Point<Scalar, $dim>
where
Scalar: num_traits::NumAssign + Copy + PartialOrd + nalgebra::Scalar,
{
type Scalar = Scalar;
type Dim = typenum::[<U $dim>];
fn at(&self, k: usize) -> Scalar {
self[k]
}
}
}
)*
};
}
macro_rules! impl_kdpoint_for_nalgebra_vector {
($($dim:literal),*) => {
$(
paste::paste! {
impl<Scalar, Storage> KdPoint for nalgebra::Vector<Scalar, nalgebra::Const<$dim>, Storage>
where
Scalar: num_traits::NumAssign + Copy + PartialOrd + nalgebra::Scalar,
Storage: nalgebra::StorageMut<Scalar, nalgebra::Const<$dim>>
{
type Scalar = Scalar;
type Dim = typenum::[<U $dim>];
fn at(&self, k: usize) -> Scalar {
self[k]
}
}
}
)*
};
}
impl_kdpoint_for_nalgebra_point!(1, 2, 3, 4, 5, 6, 7, 8);
impl_kdpoint_for_nalgebra_vector!(1, 2, 3, 4, 5, 6, 7, 8);

134
vendor/kd-tree/src/nearest.rs vendored Normal file
View File

@@ -0,0 +1,134 @@
use crate::{ItemAndDistance, KdPoint};
pub fn kd_nearest<'a, T: KdPoint>(
kdtree: &'a [T],
query: &impl KdPoint<Scalar = T::Scalar, Dim = T::Dim>,
) -> ItemAndDistance<'a, T, T::Scalar> {
kd_nearest_by(kdtree, query, |item, k| item.at(k))
}
pub fn kd_nearest_by<'a, T, P: KdPoint>(
kdtree: &'a [T],
query: &P,
get: impl Fn(&T, usize) -> P::Scalar + Copy,
) -> ItemAndDistance<'a, T, P::Scalar> {
fn distance_squared<P: KdPoint, T>(
p1: &P,
p2: &T,
get: impl Fn(&T, usize) -> P::Scalar,
) -> P::Scalar {
let mut squared_distance = <P::Scalar as num_traits::Zero>::zero();
for i in 0..P::dim() {
let diff = p1.at(i) - get(p2, i);
squared_distance += diff * diff;
}
squared_distance
}
fn recurse<'a, T, Q: KdPoint>(
nearest: &mut ItemAndDistance<'a, T, Q::Scalar>,
kdtree: &'a [T],
get: impl Fn(&T, usize) -> Q::Scalar + Copy,
query: &Q,
axis: usize,
) {
let mid_idx = kdtree.len() / 2;
let item = &kdtree[mid_idx];
let squared_distance = distance_squared(query, item, get);
if squared_distance < nearest.squared_distance {
nearest.item = item;
nearest.squared_distance = squared_distance;
use num_traits::Zero;
if nearest.squared_distance.is_zero() {
return;
}
}
let mid_pos = get(item, axis);
let [branch1, branch2] = if query.at(axis) < mid_pos {
[&kdtree[..mid_idx], &kdtree[mid_idx + 1..]]
} else {
[&kdtree[mid_idx + 1..], &kdtree[..mid_idx]]
};
if !branch1.is_empty() {
recurse(nearest, branch1, get, query, (axis + 1) % Q::dim());
}
if !branch2.is_empty() {
let diff = query.at(axis) - mid_pos;
if diff * diff < nearest.squared_distance {
recurse(nearest, branch2, get, query, (axis + 1) % Q::dim());
}
}
}
assert!(!kdtree.is_empty());
let mut nearest = ItemAndDistance {
item: &kdtree[0],
squared_distance: distance_squared(query, &kdtree[0], get),
};
recurse(&mut nearest, kdtree, get, query, 0);
nearest
}
#[allow(dead_code)]
pub fn kd_nearest_with<T, Scalar>(
kdtree: &[T],
dim: usize,
kd_difference: impl Fn(&T, usize) -> Scalar + Copy,
) -> ItemAndDistance<T, Scalar>
where
Scalar: num_traits::NumAssign + Copy + PartialOrd,
{
fn squared_distance<T, Scalar: num_traits::NumAssign + Copy>(
item: &T,
dim: usize,
kd_difference: impl Fn(&T, usize) -> Scalar + Copy,
) -> Scalar {
let mut squared_distance = Scalar::zero();
for k in 0..dim {
let diff = kd_difference(item, k);
squared_distance += diff * diff;
}
squared_distance
}
fn recurse<'a, T, Scalar>(
nearest: &mut ItemAndDistance<'a, T, Scalar>,
kdtree: &'a [T],
axis: usize,
dim: usize,
kd_difference: impl Fn(&T, usize) -> Scalar + Copy,
) where
Scalar: num_traits::NumAssign + Copy + PartialOrd,
{
let mid_idx = kdtree.len() / 2;
let mid = &kdtree[mid_idx];
let squared_distance = squared_distance(mid, dim, kd_difference);
if squared_distance < nearest.squared_distance {
*nearest = ItemAndDistance {
item: mid,
squared_distance,
};
if nearest.squared_distance.is_zero() {
return;
}
}
let [branch1, branch2] = if kd_difference(mid, axis) < Scalar::zero() {
[&kdtree[..mid_idx], &kdtree[mid_idx + 1..]]
} else {
[&kdtree[mid_idx + 1..], &kdtree[..mid_idx]]
};
if !branch1.is_empty() {
recurse(nearest, branch1, (axis + 1) % dim, dim, kd_difference);
}
if !branch2.is_empty() {
let diff = kd_difference(mid, axis);
if diff * diff < nearest.squared_distance {
recurse(nearest, branch2, (axis + 1) % dim, dim, kd_difference);
}
}
}
assert!(!kdtree.is_empty());
let mut nearest = ItemAndDistance {
item: &kdtree[0],
squared_distance: squared_distance(&kdtree[0], dim, kd_difference),
};
recurse(&mut nearest, kdtree, 0, dim, kd_difference);
nearest
}

82
vendor/kd-tree/src/nearests.rs vendored Normal file
View File

@@ -0,0 +1,82 @@
use crate::{ItemAndDistance, KdPoint};
pub fn kd_nearests<'a, T: KdPoint>(
kdtree: &'a [T],
query: &impl KdPoint<Scalar = T::Scalar, Dim = T::Dim>,
num: usize,
) -> Vec<ItemAndDistance<'a, T, T::Scalar>> {
kd_nearests_by(kdtree, query, num, |item, k| item.at(k))
}
pub fn kd_nearests_by<'a, T, P: KdPoint>(
kdtree: &'a [T],
query: &P,
num: usize,
get: impl Fn(&T, usize) -> P::Scalar + Copy,
) -> Vec<ItemAndDistance<'a, T, P::Scalar>> {
fn distance_squared<P: KdPoint, T>(
p1: &P,
p2: &T,
get: impl Fn(&T, usize) -> P::Scalar,
) -> P::Scalar {
let mut squared_distance = <P::Scalar as num_traits::Zero>::zero();
for i in 0..P::dim() {
let diff = p1.at(i) - get(p2, i);
squared_distance += diff * diff;
}
squared_distance
}
fn recurse<'a, T, Q: KdPoint>(
nearests: &mut Vec<ItemAndDistance<'a, T, Q::Scalar>>,
kdtree: &'a [T],
get: impl Fn(&T, usize) -> Q::Scalar + Copy,
query: &Q,
axis: usize,
) {
let mid_idx = kdtree.len() / 2;
let item = &kdtree[mid_idx];
let squared_distance = distance_squared(query, item, get);
if nearests.len() < nearests.capacity()
|| squared_distance < nearests.last().unwrap().squared_distance
{
if nearests.len() == nearests.capacity() {
nearests.pop();
}
let i = nearests
.binary_search_by(|item| {
item.squared_distance
.partial_cmp(&squared_distance)
.unwrap_or(std::cmp::Ordering::Equal)
})
.unwrap_or_else(|i| i);
nearests.insert(
i,
ItemAndDistance {
item,
squared_distance,
},
);
}
let mid_pos = get(item, axis);
let [branch1, branch2] = if query.at(axis) < mid_pos {
[&kdtree[..mid_idx], &kdtree[mid_idx + 1..]]
} else {
[&kdtree[mid_idx + 1..], &kdtree[..mid_idx]]
};
if !branch1.is_empty() {
recurse(nearests, branch1, get, query, (axis + 1) % Q::dim());
}
if !branch2.is_empty() {
let diff = query.at(axis) - mid_pos;
if diff * diff < nearests.last().unwrap().squared_distance {
recurse(nearests, branch2, get, query, (axis + 1) % Q::dim());
}
}
}
if num == 0 || kdtree.is_empty() {
return Vec::new();
}
let mut nearests = Vec::with_capacity(num);
recurse(&mut nearests, kdtree, get, query, 0);
nearests
}

49
vendor/kd-tree/src/sort.rs vendored Normal file
View File

@@ -0,0 +1,49 @@
use std::cmp::Ordering;
pub fn kd_sort_by<T>(
items: &mut [T],
dim: usize,
kd_compare: impl Fn(&T, &T, usize) -> Ordering + Copy,
) {
fn recurse<T>(
items: &mut [T],
axis: usize,
dim: usize,
kd_compare: impl Fn(&T, &T, usize) -> Ordering + Copy,
) {
if items.len() >= 2 {
items.select_nth_unstable_by(items.len() / 2, |x, y| kd_compare(x, y, axis));
let mid = items.len() / 2;
let axis = (axis + 1) % dim;
recurse(&mut items[..mid], axis, dim, kd_compare);
recurse(&mut items[mid + 1..], axis, dim, kd_compare);
}
}
recurse(items, 0, dim, kd_compare);
}
#[cfg(feature = "rayon")]
pub fn kd_par_sort_by<T: Send>(
items: &mut [T],
dim: usize,
kd_compare: impl Fn(&T, &T, usize) -> Ordering + Copy + Send,
) {
fn recurse<T: Send>(
items: &mut [T],
axis: usize,
dim: usize,
kd_compare: impl Fn(&T, &T, usize) -> Ordering + Copy + Send,
) {
if items.len() >= 2 {
items.select_nth_unstable_by(items.len() / 2, |x, y| kd_compare(x, y, axis));
let mid = items.len() / 2;
let axis = (axis + 1) % dim;
let (lhs, rhs) = items.split_at_mut(mid);
rayon::join(
move || recurse(lhs, axis, dim, kd_compare),
move || recurse(&mut rhs[1..], axis, dim, kd_compare),
);
}
}
recurse(items, 0, dim, kd_compare);
}

207
vendor/kd-tree/src/tests.rs vendored Normal file
View File

@@ -0,0 +1,207 @@
#![cfg(test)]
use super::*;
#[test]
fn test_nearest() {
let mut gen3d = random3d_generator();
let kdtree = KdTree::build_by_ordered_float(vec(10000, |_| gen3d()));
for _ in 0..100 {
let query = gen3d();
let found = kdtree.nearest(&query).unwrap().item;
let expected = kdtree
.iter()
.min_by_key(|p| ordered_float::OrderedFloat(squared_distance(p, &query)))
.unwrap();
assert_eq!(found, expected);
}
}
#[test]
fn test_nearests() {
test_nearests_by(random3d_generator());
test_nearests_by(random3d_10th_generator());
}
fn test_nearests_by(mut gen3d: impl FnMut() -> [f64; 3]) {
let kdtree = KdTree::build_by_ordered_float(vec(10000, |_| gen3d()));
const NUM: usize = 5;
for _ in 0..100 {
let query = gen3d();
let neighborhood = kdtree.nearests(&query, NUM);
assert_eq!(neighborhood.len(), NUM);
for i in 1..neighborhood.len() {
assert!(neighborhood[i - 1].squared_distance <= neighborhood[i].squared_distance);
}
let neighborhood_radius = neighborhood
.iter()
.max_by_key(|entry| ordered_float::OrderedFloat(entry.squared_distance))
.unwrap()
.squared_distance;
let neighborhood_contains = |p: &[f64; 3]| {
neighborhood
.iter()
.any(|entry| std::ptr::eq(entry.item as _, p as _))
};
assert!(kdtree.iter().all(
|p| neighborhood_contains(p) || neighborhood_radius <= squared_distance(p, &query)
));
}
}
#[test]
fn test_within() {
let mut gen3d = random3d_generator();
let kdtree = KdTree::build_by_ordered_float(vec(10000, |_| gen3d()));
for _ in 0..100 {
let mut p1 = gen3d();
let mut p2 = gen3d();
for k in 0..3 {
if p1[k] > p2[k] {
std::mem::swap(&mut p1[k], &mut p2[k]);
}
}
let found = kdtree.within(&[p1, p2]);
let count = kdtree
.iter()
.filter(|p| (0..3).all(|k| p1[k] <= p[k] && p[k] <= p2[k]))
.count();
assert_eq!(found.len(), count);
}
}
#[test]
fn test_within_against_empty() {
let empty: KdTree<[f64; 3]> = KdTree::build_by_ordered_float(vec![]);
assert!(empty.within(&[[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]).is_empty());
}
#[test]
fn test_within_radius() {
let mut gen3d = random3d_generator();
let kdtree = KdTree::build_by_ordered_float(vec(10000, |_| gen3d()));
const RADIUS: f64 = 0.1;
for _ in 0..100 {
let query = gen3d();
let found = kdtree.within_radius(&query, RADIUS);
let count = kdtree
.iter()
.filter(|p| squared_distance(p, &query) < RADIUS * RADIUS)
.count();
assert_eq!(found.len(), count);
}
}
fn squared_distance<T: num_traits::Num + Copy>(p1: &[T; 3], p2: &[T; 3]) -> T {
let dx = p1[0] - p2[0];
let dy = p1[1] - p2[1];
let dz = p1[2] - p2[2];
dx * dx + dy * dy + dz * dz
}
fn random3d_generator() -> impl FnMut() -> [f64; 3] {
use rand::Rng;
let mut rng = rand::rng();
move || [rng.random(), rng.random(), rng.random()]
}
fn random3d_10th_generator() -> impl FnMut() -> [f64; 3] {
// generates a random number between 0 and 1 with 0.1 step
fn random_10th(rng: &mut impl rand::Rng) -> f64 {
f64::from(rng.random_range(0u8..=10u8)) / 10.0
}
let mut rng = rand::rng();
move || {
[
random_10th(&mut rng),
random_10th(&mut rng),
random_10th(&mut rng),
]
}
}
fn vec<T>(count: usize, mut f: impl FnMut(usize) -> T) -> Vec<T> {
let mut items = Vec::with_capacity(count);
for i in 0..count {
items.push(f(i));
}
items
}
#[cfg(feature = "nalgebra")]
#[test]
fn test_nalgebra_point() {
use ::nalgebra as na;
let mut gen3d = random3d_generator();
let kdtree: KdTree<na::Point3<f64>> =
KdTree::build_by_ordered_float(vec(10000, |_| gen3d().into()));
for _ in 0..100 {
let query: na::Point3<f64> = gen3d().into();
let found = kdtree.nearest(&query).unwrap().item;
let expected = kdtree
.iter()
.min_by_key(|p| ordered_float::OrderedFloat((query - **p).norm()))
.unwrap();
assert_eq!(found, expected);
}
}
#[cfg(feature = "nalgebra")]
#[test]
fn test_nalgebra_vector() {
use ::nalgebra as na;
let mut gen3d = random3d_generator();
let kdtree: KdTree<na::Vector3<f64>> =
KdTree::build_by_ordered_float(vec(10000, |_| gen3d().into()));
for _ in 0..100 {
let query: na::Vector3<f64> = gen3d().into();
let found = kdtree.nearest(&query).unwrap().item;
let expected = kdtree
.iter()
.min_by_key(|p| ordered_float::OrderedFloat((query - **p).norm()))
.unwrap();
assert_eq!(found, expected);
}
}
#[cfg(feature = "serde")]
#[test]
fn test_serde() {
let src: KdTree3<[i32; 3]> = KdTree::build(vec![[1, 2, 3], [4, 5, 6]]);
let json = serde_json::to_string(&src).unwrap();
assert_eq!(json, "[[1,2,3],[4,5,6]]");
let dst: KdTree3<[i32; 3]> = serde_json::from_str(&json).unwrap();
assert_eq!(src, dst);
}
#[cfg(feature = "nalgebra-serde")]
#[test]
fn test_nalgebra_serde() {
use ::nalgebra as na;
let src: KdTree<na::Point3<f64>> = KdTree::build_by_ordered_float(vec![
na::Point3::new(1.0, 2.0, 3.0),
na::Point3::new(4.0, 5.0, 6.0),
]);
let json = serde_json::to_string(&src).unwrap();
assert_eq!(json, "[[1.0,2.0,3.0],[4.0,5.0,6.0]]");
let dst: KdTree3<na::Point3<f64>> = serde_json::from_str(&json).unwrap();
assert_eq!(src, dst);
}
#[cfg(feature = "rayon")]
#[test]
fn test_rayon() {
let points = {
let mut gen3d = random3d_generator();
vec(1000, |_| gen3d())
};
let kdtree1 = KdTree::build_by_ordered_float(points.clone());
let kdtree2 = KdTree::par_build_by_ordered_float(points.clone());
assert_eq!(kdtree1, kdtree2);
}

42
vendor/kd-tree/src/within.rs vendored Normal file
View File

@@ -0,0 +1,42 @@
use std::cmp::Ordering;
pub fn kd_within_by_cmp<T>(
kdtree: &[T],
dim: usize,
compare: impl Fn(&T, usize) -> Ordering + Copy,
) -> Vec<&T> {
fn recurse<'a, T>(
results: &mut Vec<&'a T>,
kdtree: &'a [T],
axis: usize,
dim: usize,
compare: impl Fn(&T, usize) -> Ordering + Copy,
) {
if kdtree.is_empty() {
return;
}
let axis = axis % dim;
let (lower, item, upper) = {
let mid = kdtree.len() / 2;
(&kdtree[..mid], &kdtree[mid], &kdtree[mid + 1..])
};
match compare(item, axis) {
Ordering::Equal => {
if (1..dim).all(|k| compare(item, (axis + k) % dim) == Ordering::Equal) {
results.push(item);
}
recurse(results, lower, axis + 1, dim, compare);
recurse(results, upper, axis + 1, dim, compare);
}
Ordering::Less => {
recurse(results, upper, axis + 1, dim, compare);
}
Ordering::Greater => {
recurse(results, lower, axis + 1, dim, compare);
}
}
}
let mut results = Vec::new();
recurse(&mut results, kdtree, 0, dim, compare);
results
}