use crate::shared::grow::{
any_cardinal_edge_ok, choose_unambiguous, collect_candidates, collect_labelled_neighbours,
enqueue_cardinal_neighbours, is_extrapolating, predict_from_neighbours, CandidateChoice,
GrowParams, GrowResult, SquareAttachPolicy,
};
use kiddo::KdTree;
use nalgebra::Point2;
use std::collections::{HashSet, VecDeque};
#[non_exhaustive]
#[derive(Clone, Debug, Default)]
pub struct BfsExtensionStats {
pub attached: usize,
pub rejected_no_candidate: usize,
pub rejected_ambiguous: usize,
pub rejected_edge: usize,
pub attached_indices: Vec<usize>,
pub attached_cells: Vec<(i32, i32)>,
}
#[cfg_attr(
feature = "tracing",
tracing::instrument(
level = "info",
skip_all,
fields(num_corners = positions.len(), num_labelled = grow.labelled.len(), cell_size = cell_size),
)
)]
pub fn extend_from_labelled<V: SquareAttachPolicy>(
positions: &[Point2<f32>],
grow: &mut GrowResult,
cell_size: f32,
params: &GrowParams,
policy: &V,
) -> BfsExtensionStats {
let mut tree: KdTree<f32, 2> = KdTree::new();
let mut tree_slot_to_corner: Vec<usize> = Vec::new();
for (idx, pos) in positions.iter().enumerate() {
if policy.is_eligible(idx) && !grow.by_corner.contains_key(&idx) {
tree.add(&[pos.x, pos.y], tree_slot_to_corner.len() as u64);
tree_slot_to_corner.push(idx);
}
}
let mut boundary: VecDeque<(i32, i32)> = VecDeque::new();
let mut seen_boundary: HashSet<(i32, i32)> = HashSet::new();
for &pos in grow.labelled.keys() {
enqueue_cardinal_neighbours(pos, &grow.labelled, &mut boundary, &mut seen_boundary);
}
let mut stats = BfsExtensionStats::default();
while let Some(pos) = boundary.pop_front() {
if grow.labelled.contains_key(&pos) {
continue;
}
let neighbours = collect_labelled_neighbours(pos, 1, &grow.labelled, positions);
if neighbours.is_empty() {
stats.rejected_no_candidate += 1;
grow.holes.insert(pos);
continue;
}
let prediction = predict_from_neighbours(
pos,
&neighbours,
grow.axis_i,
grow.axis_j,
cell_size,
&grow.labelled,
positions,
);
let search_r = params.attach_search_rel * cell_size;
let extrapolating = is_extrapolating(pos, &neighbours);
let local_search_r = if extrapolating {
search_r * params.boundary_search_factor
} else {
search_r
};
let required_label = policy.required_label_at(pos.0, pos.1);
let candidates = collect_candidates(
&tree,
&tree_slot_to_corner,
prediction,
local_search_r,
policy,
required_label,
&grow.by_corner,
);
let choice = choose_unambiguous(
&candidates,
params.attach_ambiguity_factor,
prediction,
positions,
policy,
pos,
&neighbours,
);
match choice {
CandidateChoice::None => {
stats.rejected_no_candidate += 1;
grow.holes.insert(pos);
}
CandidateChoice::Ambiguous => {
stats.rejected_ambiguous += 1;
grow.ambiguous.insert(pos);
}
CandidateChoice::Unique(c_idx) => {
if !any_cardinal_edge_ok(c_idx, pos, &grow.labelled, policy) {
stats.rejected_edge += 1;
grow.holes.insert(pos);
} else {
grow.labelled.insert(pos, c_idx);
grow.by_corner.insert(c_idx, pos);
enqueue_cardinal_neighbours(
pos,
&grow.labelled,
&mut boundary,
&mut seen_boundary,
);
stats.attached += 1;
stats.attached_indices.push(c_idx);
stats.attached_cells.push(pos);
}
}
}
}
stats
}