use super::*;
use crate::dist::SquaredEuclidean;
use crate::leaf_strategy::FlatVec;
use crate::leaf_strategy::VecOfArenas;
use crate::leaf_strategy::VecOfArrays;
use crate::Donnelly;
use crate::Eytzinger;
#[test]
fn construction_index_selection_is_adaptive() {
assert!(construction_index_fits_u32(0));
assert!(construction_index_fits_u32(u32::MAX as usize));
#[cfg(target_pointer_width = "64")]
assert!(!construction_index_fits_u32(u32::MAX as usize + 1));
}
#[test]
fn update_pivot_shifts_right_when_left_scan_hits_zero() {
type TestTree = KdTree<f32, u32, Eytzinger, FlatVec<f32, u32, 2, 32>, 2, 32>;
let source = [
[1.0f32, 10.0],
[1.0, 20.0],
[1.0, 30.0],
[2.0, 40.0],
[3.0, 50.0],
];
let mut sort_index = [0u32, 1, 2, 3, 4];
let pivot = TestTree::update_pivot(
&source,
&|point: &[f32; 2], dim| point[dim],
&mut sort_index,
0,
1,
)
.unwrap();
assert_eq!(pivot, 3);
assert_eq!(sort_index, [0, 1, 2, 3, 4]);
assert_eq!(source[sort_index[pivot - 1].as_usize()][0], 1.0);
assert_eq!(source[sort_index[pivot].as_usize()][0], 2.0);
}
#[test]
fn replace_item_updates_flat_vec_tree_without_changing_size() {
type TestTree = KdTree<f32, u32, Eytzinger, FlatVec<f32, u32, 2, 32>, 2, 32>;
let entries = [
(10u32, [1.0f32, 10.0]),
(11u32, [2.0, 20.0]),
(12u32, [1.0, 10.0]),
];
let mut tree = TestTree::new_from_entries(&entries).unwrap();
assert_eq!(tree.size(), 3);
tree.replace_item(&[1.0, 10.0], 10, 99).unwrap();
assert_eq!(tree.size(), 3);
let iterated = tree.iter().collect::<Vec<_>>();
assert_eq!(iterated[0], (99, [1.0, 10.0]));
assert_eq!(iterated[1], (11, [2.0, 20.0]));
assert_eq!(iterated[2], (12, [1.0, 10.0]));
}
#[test]
fn replace_item_returns_entry_not_found_when_exact_match_is_missing() {
type TestTree = KdTree<f32, u32, Eytzinger, VecOfArrays<f32, u32, 2, 32>, 2, 32>;
let entries = [(10u32, [1.0f32, 10.0]), (11u32, [2.0, 20.0])];
let mut tree = TestTree::new_from_entries(&entries).unwrap();
assert_eq!(
tree.replace_item(&[1.0, 10.0], 99, 100),
Err(MutationError::EntryNotFound)
);
assert_eq!(
tree.replace_item(&[9.0, 90.0], 10, 100),
Err(MutationError::EntryNotFound)
);
}
#[test]
fn replace_item_updates_vec_of_arenas_tree() {
type TestTree = KdTree<f64, u32, Eytzinger, VecOfArenas<f64, u32, 2, 32>, 2, 32>;
let entries = [
(20u32, [1.0f64, 10.0]),
(21u32, [2.0, 20.0]),
(22u32, [3.0, 30.0]),
];
let mut tree = TestTree::new_from_entries(&entries).unwrap();
tree.replace_item(&[2.0, 20.0], 21, 77).unwrap();
let iterated = tree.iter().collect::<Vec<_>>();
assert_eq!(
iterated,
vec![(20, [1.0, 10.0]), (77, [2.0, 20.0]), (22, [3.0, 30.0])]
);
}
#[test]
fn irregular_immutable_soft_layout_preserves_arithmetic_resolution() {
type TestTree = KdTree<f32, u32, Eytzinger, FlatVec<f32, u32, 2, 2>, 2, 2>;
let points = vec![
[3.0, 0.0],
[1.0, 0.6],
[1.0, 1.4],
[3.0, 3.3],
[3.0, 3.8],
[0.0, 1.8],
[3.0, 1.5],
[3.0, 2.7],
[1.0, 3.3],
];
let query = [2.9142656, 5.220647];
let tree = TestTree::new_from_slice(&points).unwrap();
assert!(tree.stem_leaf_resolution.uses_arithmetic());
assert_eq!(tree.leaf_count(), 8);
assert_eq!(tree.max_leaf_len(), 3);
assert_eq!(
(0..tree.leaf_count())
.map(|leaf_idx| {
<FlatVec<f32, u32, 2, 2> as LeafStrategy<f32, u32, Eytzinger, 2, 2>>::leaf_len(
&tree.leaves,
leaf_idx,
)
})
.collect::<Vec<_>>(),
vec![2, 0, 1, 1, 3, 0, 2, 0]
);
let result = tree
.query(&query)
.nearest_one::<SquaredEuclidean<f32>>()
.execute();
assert_eq!(result.item, 4);
assert!((result.distance - 2.025588).abs() < 1.0e-6);
}
#[test]
fn irregular_hard_terminal_layout_is_detected_and_mapped() {
type TestTree = KdTree<f32, u32, Eytzinger, VecOfArrays<f32, u32, 2, 2>, 2, 2>;
let terminal_stem_indices = vec![8usize, 10, 3];
assert!(!TestTree::terminal_stem_indices_match_arithmetic_layout(
&terminal_stem_indices,
2,
));
let stem_leaf_resolution =
TestTree::mapped_stem_leaf_resolution_from_terminals(&terminal_stem_indices);
assert!(!stem_leaf_resolution.uses_arithmetic());
assert_eq!(stem_leaf_resolution.resolve_terminal_stem_idx(8, 0), 0);
assert_eq!(stem_leaf_resolution.resolve_terminal_stem_idx(10, 0), 1);
assert_eq!(stem_leaf_resolution.resolve_terminal_stem_idx(3, 0), 2);
}
#[test]
fn unsplittable_immutable_hard_bucket_returns_error() {
type TestTree = KdTree<f32, u32, Eytzinger, VecOfArrays<f32, u32, 2, 2>, 2, 2>;
let points = vec![[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]];
assert!(matches!(
TestTree::new_from_slice(&points),
Err(ConstructionError::UnsplittableBucket { split_dim: 0 })
));
}
#[cfg(feature = "multi-threaded")]
#[test]
fn mutable_split_preserves_point_item_associations_when_pivot_scan_retries() {
type TestTree = KdTree<f32, u32, Eytzinger, VecOfArrays<f32, u32, 2, 16>, 2, 16>;
const EXTREMES: [[f32; 2]; 6] = [
[-1.0, -1.0],
[1.0, 1.0],
[-1.0, 1.0],
[1.0, -1.0],
[0.0, -0.0],
[-0.0, 0.0],
];
let mut tree = TestTree::default();
let mut expected = Vec::new();
for item in 0..33u32 {
let point = EXTREMES[item as usize % EXTREMES.len()];
tree.add(&point, item).unwrap();
expected.push((item, point.map(f32::to_bits)));
}
let final_point = [1.0, 4.0 / 9.0];
tree.add(&final_point, 33).unwrap();
expected.push((33, final_point.map(f32::to_bits)));
let mut got = tree
.iter()
.map(|(item, point)| (item, point.map(f32::to_bits)))
.collect::<Vec<_>>();
got.sort_unstable();
assert_eq!(got, expected);
}
#[test]
fn rejected_mutable_split_preserves_existing_point_item_associations() {
type TestTree = KdTree<f32, u32, Eytzinger, VecOfArrays<f32, u32, 2, 4>, 2, 4>;
let mut tree = TestTree::default();
let mut expected = Vec::new();
for item in 0..4u32 {
let point = [2.0, 200.0 + item as f32];
tree.add(&point, item).unwrap();
expected.push((item, point));
}
assert_eq!(
tree.add(&[2.0, 204.0], 4),
Err(ConstructionError::UnsplittableBucket { split_dim: 0 })
);
assert_eq!(tree.size(), expected.len());
let mut got = tree.iter().collect::<Vec<_>>();
got.sort_unstable_by_key(|(item, _)| *item);
assert_eq!(got, expected);
}
#[test]
fn parallel_construction_threshold_is_inclusive() {
type TestTree = KdTree<f32, u32, Eytzinger, FlatVec<f32, u32, 2, 32>, 2, 32>;
let policy = ParallelConstruction::with_threshold(1_024);
let default = TestTree::builder();
let forced = TestTree::builder().with_parallel_construction();
let zero_threshold = TestTree::builder().with_parallel_construction_threshold(0);
assert!(!policy.should_parallelize(1_023));
assert!(policy.should_parallelize(1_024));
assert!(policy.should_parallelize(1_025));
assert_eq!(
default.policy.threshold(),
DEFAULT_PARALLEL_CONSTRUCTION_THRESHOLD
);
assert_eq!(forced.policy.threshold(), 1);
assert_eq!(zero_threshold.policy.threshold(), 1);
assert_eq!(DEFAULT_PARALLEL_CONSTRUCTION_THRESHOLD, 262_144);
}
#[cfg(feature = "multi-threaded")]
#[test]
fn parallel_soft_construction_matches_sequential_construction() {
type FlatTree = KdTree<f32, u32, Eytzinger, FlatVec<f32, u32, 2, 32>, 2, 32>;
type ArenaTree = KdTree<f32, u32, Eytzinger, VecOfArenas<f32, u32, 2, 32>, 2, 32>;
let points = (0..4_096)
.map(|idx| {
[
((idx * 17) % 997) as f32,
((idx * 53 + idx / 7) % 991) as f32,
]
})
.collect::<Vec<_>>();
macro_rules! assert_parallel_matches {
($tree:ty) => {{
let sequential = <$tree>::builder()
.with_serial_construction()
.build_from_slice(&points)
.unwrap();
let parallel = <$tree>::builder()
.with_parallel_construction()
.build_from_slice(&points)
.unwrap();
let threshold_parallel = <$tree>::builder()
.with_parallel_construction_threshold(points.len())
.build_from_slice(&points)
.unwrap();
assert_eq!(sequential.stems.as_slice(), parallel.stems.as_slice());
assert_eq!(
sequential.stems.as_slice(),
threshold_parallel.stems.as_slice()
);
assert_eq!(sequential.size(), parallel.size());
assert_eq!(sequential.leaf_count(), parallel.leaf_count());
assert_eq!(sequential.max_leaf_len(), parallel.max_leaf_len());
assert_eq!(
sequential.iter().collect::<Vec<_>>(),
parallel.iter().collect::<Vec<_>>()
);
for query in [[0.0, 0.0], [500.0, 500.0], [996.0, 990.0]] {
assert_eq!(
sequential
.query(&query)
.nearest_one::<SquaredEuclidean<f32>>()
.execute(),
parallel
.query(&query)
.nearest_one::<SquaredEuclidean<f32>>()
.execute()
);
}
}};
}
assert_parallel_matches!(FlatTree);
assert_parallel_matches!(ArenaTree);
}
#[cfg(feature = "multi-threaded")]
#[test]
fn parallel_soft_construction_handles_scalar_donnelly_partial_block() {
type TestTree = KdTree<f64, u32, Donnelly<3>, FlatVec<f64, u32, 2, 4>, 2, 4>;
let points = (0..64)
.map(|idx| [((idx * 17) % 67) as f64, ((idx * 29 + idx / 5) % 71) as f64])
.collect::<Vec<_>>();
let sequential = TestTree::builder()
.with_serial_construction()
.build_from_slice(&points)
.unwrap();
let parallel = TestTree::builder()
.with_parallel_construction()
.build_from_slice(&points)
.unwrap();
assert_eq!(sequential.stems.as_slice(), parallel.stems.as_slice());
assert_eq!(sequential.leaf_count(), 16);
assert_eq!(sequential.leaf_count(), parallel.leaf_count());
assert_eq!(sequential.max_stem_level(), 3);
assert_eq!(sequential.max_stem_level(), parallel.max_stem_level());
assert_eq!(sequential.max_leaf_len(), parallel.max_leaf_len());
assert_eq!(
sequential.iter().collect::<Vec<_>>(),
parallel.iter().collect::<Vec<_>>()
);
for query in [[0.0, 0.0], [31.0, 37.0], [66.0, 70.0]] {
assert_eq!(
sequential
.query(&query)
.nearest_one::<SquaredEuclidean<f64>>()
.execute(),
parallel
.query(&query)
.nearest_one::<SquaredEuclidean<f64>>()
.execute()
);
}
}
#[cfg(feature = "multi-threaded")]
#[test]
fn parallel_constructor_preserves_hard_bucket_construction() {
type TestTree = KdTree<f32, u32, Eytzinger, VecOfArrays<f32, u32, 2, 32>, 2, 32>;
let points = (0..4_096)
.map(|idx| [idx as f32, ((idx * 31) % 4_099) as f32])
.collect::<Vec<_>>();
let sequential = TestTree::new_from_slice(&points).unwrap();
let parallel = TestTree::builder()
.with_parallel_construction()
.build_from_slice(&points)
.unwrap();
assert_eq!(sequential.stems.as_slice(), parallel.stems.as_slice());
assert_eq!(
sequential.iter().collect::<Vec<_>>(),
parallel.iter().collect::<Vec<_>>()
);
}
#[cfg(feature = "multi-threaded")]
#[test]
fn builder_supports_parallel_entries_sources_and_no_items() {
type ItemTree = KdTree<f32, u32, Eytzinger, FlatVec<f32, u32, 2, 32>, 2, 32>;
type NoItemsTree = KdTree<f32, (), Eytzinger, FlatVec<f32, (), 2, 32>, 2, 32>;
let entries = (0..2_048)
.map(|idx| (idx as u32 + 10, [idx as f32, (idx % 127) as f32]))
.collect::<Vec<_>>();
let from_entries = ItemTree::builder()
.with_parallel_construction()
.build_from_entries(&entries)
.unwrap();
let from_source = ItemTree::builder()
.with_parallel_construction()
.build_from_source(
&entries,
|entry, dim| entry.1[dim],
|_src_idx, entry| entry.0,
)
.unwrap();
assert_eq!(
from_entries.iter().collect::<Vec<_>>(),
from_source.iter().collect::<Vec<_>>()
);
let points = entries.iter().map(|entry| entry.1).collect::<Vec<_>>();
let no_items = NoItemsTree::builder()
.with_parallel_construction()
.build_from_slice_no_items(&points)
.unwrap();
assert_eq!(no_items.size(), points.len());
assert_eq!(no_items.iter().count(), points.len());
}
#[test]
fn serial_builder_supports_non_sync_sources() {
use std::cell::Cell;
struct Point {
coords: [Cell<f32>; 2],
}
type TestTree = KdTree<f32, u32, Eytzinger, FlatVec<f32, u32, 2, 32>, 2, 32>;
let points = [
Point {
coords: [Cell::new(1.0), Cell::new(2.0)],
},
Point {
coords: [Cell::new(3.0), Cell::new(4.0)],
},
];
let tree = TestTree::builder()
.with_serial_construction()
.build_from_source(
&points,
|point, dim| point.coords[dim].get(),
|idx, _point| idx as u32,
)
.unwrap();
assert_eq!(tree.size(), points.len());
}