use crate::shared::grow::{
collect_labelled_neighbours, predict_from_neighbours, Admit, FillEdgeCtx, GrowResult,
SquareAttachPolicy,
};
use kiddo::{KdTree, SquaredEuclidean};
use nalgebra::Point2;
#[non_exhaustive]
#[derive(Clone, Copy, Debug)]
pub struct FillParams {
pub attach_search_rel: f32,
pub attach_ambiguity_factor: f32,
pub max_iters: usize,
}
impl Default for FillParams {
fn default() -> Self {
Self {
attach_search_rel: 0.35,
attach_ambiguity_factor: 1.5,
max_iters: 1,
}
}
}
impl FillParams {
pub fn new(attach_search_rel: f32, attach_ambiguity_factor: f32, max_iters: usize) -> Self {
Self {
attach_search_rel,
attach_ambiguity_factor,
max_iters,
}
}
}
#[non_exhaustive]
#[derive(Clone, Debug, Default)]
pub struct FillStats {
pub added: usize,
pub iterations: usize,
pub attached_indices: Vec<usize>,
pub attached_cells: Vec<(i32, i32)>,
}
pub fn fill_grid_holes<V: SquareAttachPolicy>(
positions: &[Point2<f32>],
grow: &mut GrowResult,
cell_size: f32,
params: &FillParams,
policy: &V,
) -> FillStats {
let mut stats = FillStats::default();
if grow.labelled.is_empty() {
return stats;
}
for _iter in 0..params.max_iters.max(1) {
stats.iterations += 1;
let (tree, slot_to_corner) = build_fill_tree(positions, grow, policy);
let ctx = FillCtx {
positions,
cell_size,
params,
tree: &tree,
slot_to_corner: &slot_to_corner,
policy,
};
let cells = enumerate_fill_cells(grow);
let mut added_this_iter = 0usize;
for cell in cells {
if grow.labelled.contains_key(&cell) {
continue;
}
if let Some(attached_idx) = try_fill_cell(cell, grow, &ctx) {
stats.attached_indices.push(attached_idx);
stats.attached_cells.push(cell);
added_this_iter += 1;
}
}
stats.added += added_this_iter;
if added_this_iter == 0 {
break;
}
}
stats
}
fn build_fill_tree<V: SquareAttachPolicy>(
positions: &[Point2<f32>],
grow: &GrowResult,
policy: &V,
) -> (KdTree<f32, 2>, Vec<usize>) {
let mut tree: KdTree<f32, 2> = KdTree::new();
let mut slot_to_corner: Vec<usize> = Vec::new();
for (idx, pos) in positions.iter().enumerate() {
if policy.eligible_for_fill(idx) && !grow.by_corner.contains_key(&idx) {
tree.add(&[pos.x, pos.y], slot_to_corner.len() as u64);
slot_to_corner.push(idx);
}
}
(tree, slot_to_corner)
}
fn enumerate_fill_cells(grow: &GrowResult) -> Vec<(i32, i32)> {
use std::collections::HashSet;
let mut out: HashSet<(i32, i32)> = HashSet::new();
let (mut min_i, mut max_i, mut min_j, mut max_j) = (i32::MAX, i32::MIN, i32::MAX, i32::MIN);
for &(i, j) in grow.labelled.keys() {
min_i = min_i.min(i);
max_i = max_i.max(i);
min_j = min_j.min(j);
max_j = max_j.max(j);
}
for j in min_j..=max_j {
for i in min_i..=max_i {
if !grow.labelled.contains_key(&(i, j)) {
out.insert((i, j));
}
}
}
for j in min_j..=max_j {
out.insert((min_i - 1, j));
out.insert((max_i + 1, j));
}
for i in min_i..=max_i {
out.insert((i, min_j - 1));
out.insert((i, max_j + 1));
}
let mut cells: Vec<(i32, i32)> = out.into_iter().collect();
cells.sort_unstable();
cells
}
struct FillCtx<'a, V: SquareAttachPolicy> {
positions: &'a [Point2<f32>],
cell_size: f32,
params: &'a FillParams,
tree: &'a KdTree<f32, 2>,
slot_to_corner: &'a [usize],
policy: &'a V,
}
fn try_fill_cell<V: SquareAttachPolicy>(
cell: (i32, i32),
grow: &mut GrowResult,
ctx: &FillCtx<'_, V>,
) -> Option<usize> {
let positions = ctx.positions;
let cell_size = ctx.cell_size;
let params = ctx.params;
let tree = ctx.tree;
let slot_to_corner = ctx.slot_to_corner;
let policy = ctx.policy;
let neighbours = collect_labelled_neighbours(cell, 1, &grow.labelled, positions);
if neighbours.is_empty() {
return None;
}
let pred = predict_from_neighbours(
cell,
&neighbours,
grow.axis_i,
grow.axis_j,
cell_size,
&grow.labelled,
positions,
);
let required_label = policy.required_label_at(cell.0, cell.1);
let search_r = params.attach_search_rel * cell_size;
let r2 = search_r * search_r;
let mut hits: Vec<(usize, f32)> = Vec::new();
for nn in tree
.within_unsorted::<SquaredEuclidean>(&[pred.x, pred.y], r2)
.into_iter()
{
let slot = nn.item as usize;
let idx = slot_to_corner[slot];
if grow.by_corner.contains_key(&idx) {
continue;
}
if let Some(req) = required_label {
if policy.label_of(idx) != Some(req) {
continue;
}
}
if matches!(
policy.accept_candidate(idx, cell, pred, &neighbours),
Admit::Reject
) {
continue;
}
hits.push((idx, nn.distance.sqrt()));
}
hits.sort_by(|a, b| a.1.total_cmp(&b.1));
let candidate_idx = match hits.len() {
0 => return None,
1 => hits[0].0,
_ => {
let d0 = hits[0].1.max(f32::EPSILON);
let d1 = hits[1].1;
if d1 / d0 < params.attach_ambiguity_factor {
return None;
}
hits[0].0
}
};
if !any_cardinal_fill_edge_ok(
candidate_idx,
cell,
positions,
&grow.labelled,
policy,
cell_size,
) {
return None;
}
grow.labelled.insert(cell, candidate_idx);
grow.by_corner.insert(candidate_idx, cell);
Some(candidate_idx)
}
fn any_cardinal_fill_edge_ok<V: SquareAttachPolicy>(
c_idx: usize,
pos: (i32, i32),
positions: &[Point2<f32>],
labelled: &std::collections::HashMap<(i32, i32), usize>,
policy: &V,
cell_size: f32,
) -> bool {
let mut found_any = false;
for (di, dj) in [(1, 0), (-1, 0), (0, 1), (0, -1)] {
let neigh = (pos.0 + di, pos.1 + dj);
if let Some(&n_idx) = labelled.get(&neigh) {
found_any = true;
let ctx = FillEdgeCtx {
candidate_idx: c_idx,
neighbour_idx: n_idx,
at_candidate: pos,
at_neighbour: neigh,
labelled,
positions,
cell_size,
};
if policy.fill_edge_ok(ctx) {
return true;
}
}
}
!found_any
}
#[cfg(test)]
mod tests {
use super::*;
use crate::shared::grow::{Admit, LabelledNeighbour};
use std::collections::HashMap;
struct OpenValidator;
impl SquareAttachPolicy for OpenValidator {
fn is_eligible(&self, _idx: usize) -> bool {
true
}
fn required_label_at(&self, _i: i32, _j: i32) -> Option<u8> {
None
}
fn label_of(&self, _idx: usize) -> Option<u8> {
None
}
fn accept_candidate(
&self,
_idx: usize,
_at: (i32, i32),
_prediction: Point2<f32>,
_neighbours: &[LabelledNeighbour],
) -> Admit {
Admit::Accept
}
}
#[test]
fn fill_pass_attaches_interior_hole() {
let s = 20.0_f32;
let mut positions: Vec<Point2<f32>> = Vec::new();
for j in 0..3 {
for i in 0..3 {
positions.push(Point2::new(50.0 + i as f32 * s, 50.0 + j as f32 * s));
}
}
let mut labelled: HashMap<(i32, i32), usize> = HashMap::new();
let mut by_corner: HashMap<usize, (i32, i32)> = HashMap::new();
for j in 0..3 {
for i in 0..3 {
if (i, j) == (1, 1) {
continue;
}
let idx = (j * 3 + i) as usize;
labelled.insert((i, j), idx);
by_corner.insert(idx, (i, j));
}
}
let mut grow = GrowResult {
labelled,
by_corner,
ambiguous: Default::default(),
holes: Default::default(),
axis_i: nalgebra::Vector2::new(1.0, 0.0),
axis_j: nalgebra::Vector2::new(0.0, 1.0),
rebase_i_mod2: 0,
rebase_j_mod2: 0,
};
let stats = fill_grid_holes(
&positions,
&mut grow,
s,
&FillParams::default(),
&OpenValidator,
);
assert_eq!(stats.added, 1);
assert_eq!(grow.labelled.get(&(1, 1)), Some(&4));
assert_eq!(stats.attached_indices, vec![4]);
assert_eq!(stats.attached_cells, vec![(1, 1)]);
}
}