1use crate::bvh::{Hit, Ray};
2use crate::math::{add, dot, length, normalize, scale, sub};
3
4#[derive(Clone, Copy, Debug, PartialEq)]
5pub enum Shape {
6 RoundedBox {
7 center: [f32; 3],
8 half: [f32; 3],
9 radius: f32,
10 material: u32,
11 },
12 RoundCone {
13 a: [f32; 3],
14 b: [f32; 3],
15 radius_a: f32,
16 radius_b: f32,
17 material: u32,
18 },
19 Ellipsoid {
20 center: [f32; 3],
21 radii: [f32; 3],
22 material: u32,
23 },
24 Capsule {
25 a: [f32; 3],
26 b: [f32; 3],
27 radius: f32,
28 material: u32,
29 },
30 SmoothUnion {
31 left: usize,
32 right: usize,
33 radius: f32,
34 material: u32,
35 },
36}
37
38impl Shape {
39 pub fn material(self) -> u32 {
40 match self {
41 Self::RoundedBox { material, .. }
42 | Self::RoundCone { material, .. }
43 | Self::Ellipsoid { material, .. }
44 | Self::Capsule { material, .. }
45 | Self::SmoothUnion { material, .. } => material,
46 }
47 }
48 pub fn distance(self, p: [f32; 3], shapes: &[Shape]) -> f32 {
49 match self {
50 Self::RoundedBox {
51 center,
52 half,
53 radius,
54 ..
55 } => {
56 let q = [0, 1, 2].map(|i| (p[i] - center[i]).abs() - (half[i] - radius));
57 length(q.map(|v| v.max(0.0))) + q[0].max(q[1]).max(q[2]).min(0.0) - radius
58 }
59 Self::RoundCone {
60 a,
61 b,
62 radius_a,
63 radius_b,
64 ..
65 } => round_cone(p, a, b, radius_a, radius_b),
66 Self::Ellipsoid { center, radii, .. } => {
67 let q = sub(p, center);
68 let k0 = length([0, 1, 2].map(|i| q[i] / radii[i]));
69 let k1 = length([0, 1, 2].map(|i| q[i] / (radii[i] * radii[i])));
70 k0 * (k0 - 1.0) / k1.max(1e-8)
71 }
72 Self::Capsule { a, b, radius, .. } => {
73 let ba = sub(b, a);
74 let h = (dot(sub(p, a), ba) / dot(ba, ba).max(1e-9)).clamp(0.0, 1.0);
75 length(sub(p, add(a, scale(ba, h)))) - radius
76 }
77 Self::SmoothUnion {
78 left,
79 right,
80 radius,
81 ..
82 } => {
83 if left >= shapes.len() || right >= shapes.len() {
84 return f32::INFINITY;
85 }
86 let a = shapes[left].distance(p, shapes);
87 let b = shapes[right].distance(p, shapes);
88 let h = (0.5 + 0.5 * (b - a) / radius.max(1e-6)).clamp(0.0, 1.0);
89 b * (1.0 - h) + a * h - radius * h * (1.0 - h)
90 }
91 }
92 }
93}
94
95fn round_cone(p: [f32; 3], a: [f32; 3], b: [f32; 3], r1: f32, r2: f32) -> f32 {
96 let ba = sub(b, a);
97 let axis_length = length(ba);
98 if axis_length <= (r1 - r2).abs() {
99 return if r1 >= r2 {
100 length(sub(p, a)) - r1
101 } else {
102 length(sub(p, b)) - r2
103 };
104 }
105 let l2 = dot(ba, ba);
106 let rr = r1 - r2;
107 let a2 = l2 - rr * rr;
108 let pa = sub(p, a);
109 let y = dot(pa, ba);
110 let z = y - l2;
111 let xv = sub(scale(pa, l2), scale(ba, y));
112 let x2 = dot(xv, xv);
113 let y2 = y * y * l2;
114 let z2 = z * z * l2;
115 let k = rr.signum() * rr * rr * x2;
116 if z.signum() * a2 * z2 > k {
117 return (x2 + z2).sqrt() / l2 - r2;
118 }
119 if y.signum() * a2 * y2 < k {
120 return (x2 + y2).sqrt() / l2 - r1;
121 }
122 ((x2 * a2 / l2).max(0.0)).sqrt() / l2 + y * rr / l2 - r1
123}
124
125pub fn intersect(shapes: &[Shape], ray: Ray, limit: f32) -> Option<Hit> {
126 let mut best: Option<Hit> = None;
127 let mut consumed = vec![false; shapes.len()];
128 for shape in shapes {
129 if let Shape::SmoothUnion { left, right, .. } = shape {
130 if *left < shapes.len() {
131 consumed[*left] = true;
132 }
133 if *right < shapes.len() {
134 consumed[*right] = true;
135 }
136 }
137 }
138 for (index, shape) in shapes.iter().copied().enumerate() {
139 if consumed[index] {
140 continue;
141 }
142 let max = best.map_or(limit, |hit| hit.distance);
143 if let Some(distance) = intersect_one(shape, shapes, ray, max) {
144 let p = add(ray.origin, scale(ray.direction, distance));
145 let e = 0.0002;
146 let normal = normalize([0, 1, 2].map(|axis| {
147 let mut a = p;
148 let mut b = p;
149 a[axis] += e;
150 b[axis] -= e;
151 shape.distance(a, shapes) - shape.distance(b, shapes)
152 }));
153 best = Some(Hit {
154 distance,
155 normal,
156 material: shape.material(),
157 index: index as u32,
158 });
159 }
160 }
161 best
162}
163
164fn intersect_one(shape: Shape, shapes: &[Shape], ray: Ray, limit: f32) -> Option<f32> {
165 if let Shape::Ellipsoid { center, radii, .. } = shape {
166 let o = [0, 1, 2].map(|i| (ray.origin[i] - center[i]) / radii[i]);
167 let d = [0, 1, 2].map(|i| ray.direction[i] / radii[i]);
168 let a = dot(d, d);
169 let b = dot(o, d);
170 let c = dot(o, o) - 1.0;
171 let discr = b * b - a * c;
172 if discr < 0.0 {
173 return None;
174 }
175 let root = discr.sqrt();
176 let t = (-b - root) / a;
177 let t = if t > 1e-5 { t } else { (-b + root) / a };
178 return (t > 1e-5 && t < limit).then_some(t);
179 }
180 if !matches!(shape, Shape::SmoothUnion { .. }) {
181 return analytic_hit(shape, ray, limit);
182 }
183 let mut t = 0.0f32;
184 let mut last = shape.distance(ray.origin, shapes);
185 if !last.is_finite() {
186 return None;
187 }
188 for _ in 0..512 {
189 let previous = t;
190 t += last.abs().max(0.00005);
191 if t >= limit {
192 return None;
193 }
194 let distance = shape.distance(add(ray.origin, scale(ray.direction, t)), shapes);
195 if !distance.is_finite() {
196 return None;
197 }
198 if distance.abs() < 1e-5 {
199 return Some(t);
200 }
201 if (last > 0.0 && distance < 0.0) || (last < 0.0 && distance > 0.0) {
202 let mut lo = previous;
203 let mut hi = t;
204 for _ in 0..24 {
205 let mid = (lo + hi) * 0.5;
206 let value = shape.distance(add(ray.origin, scale(ray.direction, mid)), shapes);
207 if value.signum() == last.signum() {
208 lo = mid;
209 } else {
210 hi = mid;
211 }
212 }
213 return Some((lo + hi) * 0.5);
214 }
215 last = distance;
216 }
217 None
218}
219
220fn roots(a: f32, b: f32, c: f32) -> [f32; 2] {
221 if a.abs() < 1e-10 {
222 return if b.abs() < 1e-10 {
223 [f32::INFINITY; 2]
224 } else {
225 [-c / (2.0 * b), f32::INFINITY]
226 };
227 }
228 let discr = b * b - a * c;
229 if discr < 0.0 {
230 return [f32::INFINITY; 2];
231 }
232 let root = discr.sqrt();
233 [(-b - root) / a, (-b + root) / a]
234}
235
236fn analytic_hit(shape: Shape, ray: Ray, limit: f32) -> Option<f32> {
237 let mut best = limit;
238 let mut accept = |t: f32| {
239 if t > 1e-5 && t < best {
240 let p = add(ray.origin, scale(ray.direction, t));
241 if shape.distance(p, &[]).abs() < 0.002 {
242 best = t;
243 }
244 }
245 };
246 match shape {
247 Shape::RoundedBox {
248 center,
249 half,
250 radius,
251 ..
252 } => {
253 let core = half.map(|v| (v - radius).max(0.0));
254 for axis in 0..3 {
255 if ray.direction[axis].abs() > 1e-9 {
256 for side in [-1.0, 1.0] {
257 let t = (center[axis] + side * (core[axis] + radius) - ray.origin[axis])
258 / ray.direction[axis];
259 accept(t);
260 }
261 }
262 }
263 for axis in 0..3 {
264 let j = (axis + 1) % 3;
265 let k = (axis + 2) % 3;
266 for sj in [-1.0, 1.0] {
267 for sk in [-1.0, 1.0] {
268 let oj = ray.origin[j] - center[j] - sj * core[j];
269 let ok = ray.origin[k] - center[k] - sk * core[k];
270 let a = ray.direction[j] * ray.direction[j]
271 + ray.direction[k] * ray.direction[k];
272 let b = oj * ray.direction[j] + ok * ray.direction[k];
273 let c = oj * oj + ok * ok - radius * radius;
274 for t in roots(a, b, c) {
275 if (ray.origin[axis] + ray.direction[axis] * t - center[axis]).abs()
276 <= core[axis] + 0.002
277 {
278 accept(t);
279 }
280 }
281 }
282 }
283 }
284 for sx in [-1.0, 1.0] {
285 for sy in [-1.0, 1.0] {
286 for sz in [-1.0, 1.0] {
287 let corner = add(center, [sx * core[0], sy * core[1], sz * core[2]]);
288 let o = sub(ray.origin, corner);
289 for t in roots(
290 dot(ray.direction, ray.direction),
291 dot(o, ray.direction),
292 dot(o, o) - radius * radius,
293 ) {
294 accept(t);
295 }
296 }
297 }
298 }
299 }
300 Shape::Capsule { a, b, radius, .. } => {
301 let ba = sub(b, a);
302 let oa = sub(ray.origin, a);
303 let baba = dot(ba, ba);
304 let bard = dot(ba, ray.direction);
305 let baoa = dot(ba, oa);
306 let rdoa = dot(ray.direction, oa);
307 let oaoa = dot(oa, oa);
308 let qa = baba * dot(ray.direction, ray.direction) - bard * bard;
309 let qb = baba * rdoa - baoa * bard;
310 let qc = baba * oaoa - baoa * baoa - radius * radius * baba;
311 for t in roots(qa, qb, qc) {
312 let y = baoa + t * bard;
313 if y >= 0.0 && y <= baba {
314 accept(t);
315 }
316 }
317 for center in [a, b] {
318 let o = sub(ray.origin, center);
319 for t in roots(
320 dot(ray.direction, ray.direction),
321 dot(o, ray.direction),
322 dot(o, o) - radius * radius,
323 ) {
324 accept(t);
325 }
326 }
327 }
328 Shape::RoundCone {
329 a,
330 b,
331 radius_a,
332 radius_b,
333 ..
334 } => {
335 let ba = sub(b, a);
336 let length = length(ba);
337 if length > 1e-8 {
338 let axis = scale(ba, 1.0 / length);
339 let slope = (radius_a - radius_b) / length;
340 if slope.abs() < 1.0 {
341 let side = (1.0 - slope * slope).sqrt();
342 let oa = sub(ray.origin, a);
343 let y0 = dot(oa, axis);
344 let yd = dot(ray.direction, axis);
345 let radial0 = sub(oa, scale(axis, y0));
346 let radiald = sub(ray.direction, scale(axis, yd));
347 let e = radius_a - slope * y0;
348 let qa = side * side * dot(radiald, radiald) - slope * slope * yd * yd;
349 let qb = side * side * dot(radial0, radiald) + e * slope * yd;
350 let qc = side * side * dot(radial0, radial0) - e * e;
351 for t in roots(qa, qb, qc) {
352 let y = y0 + t * yd;
353 if y >= slope * radius_a - 0.002 && y <= length + slope * radius_b + 0.002 {
354 accept(t);
355 }
356 }
357 }
358 }
359 for (center, radius) in [(a, radius_a), (b, radius_b)] {
360 let o = sub(ray.origin, center);
361 for t in roots(
362 dot(ray.direction, ray.direction),
363 dot(o, ray.direction),
364 dot(o, o) - radius * radius,
365 ) {
366 accept(t);
367 }
368 }
369 }
370 _ => {}
371 }
372 (best < limit).then_some(best)
373}
374
375#[cfg(test)]
376mod tests {
377 use super::*;
378 fn hit(shape: Shape, origin: [f32; 3], direction: [f32; 3]) -> f32 {
379 intersect(&[shape], Ray { origin, direction }, 100.0)
380 .unwrap()
381 .distance
382 }
383 #[test]
384 fn ellipsoid_matches_quadratic() {
385 let shape = Shape::Ellipsoid {
386 center: [0.0; 3],
387 radii: [2.0, 1.0, 1.0],
388 material: 0,
389 };
390 assert!((hit(shape, [4.0, 0.0, 0.0], [-1.0, 0.0, 0.0]) - 2.0).abs() < 1e-5);
391 assert!((hit(shape, [0.0, 0.0, 0.0], [1.0, 0.0, 0.0]) - 2.0).abs() < 1e-5);
392 }
393 #[test]
394 fn sdf_shapes_match_axis_intersections() {
395 let shapes = [
396 (
397 Shape::RoundedBox {
398 center: [0.0; 3],
399 half: [1.0; 3],
400 radius: 0.2,
401 material: 0,
402 },
403 1.0,
404 ),
405 (
406 Shape::RoundCone {
407 a: [0.0, -1.0, 0.0],
408 b: [0.0, 1.0, 0.0],
409 radius_a: 0.5,
410 radius_b: 0.5,
411 material: 0,
412 },
413 0.5,
414 ),
415 (
416 Shape::Capsule {
417 a: [0.0, -1.0, 0.0],
418 b: [0.0, 1.0, 0.0],
419 radius: 0.5,
420 material: 0,
421 },
422 0.5,
423 ),
424 ];
425 for (shape, radius) in shapes {
426 assert!((hit(shape, [3.0, 0.0, 0.0], [-1.0, 0.0, 0.0]) - (3.0 - radius)).abs() < 0.001);
427 }
428 let shapes = [
429 Shape::Ellipsoid {
430 center: [-1.0, 0.0, 0.0],
431 radii: [1.0; 3],
432 material: 0,
433 },
434 Shape::Ellipsoid {
435 center: [1.0, 0.0, 0.0],
436 radii: [1.0; 3],
437 material: 0,
438 },
439 Shape::SmoothUnion {
440 left: 0,
441 right: 1,
442 radius: 0.5,
443 material: 0,
444 },
445 ];
446 assert!(shapes[2].distance([0.0; 3], &shapes) < 0.0);
447 let exit = intersect(
448 &shapes,
449 Ray {
450 origin: [0.0; 3],
451 direction: [1.0, 0.0, 0.0],
452 },
453 10.0,
454 )
455 .unwrap();
456 assert!((exit.distance - 2.0).abs() < 0.001);
457 }
458 #[test]
459 fn curved_patches_have_analytic_roots() {
460 let box_shape = Shape::RoundedBox {
461 center: [0.0; 3],
462 half: [1.0; 3],
463 radius: 0.2,
464 material: 0,
465 };
466 let direction = normalize([-1.0, -1.0, -1.0]);
467 let expected = 1.2 * 3.0f32.sqrt() - 0.2;
468 assert!((hit(box_shape, [2.0; 3], direction) - expected).abs() < 0.0001);
469 let capsule = Shape::Capsule {
470 a: [0.0, -1.0, 0.0],
471 b: [0.0, 1.0, 0.0],
472 radius: 0.5,
473 material: 0,
474 };
475 assert!((hit(capsule, [0.0, 3.0, 0.0], [0.0, -1.0, 0.0]) - 1.5).abs() < 0.0001);
476 let cone = Shape::RoundCone {
477 a: [0.0, -1.0, 0.0],
478 b: [0.0, 1.0, 0.0],
479 radius_a: 0.5,
480 radius_b: 0.2,
481 material: 0,
482 };
483 assert!((hit(cone, [3.0, -1.0, 0.0], [-1.0, 0.0, 0.0]) - 2.5).abs() < 0.001);
484 assert!(
485 (hit(cone, [3.0, 1.0, 0.0], [-1.0, 0.0, 0.0])
486 - (3.0 - 0.2 / (1.0f32 - 0.15 * 0.15).sqrt()))
487 .abs()
488 < 0.001
489 );
490 }
491}