1use kiddo::{KdTree, SquaredEuclidean};
40use nalgebra::{Point2, Vector2};
41use std::collections::{HashMap, HashSet, VecDeque};
42
43pub use crate::square::seed::Seed;
44
45#[derive(Clone, Copy, Debug, PartialEq, Eq)]
47pub enum Admit {
48 Accept,
50 Reject,
53}
54
55#[derive(Clone, Copy, Debug)]
58pub struct LabelledNeighbour {
59 pub idx: usize,
60 pub at: (i32, i32),
61 pub position: Point2<f32>,
62}
63
64pub trait GrowValidator {
71 fn is_eligible(&self, idx: usize) -> bool;
74
75 fn required_label_at(&self, i: i32, j: i32) -> Option<u8>;
78
79 fn label_of(&self, idx: usize) -> Option<u8>;
83
84 fn accept_candidate(
89 &self,
90 idx: usize,
91 at: (i32, i32),
92 prediction: Point2<f32>,
93 neighbours: &[LabelledNeighbour],
94 ) -> Admit;
95
96 fn edge_ok(
104 &self,
105 _candidate_idx: usize,
106 _neighbour_idx: usize,
107 _at_candidate: (i32, i32),
108 _at_neighbour: (i32, i32),
109 ) -> bool {
110 true
111 }
112}
113
114#[non_exhaustive]
116#[derive(Clone, Copy, Debug)]
117pub struct GrowParams {
118 pub attach_search_rel: f32,
122 pub attach_ambiguity_factor: f32,
125 pub boundary_search_factor: f32,
133}
134
135impl Default for GrowParams {
136 fn default() -> Self {
137 Self {
138 attach_search_rel: 0.35,
139 attach_ambiguity_factor: 1.5,
140 boundary_search_factor: 2.0,
141 }
142 }
143}
144
145impl GrowParams {
146 pub fn new(attach_search_rel: f32, attach_ambiguity_factor: f32) -> Self {
147 Self {
148 attach_search_rel,
149 attach_ambiguity_factor,
150 ..Self::default()
151 }
152 }
153}
154
155#[derive(Debug, Default)]
157pub struct GrowResult {
158 pub labelled: HashMap<(i32, i32), usize>,
161 pub by_corner: HashMap<usize, (i32, i32)>,
163 pub ambiguous: HashSet<(i32, i32)>,
165 pub holes: HashSet<(i32, i32)>,
167 pub grid_u: Vector2<f32>,
169 pub grid_v: Vector2<f32>,
170}
171
172#[cfg_attr(
183 feature = "tracing",
184 tracing::instrument(
185 level = "info",
186 skip_all,
187 fields(num_corners = positions.len(), cell_size = cell_size),
188 )
189)]
190pub fn bfs_grow<V: GrowValidator>(
191 positions: &[Point2<f32>],
192 seed: Seed,
193 cell_size: f32,
194 params: &GrowParams,
195 validator: &V,
196) -> GrowResult {
197 let grid_u = {
199 let raw = positions[seed.b] - positions[seed.a];
200 let n = raw.norm().max(1e-6);
201 raw / n
202 };
203 let grid_v = {
204 let raw = positions[seed.c] - positions[seed.a];
205 let n = raw.norm().max(1e-6);
206 raw / n
207 };
208
209 let mut tree: KdTree<f32, 2> = KdTree::new();
211 let mut tree_slot_to_corner: Vec<usize> = Vec::new();
212 for (idx, pos) in positions.iter().enumerate() {
213 if validator.is_eligible(idx) {
214 tree.add(&[pos.x, pos.y], tree_slot_to_corner.len() as u64);
215 tree_slot_to_corner.push(idx);
216 }
217 }
218
219 let mut labelled: HashMap<(i32, i32), usize> = HashMap::new();
220 let mut by_corner: HashMap<usize, (i32, i32)> = HashMap::new();
221 let mut ambiguous: HashSet<(i32, i32)> = HashSet::new();
222 let mut holes: HashSet<(i32, i32)> = HashSet::new();
223
224 for (ij, idx) in [
225 ((0, 0), seed.a),
226 ((1, 0), seed.b),
227 ((0, 1), seed.c),
228 ((1, 1), seed.d),
229 ] {
230 labelled.insert(ij, idx);
231 by_corner.insert(idx, ij);
232 }
233
234 let mut boundary: VecDeque<(i32, i32)> = VecDeque::new();
235 let mut seen_boundary: HashSet<(i32, i32)> = HashSet::new();
236 for ij in labelled.keys().copied().collect::<Vec<_>>() {
237 enqueue_cardinal_neighbours(ij, &labelled, &mut boundary, &mut seen_boundary);
238 }
239
240 while let Some(pos) = boundary.pop_front() {
241 if labelled.contains_key(&pos) {
242 continue;
243 }
244 let (decision, _neighbours) = process_boundary_cell(
245 pos,
246 positions,
247 &labelled,
248 &by_corner,
249 &tree,
250 &tree_slot_to_corner,
251 grid_u,
252 grid_v,
253 cell_size,
254 params,
255 validator,
256 );
257 match decision {
258 BoundaryDecision::Hole | BoundaryDecision::EdgeRejected => {
259 holes.insert(pos);
260 }
261 BoundaryDecision::Ambiguous => {
262 ambiguous.insert(pos);
263 }
264 BoundaryDecision::Attach(c_idx) => {
265 labelled.insert(pos, c_idx);
266 by_corner.insert(c_idx, pos);
267 enqueue_cardinal_neighbours(pos, &labelled, &mut boundary, &mut seen_boundary);
268 }
269 }
270 }
271
272 let (min_i, min_j) = labelled
274 .keys()
275 .fold((i32::MAX, i32::MAX), |(a, b), &(i, j)| (a.min(i), b.min(j)));
276 if min_i != 0 || min_j != 0 {
277 let rebased: HashMap<(i32, i32), usize> = labelled
278 .into_iter()
279 .map(|((i, j), idx)| ((i - min_i, j - min_j), idx))
280 .collect();
281 let rebased_by_corner: HashMap<usize, (i32, i32)> =
282 rebased.iter().map(|(&ij, &idx)| (idx, ij)).collect();
283 labelled = rebased;
284 by_corner = rebased_by_corner;
285 }
286 let rebase_pos = |(i, j)| (i - min_i, j - min_j);
287 let ambiguous: HashSet<(i32, i32)> = ambiguous.into_iter().map(rebase_pos).collect();
288 let holes: HashSet<(i32, i32)> = holes.into_iter().map(rebase_pos).collect();
289
290 GrowResult {
291 labelled,
292 by_corner,
293 ambiguous,
294 holes,
295 grid_u,
296 grid_v,
297 }
298}
299
300pub(super) fn enqueue_cardinal_neighbours(
301 pos: (i32, i32),
302 labelled: &HashMap<(i32, i32), usize>,
303 boundary: &mut VecDeque<(i32, i32)>,
304 seen: &mut HashSet<(i32, i32)>,
305) {
306 for (di, dj) in [(1, 0), (-1, 0), (0, 1), (0, -1)] {
307 let neigh = (pos.0 + di, pos.1 + dj);
308 if !labelled.contains_key(&neigh) && seen.insert(neigh) {
309 boundary.push_back(neigh);
310 }
311 }
312}
313
314pub(super) fn collect_labelled_neighbours(
315 pos: (i32, i32),
316 window_half: i32,
317 labelled: &HashMap<(i32, i32), usize>,
318 positions: &[Point2<f32>],
319) -> Vec<LabelledNeighbour> {
320 let mut out = Vec::new();
321 for dj in -window_half..=window_half {
322 for di in -window_half..=window_half {
323 if di == 0 && dj == 0 {
324 continue;
325 }
326 let at = (pos.0 + di, pos.1 + dj);
327 if let Some(&idx) = labelled.get(&at) {
328 out.push(LabelledNeighbour {
329 idx,
330 at,
331 position: positions[idx],
332 });
333 }
334 }
335 }
336 out
337}
338
339pub fn predict_from_neighbours(
375 target: (i32, i32),
376 neighbours: &[LabelledNeighbour],
377 u: Vector2<f32>,
378 v: Vector2<f32>,
379 cell_size: f32,
380 labelled: &HashMap<(i32, i32), usize>,
381 positions: &[Point2<f32>],
382) -> Point2<f32> {
383 debug_assert!(!neighbours.is_empty());
384 let global_i_step = u * cell_size;
385 let global_j_step = v * cell_size;
386
387 let mut sum_x = 0.0_f32;
388 let mut sum_y = 0.0_f32;
389 let mut sum_w = 0.0_f32;
390 for n in neighbours {
391 let di = (target.0 - n.at.0) as f32;
392 let dj = (target.1 - n.at.1) as f32;
393 let d2 = di * di + dj * dj;
394 let w = if d2 > 0.0 { 1.0 / d2 } else { 1.0 };
395
396 let i_step = local_step_at(n.at, (1, 0), labelled, positions).unwrap_or(global_i_step);
397 let j_step = local_step_at(n.at, (0, 1), labelled, positions).unwrap_or(global_j_step);
398
399 let off = i_step * di + j_step * dj;
400 sum_x += w * (n.position.x + off.x);
401 sum_y += w * (n.position.y + off.y);
402 sum_w += w;
403 }
404 Point2::new(sum_x / sum_w, sum_y / sum_w)
405}
406
407pub(super) fn is_extrapolating(target: (i32, i32), neighbours: &[LabelledNeighbour]) -> bool {
417 let mut has_neg_di = false;
418 let mut has_pos_di = false;
419 let mut has_neg_dj = false;
420 let mut has_pos_dj = false;
421 for n in neighbours {
422 let di = target.0 - n.at.0;
423 let dj = target.1 - n.at.1;
424 if di > 0 {
425 has_neg_di = true; } else if di < 0 {
427 has_pos_di = true;
428 }
429 if dj > 0 {
430 has_neg_dj = true;
431 } else if dj < 0 {
432 has_pos_dj = true;
433 }
434 }
435 !(has_neg_di && has_pos_di && has_neg_dj && has_pos_dj)
436}
437
438pub(super) fn local_step_at(
443 at: (i32, i32),
444 step: (i32, i32),
445 labelled: &HashMap<(i32, i32), usize>,
446 positions: &[Point2<f32>],
447) -> Option<Vector2<f32>> {
448 let here = labelled.get(&at).map(|&i| positions[i])?;
449 let fwd = (at.0 + step.0, at.1 + step.1);
450 let bwd = (at.0 - step.0, at.1 - step.1);
451 let fwd_pos = labelled.get(&fwd).map(|&i| positions[i]);
452 let bwd_pos = labelled.get(&bwd).map(|&i| positions[i]);
453 match (fwd_pos, bwd_pos) {
454 (Some(f), Some(b)) => {
455 let v = (f - b) * 0.5;
456 Some(v)
457 }
458 (Some(f), None) => Some(f - here),
459 (None, Some(b)) => Some(here - b),
460 (None, None) => None,
461 }
462}
463
464pub(super) fn collect_candidates<V: GrowValidator>(
465 tree: &KdTree<f32, 2>,
466 slot_to_corner: &[usize],
467 prediction: Point2<f32>,
468 search_r: f32,
469 validator: &V,
470 required_label: Option<u8>,
471 by_corner: &HashMap<usize, (i32, i32)>,
472) -> Vec<(usize, f32)> {
473 let r2 = search_r * search_r;
474 let mut out: Vec<(usize, f32)> = Vec::new();
475 for nn in tree
476 .within_unsorted::<SquaredEuclidean>(&[prediction.x, prediction.y], r2)
477 .into_iter()
478 {
479 let idx = slot_to_corner[nn.item as usize];
480 if by_corner.contains_key(&idx) {
481 continue;
482 }
483 if let Some(req) = required_label {
484 let Some(got) = validator.label_of(idx) else {
485 continue;
486 };
487 if got != req {
488 continue;
489 }
490 }
491 let d = nn.distance.sqrt();
492 out.push((idx, d));
493 }
494 out.sort_by(|a, b| a.1.total_cmp(&b.1));
495 out
496}
497
498pub(super) enum CandidateChoice {
499 None,
500 Ambiguous,
501 Unique(usize),
502}
503
504pub(super) fn choose_unambiguous<V: GrowValidator>(
505 candidates: &[(usize, f32)],
506 ambiguity_factor: f32,
507 prediction: Point2<f32>,
508 positions: &[Point2<f32>],
509 validator: &V,
510 at: (i32, i32),
511 neighbours: &[LabelledNeighbour],
512) -> CandidateChoice {
513 if candidates.is_empty() {
517 return CandidateChoice::None;
518 }
519 if candidates.len() >= 2 {
520 let (_, d0) = candidates[0];
521 let (_, d1) = candidates[1];
522 if d0 <= f32::EPSILON {
523 return CandidateChoice::Ambiguous;
524 }
525 if d1 / d0 < ambiguity_factor {
526 return CandidateChoice::Ambiguous;
527 }
528 }
529 for &(idx, _dist) in candidates {
530 let pos = positions[idx];
531 let _ = pos; match validator.accept_candidate(idx, at, prediction, neighbours) {
533 Admit::Accept => return CandidateChoice::Unique(idx),
534 Admit::Reject => continue,
535 }
536 }
537 CandidateChoice::None
538}
539
540pub(super) fn any_cardinal_edge_ok<V: GrowValidator>(
541 c_idx: usize,
542 pos: (i32, i32),
543 labelled: &HashMap<(i32, i32), usize>,
544 validator: &V,
545) -> bool {
546 let mut found_any = false;
547 for (di, dj) in [(1, 0), (-1, 0), (0, 1), (0, -1)] {
548 let neigh = (pos.0 + di, pos.1 + dj);
549 if let Some(&n_idx) = labelled.get(&neigh) {
550 found_any = true;
551 if validator.edge_ok(c_idx, n_idx, pos, neigh) {
552 return true;
553 }
554 }
555 }
556 !found_any
559}
560
561pub(super) enum BoundaryDecision {
563 Hole,
565 Ambiguous,
567 EdgeRejected,
569 Attach(usize),
571}
572
573#[allow(clippy::too_many_arguments)]
590pub(super) fn process_boundary_cell<V: GrowValidator>(
591 pos: (i32, i32),
592 positions: &[Point2<f32>],
593 labelled: &HashMap<(i32, i32), usize>,
594 by_corner: &HashMap<usize, (i32, i32)>,
595 tree: &KdTree<f32, 2>,
596 tree_slot_to_corner: &[usize],
597 grid_u: Vector2<f32>,
598 grid_v: Vector2<f32>,
599 cell_size: f32,
600 params: &GrowParams,
601 validator: &V,
602) -> (BoundaryDecision, Vec<LabelledNeighbour>) {
603 let neighbours = collect_labelled_neighbours(pos, 1, labelled, positions);
604 if neighbours.is_empty() {
605 return (BoundaryDecision::Hole, neighbours);
606 }
607
608 let prediction = predict_from_neighbours(
609 pos,
610 &neighbours,
611 grid_u,
612 grid_v,
613 cell_size,
614 labelled,
615 positions,
616 );
617
618 let search_r = params.attach_search_rel * cell_size;
619 let extrapolating = is_extrapolating(pos, &neighbours);
620 let local_search_r = if extrapolating {
621 search_r * params.boundary_search_factor
622 } else {
623 search_r
624 };
625
626 let required_label = validator.required_label_at(pos.0, pos.1);
627 let candidates = collect_candidates(
628 tree,
629 tree_slot_to_corner,
630 prediction,
631 local_search_r,
632 validator,
633 required_label,
634 by_corner,
635 );
636
637 let choice = choose_unambiguous(
638 &candidates,
639 params.attach_ambiguity_factor,
640 prediction,
641 positions,
642 validator,
643 pos,
644 &neighbours,
645 );
646
647 let decision = match choice {
648 CandidateChoice::None => BoundaryDecision::Hole,
649 CandidateChoice::Ambiguous => BoundaryDecision::Ambiguous,
650 CandidateChoice::Unique(c_idx) => {
651 if !any_cardinal_edge_ok(c_idx, pos, labelled, validator) {
652 BoundaryDecision::EdgeRejected
653 } else {
654 BoundaryDecision::Attach(c_idx)
655 }
656 }
657 };
658 (decision, neighbours)
659}
660
661#[cfg(test)]
662mod tests {
663 use super::*;
664
665 struct OpenValidator;
668
669 impl GrowValidator for OpenValidator {
670 fn is_eligible(&self, _idx: usize) -> bool {
671 true
672 }
673 fn required_label_at(&self, _i: i32, _j: i32) -> Option<u8> {
674 None
675 }
676 fn label_of(&self, _idx: usize) -> Option<u8> {
677 None
678 }
679 fn accept_candidate(
680 &self,
681 _idx: usize,
682 _at: (i32, i32),
683 _prediction: Point2<f32>,
684 _neighbours: &[LabelledNeighbour],
685 ) -> Admit {
686 Admit::Accept
687 }
688 }
689
690 #[test]
691 fn predict_weights_diagonal_less_than_cardinal() {
692 let s = 10.0_f32;
711 let u = Vector2::new(1.0, 0.0);
712 let v = Vector2::new(0.0, 1.0);
713 let target = (5, 5);
714 let cardinal = LabelledNeighbour {
715 idx: 0,
716 at: (5, 4),
717 position: Point2::new(50.0, 40.0),
718 };
719 let diagonal = LabelledNeighbour {
720 idx: 1,
721 at: (3, 3),
722 position: Point2::new(30.0, 34.0),
723 };
724 let positions = vec![cardinal.position, diagonal.position];
725 let mut labelled = HashMap::new();
726 labelled.insert(cardinal.at, 0usize);
727 labelled.insert(diagonal.at, 1usize);
728 let pred = predict_from_neighbours(
729 target,
730 &[cardinal, diagonal],
731 u,
732 v,
733 s,
734 &labelled,
735 &positions,
736 );
737 let expected_y = (50.0 + 0.125 * 54.0) / 1.125;
738 assert!(
739 (pred.x - 50.0).abs() < 1e-4,
740 "predicted x {} should equal 50",
741 pred.x
742 );
743 assert!(
744 (pred.y - expected_y).abs() < 1e-4,
745 "predicted y {} should equal {} (1/d² weighted)",
746 pred.y,
747 expected_y
748 );
749 let equal_weight_y = (50.0 + 54.0) * 0.5;
750 assert!(
751 (pred.y - 50.0) < (equal_weight_y - 50.0),
752 "weighted bias {} should be smaller than equal-weight bias {}",
753 pred.y - 50.0,
754 equal_weight_y - 50.0,
755 );
756 }
757
758 #[test]
759 fn predict_with_only_cardinal_recovers_exact_offset() {
760 let s = 12.0_f32;
761 let u = Vector2::new(1.0, 0.0);
762 let v = Vector2::new(0.0, 1.0);
763 let target = (2, 2);
764 let neighbour = LabelledNeighbour {
765 idx: 0,
766 at: (1, 2),
767 position: Point2::new(s, 2.0 * s),
768 };
769 let positions = vec![neighbour.position];
770 let mut labelled = HashMap::new();
771 labelled.insert(neighbour.at, 0usize);
772 let pred = predict_from_neighbours(target, &[neighbour], u, v, s, &labelled, &positions);
773 assert!((pred.x - 2.0 * s).abs() < 1e-4);
774 assert!((pred.y - 2.0 * s).abs() < 1e-4);
775 }
776
777 #[test]
778 fn predict_uses_local_step_when_neighbour_has_own_neighbours() {
779 let u = Vector2::new(1.0, 0.0);
797 let v = Vector2::new(0.0, 1.0);
798 let global_cell_size = 50.0_f32;
799 let neighbour = LabelledNeighbour {
800 idx: 0,
801 at: (3, 0),
802 position: Point2::new(300.0, 0.0),
803 };
804 let mut positions = vec![neighbour.position];
805 let mut labelled: HashMap<(i32, i32), usize> = HashMap::new();
806 labelled.insert((3, 0), 0);
807 positions.push(Point2::new(310.0, 0.0));
808 labelled.insert((4, 0), 1);
809 positions.push(Point2::new(320.0, 0.0));
810 labelled.insert((5, 0), 2);
811
812 let pred = predict_from_neighbours(
813 (2, 0),
814 &[neighbour],
815 u,
816 v,
817 global_cell_size,
818 &labelled,
819 &positions,
820 );
821 assert!(
824 (pred.x - 290.0).abs() < 1e-3,
825 "expected adaptive prediction at x=290, got {}",
826 pred.x
827 );
828 assert!((pred.y - 0.0).abs() < 1e-3);
829 }
830
831 #[test]
832 fn predict_falls_back_to_global_when_no_local_steps() {
833 let u = Vector2::new(1.0, 0.0);
838 let v = Vector2::new(0.0, 1.0);
839 let s = 25.0_f32;
840 let neighbour = LabelledNeighbour {
841 idx: 0,
842 at: (4, 4),
843 position: Point2::new(100.0, 100.0),
844 };
845 let positions = vec![neighbour.position];
846 let mut labelled: HashMap<(i32, i32), usize> = HashMap::new();
847 labelled.insert((4, 4), 0);
848 let pred = predict_from_neighbours((5, 4), &[neighbour], u, v, s, &labelled, &positions);
849 assert!((pred.x - (100.0 + s)).abs() < 1e-3);
850 assert!((pred.y - 100.0).abs() < 1e-3);
851 }
852
853 #[test]
854 fn open_validator_grows_clean_grid() {
855 let s = 20.0_f32;
856 let rows = 6_i32;
857 let cols = 6_i32;
858 let mut positions = Vec::new();
859 let mut seed_idx = [0usize; 4];
860 for j in 0..rows {
861 for i in 0..cols {
862 let x = i as f32 * s + 50.0;
863 let y = j as f32 * s + 50.0;
864 let k = positions.len();
865 positions.push(Point2::new(x, y));
866 if (i, j) == (0, 0) {
867 seed_idx[0] = k;
868 }
869 if (i, j) == (1, 0) {
870 seed_idx[1] = k;
871 }
872 if (i, j) == (0, 1) {
873 seed_idx[2] = k;
874 }
875 if (i, j) == (1, 1) {
876 seed_idx[3] = k;
877 }
878 }
879 }
880
881 let seed = Seed {
882 a: seed_idx[0],
883 b: seed_idx[1],
884 c: seed_idx[2],
885 d: seed_idx[3],
886 };
887 let res = bfs_grow(&positions, seed, s, &GrowParams::default(), &OpenValidator);
888 assert_eq!(res.labelled.len(), (rows * cols) as usize);
889 let (mi, mj) = res
891 .labelled
892 .keys()
893 .fold((i32::MAX, i32::MAX), |(a, b), &(i, j)| (a.min(i), b.min(j)));
894 assert_eq!((mi, mj), (0, 0));
895 }
896}