use crate::diagnostics::ContainmentStats;
use crate::index::{IndexedEnvelope, RStarBackend};
use crate::options::TouchPolicy;
use crate::polygonizer::{
bounding_rect_3d, guaranteed_interior_probe_prepared, rings_share_edge, rings_touch_at_vertex,
};
use crate::types::Polygon3D;
use crate::utils::simd::SimdRing;
#[cfg(feature = "parallel")]
use rayon::prelude::*;
use rstar::AABB;
use std::sync::OnceLock;
struct PreparedRing {
signed_area: f64,
aabb: Option<AABB<[f64; 2]>>,
diagonal: f64,
probe: OnceLock<Option<geo_types::Point<f64>>>,
}
impl PreparedRing {
fn new(polygon: &Polygon3D) -> Self {
let signed_area = Polygon3D::ring_signed_area_2d(&polygon.exterior);
let bbox = bounding_rect_3d(&polygon.exterior);
let diagonal = bbox
.as_ref()
.map(|bbox| {
let dx = bbox.max().x - bbox.min().x;
let dy = bbox.max().y - bbox.min().y;
(dx * dx + dy * dy).sqrt()
})
.unwrap_or(1.0);
let aabb = bbox.map(|bbox| {
AABB::from_corners([bbox.min().x, bbox.min().y], [bbox.max().x, bbox.max().y])
});
Self {
signed_area,
aabb,
diagonal,
probe: OnceLock::new(),
}
}
fn probe(&self, polygon: &Polygon3D, locator: &SimdRing) -> Option<geo_types::Point<f64>> {
*self.probe.get_or_init(|| {
guaranteed_interior_probe_prepared(
&polygon.exterior,
self.signed_area,
self.diagonal,
locator,
)
})
}
}
pub struct ContainmentForest {
pub tree: RStarBackend,
pub simd_shells: Vec<Option<SimdRing>>,
pub shell_areas: Vec<Option<f64>>,
prepared_shells: Vec<PreparedRing>,
}
impl ContainmentForest {
pub fn new(shells: &[Polygon3D]) -> Self {
#[cfg(feature = "parallel")]
let prepared_and_locators: Vec<_> = shells
.par_iter()
.map(|shell| {
let locator = SimdRing::new_3d(&shell.exterior);
(PreparedRing::new(shell), locator)
})
.collect();
#[cfg(not(feature = "parallel"))]
let prepared_and_locators: Vec<_> = shells
.iter()
.map(|shell| {
let locator = SimdRing::new_3d(&shell.exterior);
(PreparedRing::new(shell), locator)
})
.collect();
let mut prepared_shells = Vec::with_capacity(shells.len());
let mut simd_shells = Vec::with_capacity(shells.len());
let mut shell_areas = Vec::with_capacity(shells.len());
for (prepared, locator) in prepared_and_locators {
shell_areas.push(Some(prepared.signed_area.abs()));
prepared_shells.push(prepared);
simd_shells.push(Some(locator));
}
let mut indexed_shells = Vec::with_capacity(shells.len());
for (i, prepared) in prepared_shells.iter().enumerate() {
if let Some(aabb) = prepared.aabb {
indexed_shells.push(IndexedEnvelope { aabb, index: i });
}
}
let tree = RStarBackend::new(indexed_shells);
Self {
tree,
simd_shells,
shell_areas,
prepared_shells,
}
}
pub fn filter_polygonal(&self, shells: &[Polygon3D], touch_policy: &TouchPolicy) -> Vec<bool> {
self.filter_polygonal_impl(shells, touch_policy, None)
}
pub(crate) fn filter_polygonal_with_stats(
&self,
shells: &[Polygon3D],
touch_policy: &TouchPolicy,
) -> (Vec<bool>, ContainmentStats) {
let mut stats = ContainmentStats::default();
let keep_mask = self.filter_polygonal_impl(shells, touch_policy, Some(&mut stats));
(keep_mask, stats)
}
fn filter_polygonal_impl(
&self,
shells: &[Polygon3D],
touch_policy: &TouchPolicy,
mut stats: Option<&mut ContainmentStats>,
) -> Vec<bool> {
let mut keep_mask = vec![true; shells.len()];
let mut container_counts = vec![0; shells.len()];
for (i, shell) in shells.iter().enumerate() {
let prepared = &self.prepared_shells[i];
let aabb = match prepared.aabb.as_ref() {
Some(aabb) => aabb,
None => {
keep_mask[i] = false;
continue;
}
};
let candidates = self.tree.locate_containing_envelope(aabb);
let probe = prepared.probe(shell, self.simd_shells[i].as_ref().unwrap());
let area_i = prepared.signed_area.abs();
if let Some(probe_pt) = probe {
for j in candidates {
if let Some(stats) = stats.as_deref_mut() {
stats.envelope_candidates += 1;
}
if i == j {
continue;
}
let area_j = self.prepared_shells[j].signed_area.abs();
if !(area_j > area_i || ((area_j - area_i).abs() < 1e-9 && j < i)) {
if let Some(stats) = stats.as_deref_mut() {
stats.area_rejections += 1;
}
continue;
}
if let Some(simd_shell) = &self.simd_shells[j] {
if let Some(stats) = stats.as_deref_mut() {
stats.point_in_ring_calls += 1;
}
if simd_shell.contains(probe_pt.0) {
let touch_ok = match touch_policy {
TouchPolicy::AllowPointTouchDisallowEdgeShare => {
if let Some(stats) = stats.as_deref_mut() {
stats.shared_edge_checks += 1;
}
!rings_share_edge(&shells[j].exterior, &shell.exterior, 1e-10)
}
TouchPolicy::TreatAnyTouchAsDisjoint => {
if let Some(stats) = stats.as_deref_mut() {
stats.shared_edge_checks += 1;
}
let shares_edge = rings_share_edge(
&shells[j].exterior,
&shell.exterior,
1e-10,
);
if let Some(stats) = stats.as_deref_mut() {
stats.shared_vertex_checks += usize::from(!shares_edge);
}
!shares_edge
&& !rings_touch_at_vertex(
&shells[j].exterior,
&shell.exterior,
1e-10,
)
}
TouchPolicy::AllowEdgeShare => true,
};
container_counts[i] += usize::from(touch_ok);
}
}
}
} else {
keep_mask[i] = false;
}
}
for (keep, count) in keep_mask.iter_mut().zip(container_counts.iter()) {
if *keep && count % 2 != 0 {
*keep = false;
}
}
keep_mask
}
pub fn assign_hole(
&self,
hole_3d: &Polygon3D,
shells: &[Polygon3D],
touch_policy: &TouchPolicy,
) -> Option<usize> {
self.assign_hole_impl(hole_3d, shells, touch_policy, None)
}
pub(crate) fn assign_hole_with_stats(
&self,
hole_3d: &Polygon3D,
shells: &[Polygon3D],
touch_policy: &TouchPolicy,
) -> (Option<usize>, ContainmentStats) {
let mut stats = ContainmentStats::default();
let best_shell_idx = self.assign_hole_impl(hole_3d, shells, touch_policy, Some(&mut stats));
(best_shell_idx, stats)
}
fn assign_hole_impl(
&self,
hole_3d: &Polygon3D,
shells: &[Polygon3D],
touch_policy: &TouchPolicy,
mut stats: Option<&mut ContainmentStats>,
) -> Option<usize> {
let hole_locator = SimdRing::new_3d(&hole_3d.exterior);
let prepared_hole = PreparedRing::new(hole_3d);
let hole_aabb = prepared_hole.aabb.as_ref()?;
let candidates = self.tree.locate_containing_envelope(hole_aabb);
let mut best_shell_idx = None;
let mut min_area = f64::MAX;
let probe_point = prepared_hole.probe(hole_3d, &hole_locator)?;
let hole_area = prepared_hole.signed_area.abs();
for idx in candidates {
if let Some(stats) = stats.as_deref_mut() {
stats.envelope_candidates += 1;
}
let area = self.prepared_shells[idx].signed_area.abs();
if area <= hole_area + 1e-6 || area >= min_area {
if let Some(stats) = stats.as_deref_mut() {
stats.area_rejections += 1;
}
continue;
}
if let Some(simd_shell) = &self.simd_shells[idx] {
if let Some(stats) = stats.as_deref_mut() {
stats.point_in_ring_calls += 1;
}
if simd_shell.contains(probe_point.0) {
let touch_ok = match touch_policy {
TouchPolicy::AllowPointTouchDisallowEdgeShare => {
if let Some(stats) = stats.as_deref_mut() {
stats.shared_edge_checks += 1;
}
!rings_share_edge(&shells[idx].exterior, &hole_3d.exterior, 1e-10)
}
TouchPolicy::TreatAnyTouchAsDisjoint => {
if let Some(stats) = stats.as_deref_mut() {
stats.shared_edge_checks += 1;
}
let shares_edge =
rings_share_edge(&shells[idx].exterior, &hole_3d.exterior, 1e-10);
if let Some(stats) = stats.as_deref_mut() {
stats.shared_vertex_checks += usize::from(!shares_edge);
}
!shares_edge
&& !rings_touch_at_vertex(
&shells[idx].exterior,
&hole_3d.exterior,
1e-10,
)
}
TouchPolicy::AllowEdgeShare => true,
};
if touch_ok {
min_area = area;
best_shell_idx = Some(idx);
}
}
}
}
best_shell_idx
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::Coord3D;
#[test]
fn prepared_ring_retains_containment_metadata() {
let polygon = Polygon3D::new(
vec![
Coord3D::new(0.0, 0.0, 0.0),
Coord3D::new(10.0, 0.0, 0.0),
Coord3D::new(10.0, 10.0, 0.0),
Coord3D::new(0.0, 10.0, 0.0),
Coord3D::new(0.0, 0.0, 0.0),
],
vec![],
vec![],
vec![],
);
let locator = SimdRing::new_3d(&polygon.exterior);
let prepared = PreparedRing::new(&polygon);
assert_eq!(prepared.signed_area, 100.0);
assert_eq!(prepared.aabb.unwrap().lower(), [0.0, 0.0]);
assert!(locator.contains(prepared.probe(&polygon, &locator).unwrap().0));
}
}