1use crate::math::{cross, dot, sub};
2
3#[repr(C)]
4#[derive(Clone, Copy, Debug, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
5pub struct Node {
6 pub lo: [f32; 4],
7 pub hi: [f32; 4],
8 pub first: u32,
9 pub count: u32,
10 pub _pad: [u32; 2],
11}
12
13#[derive(Clone, Copy, Debug, PartialEq)]
14pub struct Triangle {
15 pub vertices: [[f32; 3]; 3],
16 pub material: u32,
17}
18
19#[derive(Clone, Copy, Debug, PartialEq)]
20pub struct Ray {
21 pub origin: [f32; 3],
22 pub direction: [f32; 3],
23}
24
25#[derive(Clone, Copy, Debug, PartialEq)]
26pub struct Hit {
27 pub distance: f32,
28 pub normal: [f32; 3],
29 pub material: u32,
30 pub index: u32,
31}
32
33#[derive(Clone, Debug)]
34pub struct Bvh {
35 pub nodes: Vec<Node>,
36 pub order: Vec<u32>,
37}
38
39#[derive(Clone, Copy)]
40struct Bounds {
41 lo: [f32; 3],
42 hi: [f32; 3],
43}
44
45impl Bounds {
46 fn empty() -> Self {
47 Self {
48 lo: [f32::INFINITY; 3],
49 hi: [f32::NEG_INFINITY; 3],
50 }
51 }
52 fn point(&mut self, p: [f32; 3]) {
53 for (axis, value) in p.into_iter().enumerate() {
54 self.lo[axis] = self.lo[axis].min(value);
55 self.hi[axis] = self.hi[axis].max(value);
56 }
57 }
58 fn grow(&mut self, other: Self) {
59 self.point(other.lo);
60 self.point(other.hi);
61 }
62 fn area(self) -> f32 {
63 let d = sub(self.hi, self.lo);
64 if d[0] < 0.0 {
65 0.0
66 } else {
67 2.0 * (d[0] * d[1] + d[1] * d[2] + d[2] * d[0])
68 }
69 }
70}
71
72const BINS: usize = 16;
73const LEAF: usize = 4;
74
75impl Bvh {
76 pub fn build(triangles: &[Triangle]) -> Self {
77 let boxes: Vec<_> = triangles
78 .iter()
79 .map(|triangle| {
80 let mut bounds = Bounds::empty();
81 for &vertex in &triangle.vertices {
82 bounds.point(vertex);
83 }
84 bounds
85 })
86 .collect();
87 let centers: Vec<_> = boxes
88 .iter()
89 .map(|b| [0, 1, 2].map(|axis| (b.lo[axis] + b.hi[axis]) * 0.5))
90 .collect();
91 let mut order: Vec<_> = (0..triangles.len() as u32).collect();
92 let mut nodes = Vec::new();
93 if !triangles.is_empty() {
94 split(&boxes, ¢ers, &mut order, 0, triangles.len(), &mut nodes);
95 }
96 Self { nodes, order }
97 }
98 pub fn intersect(&self, triangles: &[Triangle], ray: Ray, limit: f32) -> Option<Hit> {
99 if self.nodes.is_empty() {
100 return None;
101 }
102 let mut best = None;
103 let mut stack = vec![0usize];
104 while let Some(index) = stack.pop() {
105 let node = self.nodes[index];
106 let max = best.as_ref().map_or(limit, |hit: &Hit| hit.distance);
107 if !box_hit(node, ray, max) {
108 continue;
109 }
110 if node.count == 0 {
111 stack.push(node.first as usize);
112 stack.push(index + 1);
113 } else {
114 for &i in &self.order[node.first as usize..(node.first + node.count) as usize] {
115 if let Some(mut hit) = triangle_hit(
116 &triangles[i as usize],
117 ray,
118 best.as_ref().map_or(limit, |hit: &Hit| hit.distance),
119 ) {
120 hit.index = i;
121 best = Some(hit);
122 }
123 }
124 }
125 }
126 best
127 }
128}
129
130fn split(
131 boxes: &[Bounds],
132 centers: &[[f32; 3]],
133 order: &mut [u32],
134 start: usize,
135 end: usize,
136 nodes: &mut Vec<Node>,
137) {
138 let mut bounds = Bounds::empty();
139 let mut centroids = Bounds::empty();
140 for &index in &order[start..end] {
141 bounds.grow(boxes[index as usize]);
142 centroids.point(centers[index as usize]);
143 }
144 let here = nodes.len();
145 nodes.push(Node {
146 lo: [bounds.lo[0], bounds.lo[1], bounds.lo[2], 0.0],
147 hi: [bounds.hi[0], bounds.hi[1], bounds.hi[2], 0.0],
148 first: start as u32,
149 count: (end - start) as u32,
150 _pad: [0; 2],
151 });
152 let len = end - start;
153 if len <= LEAF {
154 return;
155 }
156 let mut best = (f32::INFINITY, 0usize, 0usize);
157 for (axis, (&lo, &hi)) in centroids.lo.iter().zip(¢roids.hi).enumerate() {
158 let extent = hi - lo;
159 if extent <= 1e-9 {
160 continue;
161 }
162 let mut bins = [(Bounds::empty(), 0usize); BINS];
163 for &index in &order[start..end] {
164 let bin = bin(centers[index as usize][axis], centroids.lo[axis], extent);
165 bins[bin].0.grow(boxes[index as usize]);
166 bins[bin].1 += 1;
167 }
168 let mut areas = [0.0; BINS];
169 let mut counts = [0; BINS];
170 let mut bound = Bounds::empty();
171 let mut count = 0;
172 for index in 0..BINS {
173 if bins[index].1 > 0 {
174 bound.grow(bins[index].0);
175 }
176 count += bins[index].1;
177 areas[index] = bound.area();
178 counts[index] = count;
179 }
180 bound = Bounds::empty();
181 count = 0;
182 for index in (1..BINS).rev() {
183 if bins[index].1 > 0 {
184 bound.grow(bins[index].0);
185 }
186 count += bins[index].1;
187 let cost = areas[index - 1] * counts[index - 1] as f32 + bound.area() * count as f32;
188 if counts[index - 1] > 0 && count > 0 && cost < best.0 {
189 best = (cost, axis, index);
190 }
191 }
192 }
193 if best.0 >= bounds.area() * len as f32 && len <= 16 {
194 return;
195 }
196 let mid = if best.0.is_finite() {
197 let axis = best.1;
198 let extent = centroids.hi[axis] - centroids.lo[axis];
199 let slice = &mut order[start..end];
200 let mut left = 0;
201 let mut right = len;
202 while left < right {
203 if bin(
204 centers[slice[left] as usize][axis],
205 centroids.lo[axis],
206 extent,
207 ) < best.2
208 {
209 left += 1;
210 } else {
211 right -= 1;
212 slice.swap(left, right);
213 }
214 }
215 if left == 0 || left == len {
216 start + len / 2
217 } else {
218 start + left
219 }
220 } else {
221 start + len / 2
222 };
223 split(boxes, centers, order, start, mid, nodes);
224 let right = nodes.len();
225 split(boxes, centers, order, mid, end, nodes);
226 nodes[here].first = right as u32;
227 nodes[here].count = 0;
228}
229
230fn bin(value: f32, lo: f32, extent: f32) -> usize {
231 (((value - lo) / extent) * BINS as f32).clamp(0.0, (BINS - 1) as f32) as usize
232}
233
234fn box_hit(node: Node, ray: Ray, limit: f32) -> bool {
235 let mut near = 0.0f32;
236 let mut far = limit;
237 for axis in 0..3 {
238 let direction = ray.direction[axis];
239 if direction.abs() < 1e-12 {
240 if ray.origin[axis] < node.lo[axis] || ray.origin[axis] > node.hi[axis] {
241 return false;
242 }
243 } else {
244 let a = (node.lo[axis] - ray.origin[axis]) / direction;
245 let b = (node.hi[axis] - ray.origin[axis]) / direction;
246 near = near.max(a.min(b));
247 far = far.min(a.max(b));
248 if near > far {
249 return false;
250 }
251 }
252 }
253 true
254}
255
256pub fn triangle_hit(triangle: &Triangle, ray: Ray, limit: f32) -> Option<Hit> {
257 let [a, b, c] = triangle.vertices;
258 let e1 = sub(b, a);
259 let e2 = sub(c, a);
260 let p = cross(ray.direction, e2);
261 let det = dot(e1, p);
262 if det.abs() < 1e-9 {
263 return None;
264 }
265 let inv = 1.0 / det;
266 let s = sub(ray.origin, a);
267 let u = dot(s, p) * inv;
268 if !(0.0..=1.0).contains(&u) {
269 return None;
270 }
271 let q = cross(s, e1);
272 let v = dot(ray.direction, q) * inv;
273 if v < 0.0 || u + v > 1.0 {
274 return None;
275 }
276 let distance = dot(e2, q) * inv;
277 if distance <= 1e-5 || distance >= limit {
278 return None;
279 }
280 Some(Hit {
281 distance,
282 normal: crate::math::normalize(cross(e1, e2)),
283 material: triangle.material,
284 index: 0,
285 })
286}
287
288#[cfg(test)]
289mod tests {
290 use super::*;
291 #[test]
292 fn random_rays_match_brute_force_and_build_is_stable() {
293 let mut state = 17u32;
294 let mut random = || {
295 state = state.wrapping_mul(1664525).wrapping_add(1013904223);
296 (state >> 8) as f32 / (1u32 << 24) as f32
297 };
298 let triangles: Vec<_> = (0..120)
299 .map(|_| {
300 let x = random() * 8.0 - 4.0;
301 let y = random() * 8.0 - 4.0;
302 let z = random() * 8.0 - 4.0;
303 Triangle {
304 vertices: [
305 [x, y, z],
306 [x + random() + 0.1, y, z],
307 [x, y + random() + 0.1, z],
308 ],
309 material: 0,
310 }
311 })
312 .collect();
313 let bvh = Bvh::build(&triangles);
314 let other = Bvh::build(&triangles);
315 assert_eq!(bvh.nodes, other.nodes);
316 assert_eq!(bvh.order, other.order);
317 for _ in 0..2000 {
318 let ray = Ray {
319 origin: [
320 random() * 12.0 - 6.0,
321 random() * 12.0 - 6.0,
322 random() * 12.0 - 6.0,
323 ],
324 direction: crate::math::normalize([
325 random() * 2.0 - 1.0,
326 random() * 2.0 - 1.0,
327 random() * 2.0 - 1.0,
328 ]),
329 };
330 let brute = triangles
331 .iter()
332 .enumerate()
333 .filter_map(|(i, t)| {
334 triangle_hit(t, ray, f32::INFINITY).map(|mut h| {
335 h.index = i as u32;
336 h
337 })
338 })
339 .min_by(|a, b| a.distance.total_cmp(&b.distance));
340 let fast = bvh.intersect(&triangles, ray, f32::INFINITY);
341 assert_eq!(fast.map(|h| h.index), brute.map(|h| h.index));
342 }
343 }
344}