1#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
2pub struct Aabb {
3 pub min: [i32; 2],
4 pub max: [i32; 2],
5}
6
7impl Aabb {
8 pub const fn new(min: [i32; 2], max: [i32; 2]) -> Self {
9 Self { min, max }
10 }
11
12 pub const fn point(at: [i32; 2]) -> Self {
13 Self {
14 min: at,
15 max: [at[0].saturating_add(1), at[1].saturating_add(1)],
16 }
17 }
18
19 pub const fn around(centre: [i32; 2], half: [i32; 2]) -> Self {
20 Self {
21 min: [
22 centre[0].saturating_sub(half[0]),
23 centre[1].saturating_sub(half[1]),
24 ],
25 max: [
26 centre[0].saturating_add(half[0]),
27 centre[1].saturating_add(half[1]),
28 ],
29 }
30 }
31
32 pub const fn is_empty(&self) -> bool {
33 self.max[0] <= self.min[0] || self.max[1] <= self.min[1]
34 }
35
36 pub const fn overlaps(&self, other: &Aabb) -> bool {
37 !self.is_empty()
38 && !other.is_empty()
39 && self.min[0] < other.max[0]
40 && other.min[0] < self.max[0]
41 && self.min[1] < other.max[1]
42 && other.min[1] < self.max[1]
43 }
44}
45
46#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
47pub struct Aabb3 {
48 pub min: [i32; 3],
49 pub max: [i32; 3],
50}
51
52impl Aabb3 {
53 pub const fn new(min: [i32; 3], max: [i32; 3]) -> Self {
54 Self { min, max }
55 }
56
57 pub const fn point(at: [i32; 3]) -> Self {
58 Self {
59 min: at,
60 max: [
61 at[0].saturating_add(1),
62 at[1].saturating_add(1),
63 at[2].saturating_add(1),
64 ],
65 }
66 }
67
68 pub const fn around(centre: [i32; 3], half: [i32; 3]) -> Self {
69 Self {
70 min: [
71 centre[0].saturating_sub(half[0]),
72 centre[1].saturating_sub(half[1]),
73 centre[2].saturating_sub(half[2]),
74 ],
75 max: [
76 centre[0].saturating_add(half[0]),
77 centre[1].saturating_add(half[1]),
78 centre[2].saturating_add(half[2]),
79 ],
80 }
81 }
82
83 pub const fn is_empty(&self) -> bool {
84 self.max[0] <= self.min[0] || self.max[1] <= self.min[1] || self.max[2] <= self.min[2]
85 }
86
87 pub const fn overlaps(&self, other: &Aabb3) -> bool {
88 !self.is_empty()
89 && !other.is_empty()
90 && self.min[0] < other.max[0]
91 && other.min[0] < self.max[0]
92 && self.min[1] < other.max[1]
93 && other.min[1] < self.max[1]
94 && self.min[2] < other.max[2]
95 && other.min[2] < self.max[2]
96 }
97}
98
99#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
100pub struct Hit {
101 pub a: u32,
102 pub b: u32,
103}
104
105#[derive(Clone, Debug, Default)]
106pub struct Overlaps(Engine<2>);
107
108impl Overlaps {
109 pub fn new() -> Self {
110 Self::default()
111 }
112
113 pub fn find(&mut self, a: &[Aabb], b: &[Aabb]) -> &[Hit] {
114 self.0.load(
115 a.iter().map(|x| (x.min, x.max)),
116 b.iter().map(|x| (x.min, x.max)),
117 );
118 self.0.run()
119 }
120
121 pub fn hits(&self) -> &[Hit] {
122 &self.0.hits
123 }
124}
125
126#[derive(Clone, Debug, Default)]
127pub struct Overlaps3(Engine<3>);
128
129impl Overlaps3 {
130 pub fn new() -> Self {
131 Self::default()
132 }
133
134 pub fn find(&mut self, a: &[Aabb3], b: &[Aabb3]) -> &[Hit] {
135 self.0.load(
136 a.iter().map(|x| (x.min, x.max)),
137 b.iter().map(|x| (x.min, x.max)),
138 );
139 self.0.run()
140 }
141
142 pub fn hits(&self) -> &[Hit] {
143 &self.0.hits
144 }
145}
146
147#[derive(Clone, Copy, Debug)]
148struct Span<const N: usize> {
149 min: [i32; N],
150 max: [i32; N],
151 index: u32,
152}
153
154impl<const N: usize> Span<N> {
155 fn overlaps(&self, other: &Span<N>) -> bool {
156 (0..N).all(|k| self.min[k] < other.max[k] && other.min[k] < self.max[k])
157 }
158}
159
160#[derive(Clone, Copy, Debug)]
161struct Grid<const N: usize> {
162 origin: [i64; N],
163 cell: [i64; N],
164 dims: [usize; N],
165}
166
167impl<const N: usize> Grid<N> {
168 fn at(&self, axis: usize, value: i32) -> usize {
169 let at = (i64::from(value) - self.origin[axis]) / self.cell[axis];
170 (at.max(0) as usize).min(self.dims[axis] - 1)
171 }
172
173 fn range(&self, span: &Span<N>) -> ([usize; N], [usize; N]) {
174 (
175 std::array::from_fn(|k| self.at(k, span.min[k])),
176 std::array::from_fn(|k| self.at(k, span.max[k] - 1)),
177 )
178 }
179
180 fn index(&self, at: [usize; N]) -> usize {
181 (0..N).rev().fold(0, |sum, k| sum * self.dims[k] + at[k])
182 }
183}
184
185fn cells_in<const N: usize>(low: [usize; N], high: [usize; N]) -> usize {
186 (0..N).map(|k| high[k] - low[k] + 1).product()
187}
188
189fn each_cell<const N: usize>(low: [usize; N], high: [usize; N], mut visit: impl FnMut([usize; N])) {
190 let mut at = low;
191 loop {
192 visit(at);
193 let mut k = 0;
194 loop {
195 if k == N {
196 return;
197 }
198 if at[k] < high[k] {
199 at[k] += 1;
200 break;
201 }
202 at[k] = low[k];
203 k += 1;
204 }
205 }
206}
207
208fn fill<const N: usize>(
209 spans: &mut Vec<Span<N>>,
210 boxes: impl Iterator<Item = ([i32; N], [i32; N])>,
211) {
212 spans.clear();
213 for (index, (min, max)) in boxes.enumerate() {
214 if (0..N).any(|k| max[k] <= min[k]) {
215 continue;
216 }
217 assert!(index < 1 << 31, "at most 2^31 boxes a side");
218 spans.push(Span {
219 min,
220 max,
221 index: index as u32,
222 });
223 }
224}
225
226#[derive(Clone, Debug)]
227struct Engine<const N: usize> {
228 boxes: [Vec<Span<N>>; 2],
229 keys: Vec<u64>,
230 active: [Vec<u32>; 2],
231 pairs: Vec<u64>,
232 starts: Vec<u32>,
233 cursor: Vec<u32>,
234 members: Vec<u32>,
235 found: Vec<u32>,
236 hits: Vec<Hit>,
237}
238
239impl<const N: usize> Default for Engine<N> {
240 fn default() -> Self {
241 Self {
242 boxes: Default::default(),
243 keys: Vec::new(),
244 active: Default::default(),
245 pairs: Vec::new(),
246 starts: Vec::new(),
247 cursor: Vec::new(),
248 members: Vec::new(),
249 found: Vec::new(),
250 hits: Vec::new(),
251 }
252 }
253}
254
255impl<const N: usize> Engine<N> {
256 fn load(
257 &mut self,
258 a: impl Iterator<Item = ([i32; N], [i32; N])>,
259 b: impl Iterator<Item = ([i32; N], [i32; N])>,
260 ) {
261 fill(&mut self.boxes[0], a);
262 fill(&mut self.boxes[1], b);
263 }
264
265 fn run(&mut self) -> &[Hit] {
266 self.hits.clear();
267 if self.boxes[0].is_empty() || self.boxes[1].is_empty() {
268 return &self.hits;
269 }
270 match self.plan() {
271 Some(grid) => self.through(grid),
272 None => self.sweep(),
273 }
274 &self.hits
275 }
276
277 fn plan(&self) -> Option<Grid<N>> {
278 let mut low = [i64::MAX; N];
279 let mut high = [i64::MIN; N];
280 let mut extent = [0i128; N];
281 for span in self.boxes.iter().flatten() {
282 for k in 0..N {
283 low[k] = low[k].min(i64::from(span.min[k]));
284 high[k] = high[k].max(i64::from(span.max[k]));
285 extent[k] += i128::from(span.max[k]) - i128::from(span.min[k]);
286 }
287 }
288 let count = self.boxes[0].len() + self.boxes[1].len();
289 let budget = (4 * count).max(64);
290 let mut cell: [i64; N] = std::array::from_fn(|k| (extent[k] / count as i128).max(1) as i64);
291 let dims = |cell: &[i64; N]| -> [usize; N] {
292 std::array::from_fn(|k| ((high[k] - low[k] + cell[k] - 1) / cell[k]).max(1) as usize)
293 };
294 let mut grid_dims = dims(&cell);
295 while grid_dims.iter().map(|&d| d as u128).product::<u128>() > budget as u128 {
296 let widest = (0..N).max_by_key(|&k| (grid_dims[k], std::cmp::Reverse(k)))?;
297 cell[widest] = cell[widest].saturating_mul(2);
298 grid_dims = dims(&cell);
299 }
300 let grid = Grid {
301 origin: low,
302 cell,
303 dims: grid_dims,
304 };
305 let covered: usize = self
306 .boxes
307 .iter()
308 .flatten()
309 .map(|span| {
310 let (from, to) = grid.range(span);
311 cells_in(from, to)
312 })
313 .sum();
314 (covered <= 8 * count + 1024).then_some(grid)
315 }
316
317 fn through(&mut self, grid: Grid<N>) {
318 let cells: usize = grid.dims.iter().product();
319 self.starts.clear();
320 self.starts.resize(cells + 1, 0);
321 for span in &self.boxes[1] {
322 let (from, to) = grid.range(span);
323 each_cell(from, to, |at| self.starts[grid.index(at) + 1] += 1);
324 }
325 for c in 0..cells {
326 self.starts[c + 1] += self.starts[c];
327 }
328 self.cursor.clear();
329 self.cursor.extend_from_slice(&self.starts[..cells]);
330 self.members.clear();
331 self.members.resize(self.starts[cells] as usize, 0);
332 for (position, span) in self.boxes[1].iter().enumerate() {
333 let (from, to) = grid.range(span);
334 each_cell(from, to, |at| {
335 let c = grid.index(at);
336 self.members[self.cursor[c] as usize] = position as u32;
337 self.cursor[c] += 1;
338 });
339 }
340 let targets = &self.boxes[1];
341 for shot in &self.boxes[0] {
342 self.found.clear();
343 let (from, to) = grid.range(shot);
344 each_cell(from, to, |at| {
345 let c = grid.index(at);
346 let edge: [i64; N] = std::array::from_fn(|k| {
347 if at[k] == from[k] {
348 i64::MIN
349 } else {
350 grid.origin[k] + at[k] as i64 * grid.cell[k]
351 }
352 });
353 for &position in &self.members[self.starts[c] as usize..self.starts[c + 1] as usize]
354 {
355 let target = &targets[position as usize];
356 if shot.overlaps(target) && (0..N).all(|k| i64::from(target.min[k]) >= edge[k])
357 {
358 self.found.push(target.index);
359 }
360 }
361 });
362 self.found.sort_unstable();
363 self.hits
364 .extend(self.found.iter().map(|&b| Hit { a: shot.index, b }));
365 }
366 }
367
368 fn sweep(&mut self) {
369 let axis = self.sweep_axis();
370 let Self {
371 boxes,
372 keys,
373 active,
374 pairs,
375 hits,
376 ..
377 } = self;
378 keys.clear();
379 for (side, spans) in boxes.iter().enumerate() {
380 for (position, span) in spans.iter().enumerate() {
381 let start = u64::from(span.min[axis] as u32 ^ 0x8000_0000);
382 keys.push((start << 32) | ((side as u64) << 31) | position as u64);
383 }
384 }
385 keys.sort_unstable();
386 active[0].clear();
387 active[1].clear();
388 pairs.clear();
389 for &key in keys.iter() {
390 let side = ((key >> 31) & 1) as usize;
391 let position = (key & 0x7fff_ffff) as u32;
392 let entry = &boxes[side][position as usize];
393 let spans = &boxes[1 - side];
394 let others = &mut active[1 - side];
395 let mut at = 0;
396 while at < others.len() {
397 let seen = &spans[others[at] as usize];
398 if seen.max[axis] <= entry.min[axis] {
399 others.swap_remove(at);
400 continue;
401 }
402 if entry.overlaps(seen) {
403 let (a, b) = if side == 0 {
404 (entry.index, seen.index)
405 } else {
406 (seen.index, entry.index)
407 };
408 pairs.push((u64::from(a) << 32) | u64::from(b));
409 }
410 at += 1;
411 }
412 active[side].push(position);
413 }
414 pairs.sort_unstable();
415 hits.extend(pairs.iter().map(|key| Hit {
416 a: (key >> 32) as u32,
417 b: *key as u32,
418 }));
419 }
420
421 fn sweep_axis(&self) -> usize {
422 let mut low = [i32::MAX; N];
423 let mut high = [i32::MIN; N];
424 for span in self.boxes.iter().flatten() {
425 for k in 0..N {
426 low[k] = low[k].min(span.min[k]);
427 high[k] = high[k].max(span.min[k]);
428 }
429 }
430 (0..N)
431 .rev()
432 .max_by_key(|&k| i64::from(high[k]) - i64::from(low[k]))
433 .unwrap_or(0)
434 }
435}