const NONE: usize = usize::MAX;
const MAX_DEPTH: u32 = 24;
const EPS2: f32 = 1e-6;
const MIN_DIST2: f32 = 1.0;
#[derive(Clone)]
struct Cell {
center: [f32; 2],
half: f32,
mass: f32,
sum: [f32; 2],
kids: [usize; 4],
body: usize,
}
pub struct QuadTree {
cells: Vec<Cell>,
next: Vec<usize>,
}
impl QuadTree {
pub fn build(pos: &[[f32; 2]]) -> Self {
let mut t = QuadTree { cells: Vec::new(), next: vec![NONE; pos.len()] };
if pos.is_empty() {
return t;
}
let (mut lo, mut hi) = ([f32::MAX; 2], [f32::MIN; 2]);
for p in pos {
lo[0] = lo[0].min(p[0]); lo[1] = lo[1].min(p[1]);
hi[0] = hi[0].max(p[0]); hi[1] = hi[1].max(p[1]);
}
let half = ((hi[0] - lo[0]).max(hi[1] - lo[1]) * 0.5).max(1.0);
let center = [(lo[0] + hi[0]) * 0.5, (lo[1] + hi[1]) * 0.5];
t.cells.push(Cell { center, half, mass: 0.0, sum: [0.0, 0.0], kids: [NONE; 4], body: NONE });
for i in 0..pos.len() {
t.insert(0, i, pos, 0);
}
t
}
fn quadrant(center: [f32; 2], p: [f32; 2]) -> usize {
(if p[0] >= center[0] { 1 } else { 0 }) | (if p[1] < center[1] { 2 } else { 0 })
}
fn child_cell(parent: &Cell, q: usize) -> Cell {
let h = parent.half * 0.5;
let dx = if q & 1 == 1 { h } else { -h };
let dy = if q & 2 == 2 { -h } else { h };
Cell {
center: [parent.center[0] + dx, parent.center[1] + dy],
half: h,
mass: 0.0,
sum: [0.0, 0.0],
kids: [NONE; 4],
body: NONE,
}
}
fn insert(&mut self, cell: usize, b: usize, pos: &[[f32; 2]], depth: u32) {
self.cells[cell].mass += 1.0;
self.cells[cell].sum[0] += pos[b][0];
self.cells[cell].sum[1] += pos[b][1];
let is_internal = self.cells[cell].kids.iter().any(|&k| k != NONE);
if is_internal {
let q = Self::quadrant(self.cells[cell].center, pos[b]);
let kid = self.ensure_kid(cell, q);
self.insert(kid, b, pos, depth + 1);
return;
}
let existing = self.cells[cell].body;
if existing == NONE {
self.cells[cell].body = b;
return;
}
if depth >= MAX_DEPTH {
self.next[b] = existing;
self.cells[cell].body = b;
return;
}
self.cells[cell].body = NONE;
for who in [existing, b] {
let q = Self::quadrant(self.cells[cell].center, pos[who]);
let kid = self.ensure_kid(cell, q);
self.insert(kid, who, pos, depth + 1);
}
}
fn ensure_kid(&mut self, cell: usize, q: usize) -> usize {
let existing = self.cells[cell].kids[q];
if existing != NONE {
return existing;
}
let c = Self::child_cell(&self.cells[cell], q);
self.cells.push(c);
let idx = self.cells.len() - 1;
self.cells[cell].kids[q] = idx;
idx
}
pub fn repulsion(&self, i: usize, pos: &[[f32; 2]], theta: f32, strength: f32) -> [f32; 2] {
let mut f = [0.0f32, 0.0];
if self.cells.is_empty() || i >= pos.len() {
return f;
}
self.walk(0, i, pos, theta, strength, &mut f);
f
}
fn walk(&self, cell: usize, i: usize, pos: &[[f32; 2]], theta: f32, strength: f32, f: &mut [f32; 2]) {
let c = &self.cells[cell];
if c.mass == 0.0 {
return;
}
let is_internal = c.kids.iter().any(|&k| k != NONE);
if !is_internal {
let mut b = c.body;
while b != NONE {
if b != i {
Self::pair(pos[i], pos[b], i, b, strength, f);
}
b = self.next[b];
}
return;
}
let centroid = [c.sum[0] / c.mass, c.sum[1] / c.mass];
let dx = centroid[0] - pos[i][0];
let dy = centroid[1] - pos[i][1];
let d2 = dx * dx + dy * dy;
if d2 > 0.0 && (c.half * 2.0) * (c.half * 2.0) < theta * theta * d2 {
let d2 = Self::clamp_near(d2);
let mag = strength * c.mass / d2;
f[0] += dx * mag;
f[1] += dy * mag;
return;
}
for &k in &c.kids {
if k != NONE {
self.walk(k, i, pos, theta, strength, f);
}
}
}
fn clamp_near(d2: f32) -> f32 {
if d2 < MIN_DIST2 { (MIN_DIST2 * d2).sqrt() } else { d2 }
}
fn pair(a: [f32; 2], b: [f32; 2], ai: usize, bi: usize, strength: f32, f: &mut [f32; 2]) {
let mut dx = b[0] - a[0];
let mut dy = b[1] - a[1];
let mut d2 = dx * dx + dy * dy;
if d2 < EPS2 {
let s = (ai.wrapping_mul(31).wrapping_add(bi)) as f32;
let (jx, jy) = ((s * 0.7).sin(), (s * 1.3).cos());
let n = (jx * jx + jy * jy).sqrt().max(1e-6);
dx = jx / n;
dy = jy / n;
d2 = MIN_DIST2;
}
let d2 = Self::clamp_near(d2);
let mag = strength / d2;
f[0] += dx * mag;
f[1] += dy * mag;
}
}
#[cfg(test)]
mod tests {
use super::*;
fn brute_force(pos: &[[f32; 2]], i: usize, strength: f32) -> [f32; 2] {
let mut f = [0.0f32, 0.0];
for (j, p) in pos.iter().enumerate() {
if j == i { continue; }
let dx = p[0] - pos[i][0];
let dy = p[1] - pos[i][1];
let d2 = (dx * dx + dy * dy).max(1e-6);
let mag = strength / d2;
f[0] += dx * mag;
f[1] += dy * mag;
}
f
}
fn scatter(n: usize) -> Vec<[f32; 2]> {
let mut s = 12345u64;
(0..n)
.map(|_| {
let mut nxt = || {
s ^= s >> 12; s ^= s << 25; s ^= s >> 27;
((s.wrapping_mul(0x2545F4914F6CDD1D) >> 33) as f32) / (1u64 << 31) as f32
};
[(nxt() - 0.5) * 1000.0, (nxt() - 0.5) * 1000.0]
})
.collect()
}
#[test]
fn theta_zero_is_exact_because_no_cell_can_be_approximated() {
let pos = scatter(60);
let t = QuadTree::build(&pos);
for i in 0..pos.len() {
let got = t.repulsion(i, &pos, 0.0, -100.0);
let want = brute_force(&pos, i, -100.0);
assert!(
(got[0] - want[0]).abs() < 0.05 && (got[1] - want[1]).abs() < 0.05,
"body {i}: tree {got:?} vs brute {want:?}"
);
}
}
#[test]
fn default_theta_approximates_brute_force_within_tolerance() {
let pos = scatter(200);
let t = QuadTree::build(&pos);
let mut worst = 0.0f32;
for i in 0..pos.len() {
let got = t.repulsion(i, &pos, 0.9, -100.0);
let want = brute_force(&pos, i, -100.0);
let mag = (want[0] * want[0] + want[1] * want[1]).sqrt().max(1e-3);
let err = ((got[0] - want[0]).powi(2) + (got[1] - want[1]).powi(2)).sqrt() / mag;
worst = worst.max(err);
}
assert!(worst < 0.30, "worst relative error {worst} exceeds 30% at theta=0.9");
}
#[test]
fn accuracy_degrades_monotonically_with_theta() {
let pos = scatter(200);
let t = QuadTree::build(&pos);
let err_at = |theta: f32| {
let mut worst = 0.0f32;
for i in 0..pos.len() {
let got = t.repulsion(i, &pos, theta, -100.0);
let want = brute_force(&pos, i, -100.0);
let mag = (want[0] * want[0] + want[1] * want[1]).sqrt().max(1e-3);
let e = ((got[0] - want[0]).powi(2) + (got[1] - want[1]).powi(2)).sqrt() / mag;
worst = worst.max(e);
}
worst
};
let (e0, e3, e5, e9) = (err_at(0.0), err_at(0.3), err_at(0.5), err_at(0.9));
assert!(e0 < 1e-4, "theta=0 must be exact, got {e0}");
assert!(e3 < e5 && e5 < e9, "error must grow with theta: {e3} {e5} {e9}");
}
#[test]
fn every_cell_holds_each_body_beneath_it_exactly_once() {
let pos = scatter(200);
let t = QuadTree::build(&pos);
let n = pos.len() as f32;
assert_eq!(t.cells[0].mass, n, "root mass must be the body count");
let mean = [
pos.iter().map(|p| p[0]).sum::<f32>() / n,
pos.iter().map(|p| p[1]).sum::<f32>() / n,
];
assert!((t.cells[0].sum[0] / n - mean[0]).abs() < 0.01);
assert!((t.cells[0].sum[1] / n - mean[1]).abs() < 0.01);
let mut reachable = 0usize;
for c in &t.cells {
if c.kids.iter().any(|&k| k != NONE) {
let km: f32 = c.kids.iter().filter(|&&k| k != NONE).map(|&k| t.cells[k].mass).sum();
assert!((km - c.mass).abs() < 1e-3, "cell mass {} vs kids {km}", c.mass);
} else {
let mut b = c.body;
let mut chain = 0.0f32;
while b != NONE {
reachable += 1;
chain += 1.0;
b = t.next[b];
}
assert!((chain - c.mass).abs() < 1e-3, "leaf mass {} vs chain {chain}", c.mass);
}
}
assert_eq!(reachable, pos.len(), "every body must sit in exactly one leaf");
}
#[test]
fn a_body_never_repels_itself() {
let pos = vec![[10.0, 10.0]];
let t = QuadTree::build(&pos);
assert_eq!(t.repulsion(0, &pos, 0.9, -100.0), [0.0, 0.0]);
}
#[test]
fn coincident_bodies_do_not_produce_nan() {
let pos = vec![[5.0, 5.0], [5.0, 5.0], [5.0, 5.0]];
let t = QuadTree::build(&pos);
for i in 0..3 {
let f = t.repulsion(i, &pos, 0.9, -100.0);
assert!(f[0].is_finite() && f[1].is_finite(), "body {i} produced {f:?}");
}
}
#[test]
fn coincident_bodies_are_nudged_not_launched() {
let pos = vec![[5.0, 5.0], [5.0, 5.0], [5.0, 5.0]];
let t = QuadTree::build(&pos);
for i in 0..3 {
let f = t.repulsion(i, &pos, 0.9, -100.0);
let mag = (f[0] * f[0] + f[1] * f[1]).sqrt();
assert!(mag <= 200.0, "body {i} force {mag} exceeds two bodies' worth of strength");
}
let t2 = QuadTree::build(&pos);
assert_eq!(t.repulsion(0, &pos, 0.9, -100.0), t2.repulsion(0, &pos, 0.9, -100.0));
}
#[test]
fn a_near_coincident_pair_is_softened_to_about_the_strength() {
let pos = vec![[0.0, 0.0], [0.01, 0.0]];
let t = QuadTree::build(&pos);
let f = t.repulsion(0, &pos, 0.9, -800.0);
let mag = (f[0] * f[0] + f[1] * f[1]).sqrt();
assert!(
(mag - 800.0).abs() < 1.0,
"expected the softened |strength| = 800, got {mag}; unsoftened this is 80000"
);
assert!(f[0] < 0.0, "and it must still point away from the other body, got {f:?}");
}
#[test]
fn coincident_bodies_are_pushed_in_different_directions() {
let pos = vec![[5.0, 5.0], [5.0, 5.0]];
let t = QuadTree::build(&pos);
let f0 = t.repulsion(0, &pos, 0.9, -100.0);
let f1 = t.repulsion(1, &pos, 0.9, -100.0);
let spread = ((f0[0] - f1[0]).powi(2) + (f0[1] - f1[1]).powi(2)).sqrt();
assert!(
spread > 1.0,
"bodies got near-identical pushes {f0:?} and {f1:?}; a rigid cluster never separates"
);
}
#[test]
fn deeply_clustered_bodies_terminate() {
let pos: Vec<[f32; 2]> = (0..50).map(|i| [1.0 + i as f32 * 1e-7, 1.0]).collect();
let t = QuadTree::build(&pos);
let f = t.repulsion(0, &pos, 0.9, -100.0);
assert!(f[0].is_finite() && f[1].is_finite());
}
#[test]
fn an_empty_set_builds_and_answers_nothing() {
let pos: Vec<[f32; 2]> = vec![];
let t = QuadTree::build(&pos);
assert_eq!(t.repulsion(0, &pos, 0.9, -100.0), [0.0, 0.0]);
}
#[test]
fn repulsion_pushes_apart_and_attraction_pulls_together() {
let pos = vec![[0.0, 0.0], [10.0, 0.0]];
let t = QuadTree::build(&pos);
let f = t.repulsion(0, &pos, 0.9, -100.0);
assert!(f[0] < 0.0, "body 0 should be pushed away from body 1 (leftward), got {f:?}");
let a = t.repulsion(0, &pos, 0.9, 100.0);
assert!(a[0] > 0.0, "positive strength should attract, got {a:?}");
}
#[test]
fn force_magnitude_falls_off_as_inverse_distance() {
let near = vec![[0.0, 0.0], [10.0, 0.0]];
let far = vec![[0.0, 0.0], [20.0, 0.0]];
let fn_ = QuadTree::build(&near).repulsion(0, &near, 0.0, -100.0)[0].abs();
let ff = QuadTree::build(&far).repulsion(0, &far, 0.0, -100.0)[0].abs();
assert!((fn_ - 10.0).abs() < 1e-3, "expected |strength|/d = 10, got {fn_}");
assert!((ff - 5.0).abs() < 1e-3, "expected |strength|/d = 5, got {ff}");
}
}