ray-tracing2/ray-tracing-scene/src/triangle_bvh.rs
2024-10-03 22:32:37 +02:00

358 lines
11 KiB
Rust

use ray_tracing_core::{
prelude::*,
scene::{Intersection, Scene},
};
type Index = u32;
#[derive(Debug)]
pub struct TriangleBVH<R: Rng> {
vertices: Vec<Pos3>,
triangles: Vec<Triangle>,
materials: Vec<BVHMaterial<R>>,
bvh: Vec<Node>,
}
#[derive(Debug)]
struct BVHMaterial<R: Rng> {
material: Option<Box<dyn Material<R>>>,
light: Option<Box<dyn Light<R>>>,
}
#[derive(Debug, Clone, Copy)]
enum Node {
Inner {
left: Index,
left_aabb: AABB,
right: Index,
right_aabb: AABB,
},
Leaf {
start: Index,
count: Index,
},
}
#[derive(Debug, Clone, Copy)]
pub struct Triangle {
vertices: [Index; 3],
material: Index,
}
impl Triangle {
fn new(vertices: [Index; 3], material: Index) -> Self {
Triangle { vertices, material }
}
}
fn triangle_intersection(ray: Ray, v: [Pos3; 3]) -> Option<Float> {
let e1 = v[1] - v[0];
let e2 = v[2] - v[0];
let ray_cross_e2 = Dir3::cross(ray.dir(), e2);
let det = e1.dot(ray_cross_e2);
if det > -f32::EPSILON && det < f32::EPSILON {
return None; // This ray is parallel to this triangle.
}
let inv_det = 1.0 / det;
let s = ray.start() - v[0];
let u = inv_det * s.dot(ray_cross_e2);
if !(0.0..=1.0).contains(&u) {
return None;
}
let s_cross_e1 = s.cross(e1);
let v = inv_det * Dir3::dot(ray.dir(), s_cross_e1);
if v < 0.0 || u + v > 1.0 {
return None;
}
// At this stage we can compute t to find out where the intersection point is on the line.
let t = inv_det * e2.dot(s_cross_e1);
if t > Float::EPSILON {
// ray intersection
Some(t)
} else {
// This means that there is a line intersection but not a ray intersection.
None
}
}
fn triangle_normal(v: [Pos3; 3]) -> Dir3 {
let e1 = v[1] - v[0];
let e2 = v[2] - v[0];
Dir3::cross(e1, e2)
}
fn calculate_aabb(vertices: &[Pos3], triangles: &[Triangle]) -> AABB {
let mut aabb = AABB::new(
vertices[triangles[0].vertices[0] as usize],
vertices[triangles[0].vertices[1] as usize],
);
aabb = aabb.extend(vertices[triangles[0].vertices[2] as usize]);
for t in triangles.iter().skip(1) {
for v in t.vertices {
aabb = aabb.extend(vertices[v as usize]);
}
}
aabb
}
fn build_bvh(
vertices: &[Pos3],
triangles: &mut [Triangle],
bvh: &mut Vec<Node>,
node: usize,
aabb: AABB,
) {
let (start, count) = if let Node::Leaf { start, count } = bvh[node] {
(start, count)
} else {
unreachable!()
};
if count < 8 {
return;
}
let size = aabb.size();
let dim = if size.x() > Float::max(size.y(), size.z()) {
0
} else if size.y() > Float::max(size.x(), size.z()) {
1
} else {
2
};
let get_key = |t: &Triangle| {
t.vertices
.iter()
.map(|&v| vertices[v as usize][dim])
.sum::<Float>()
/ 3.0
};
triangles.sort_by(|a, b| get_key(a).partial_cmp(&get_key(b)).unwrap());
let (left_range, right_range) = triangles.split_at_mut((count / 2) as usize);
let left_aabb = calculate_aabb(vertices, left_range);
let right_aabb = calculate_aabb(vertices, right_range);
let i = bvh.len() as Index;
bvh[node] = Node::Inner {
left: i,
left_aabb,
right: i + 1,
right_aabb,
};
bvh.push(Node::Leaf {
start,
count: left_range.len() as Index,
});
bvh.push(Node::Leaf {
start: start + left_range.len() as Index,
count: right_range.len() as Index,
});
build_bvh(vertices, left_range, bvh, i as usize, left_aabb);
build_bvh(vertices, right_range, bvh, i as usize + 1, right_aabb);
}
impl<R: Rng> TriangleBVH<R> {
fn new(
vertices: Vec<Pos3>,
mut triangles: Vec<Triangle>,
materials: Vec<BVHMaterial<R>>,
) -> Self {
let mut bvh = vec![Node::Leaf {
start: 0,
count: triangles.len() as Index,
}];
let aabb = calculate_aabb(&vertices, &triangles);
build_bvh(&vertices, &mut triangles, &mut bvh, 0, aabb);
Self {
vertices,
bvh,
triangles,
materials,
}
}
fn get_vertices(&self, triangle: Index) -> [Pos3; 3] {
let t = self.triangles[triangle as usize];
[
self.vertices[t.vertices[0] as usize],
self.vertices[t.vertices[1] as usize],
self.vertices[t.vertices[2] as usize],
]
}
fn intersect_bvh(
&self,
node: Index,
ray: Ray,
min: Float,
max: Float,
) -> Option<(Index, Float)> {
match self.bvh[node as usize] {
Node::Inner {
left,
left_aabb,
right,
right_aabb,
} => {
let left_intersect = left_aabb.intersect_ray(ray, min, max);
let right_intersect = right_aabb.intersect_ray(ray, min, max);
match (left_intersect, right_intersect) {
(None, None) => None,
(None, Some(_)) => self.intersect_bvh(right, ray, min, max),
(Some(_), None) => self.intersect_bvh(left, ray, min, max),
(Some(l), Some(r)) => {
let close;
let far;
if l < r {
close = left;
far = right;
} else {
close = right;
far = left;
}
if let Some(close_intersect) = self.intersect_bvh(close, ray, min, max) {
if let Some(far_intersect) = self
.intersect_bvh(far, ray, min, Float::min(max, close_intersect.1))
.filter(|far_intersect| far_intersect.1 < close_intersect.1)
{
Some(far_intersect)
} else {
Some(close_intersect)
}
} else {
self.intersect_bvh(far, ray, min, max)
}
}
}
}
Node::Leaf { start, count } => {
let mut intersection = None;
for i in start..(start + count) {
if let Some(t) = triangle_intersection(ray, self.get_vertices(i)) {
if min <= t && t <= max && !intersection.is_some_and(|(_, old_t)| t > old_t)
{
intersection.replace((i, t));
}
}
}
intersection
}
}
}
}
impl<R: Rng> Scene<R> for TriangleBVH<R> {
fn intersect(
&self,
ray: Ray,
min: Float,
max: Float,
) -> Option<ray_tracing_core::scene::Intersection<'_, R>> {
let (i, t) = self.intersect_bvh(0, ray, min, max)?;
let triangle = self.triangles[i as usize];
let material = &self.materials[triangle.material as usize];
let n = triangle_normal(self.get_vertices(i)).normalize();
Some(Intersection::new(
t,
n,
material.material.as_deref(),
material.light.as_deref(),
))
}
}
pub mod examples {
use super::{BVHMaterial, Triangle, TriangleBVH};
use ray_tracing_core::{light::AreaLight, prelude::*};
use ray_tracing_material::{diffuse::DiffuseMaterial, mirror::Mirror};
use std::fmt::Debug;
pub fn cornel<R: Rng + Debug>() -> TriangleBVH<R> {
let side_length = 1.5;
let light_size = 0.5;
let light_offset = 0.01;
let box_size = 0.3;
dbg!(TriangleBVH::new(
vec![
Pos3::new(side_length, side_length, side_length),
Pos3::new(side_length, side_length, -side_length),
Pos3::new(side_length, -side_length, side_length),
Pos3::new(side_length, -side_length, -side_length),
Pos3::new(-side_length, side_length, side_length),
Pos3::new(-side_length, side_length, -side_length),
Pos3::new(-side_length, -side_length, side_length),
Pos3::new(-side_length, -side_length, -side_length),
Pos3::new(light_size, side_length - light_offset, light_size),
Pos3::new(light_size, side_length - light_offset, -light_size),
Pos3::new(-light_size, side_length - light_offset, light_size),
Pos3::new(-light_size, side_length - light_offset, -light_size),
Pos3::new(box_size, -side_length, 0.0),
Pos3::new(0.0, -side_length, -box_size),
Pos3::new(-box_size, -side_length, 0.0),
Pos3::new(0.0, -side_length, box_size),
Pos3::new(0.0, box_size - side_length, 0.0),
],
vec![
Triangle::new([0, 1, 2], 0),
Triangle::new([1, 3, 2], 0),
Triangle::new([0, 4, 1], 0),
Triangle::new([1, 4, 5], 0),
Triangle::new([2, 3, 6], 0),
Triangle::new([3, 7, 6], 0),
Triangle::new([0, 2, 4], 1),
Triangle::new([6, 4, 2], 1),
Triangle::new([1, 5, 3], 2),
Triangle::new([7, 3, 5], 2),
Triangle::new([8, 10, 9], 3),
Triangle::new([11, 9, 10], 3),
Triangle::new([12, 13, 16], 4),
Triangle::new([13, 14, 16], 4),
Triangle::new([14, 15, 16], 4),
Triangle::new([15, 12, 16], 4),
],
vec![
BVHMaterial {
material: Some(Box::new(DiffuseMaterial::new(Color::new(0.8, 0.8, 0.8)))),
light: None,
},
BVHMaterial {
material: Some(Box::new(DiffuseMaterial::new(Color::new(0.9, 0.0, 0.0)))),
light: None,
},
BVHMaterial {
material: Some(Box::new(DiffuseMaterial::new(Color::new(0.0, 0.9, 0.0)))),
light: None,
},
BVHMaterial {
material: None,
light: Some(Box::new(AreaLight::new(Color::white() * 30.0))),
},
BVHMaterial {
material: Some(Box::new(Mirror::new(Color::new(1.0, 1.0, 1.0)))),
light: None,
},
],
))
}
}