use crate::*;
use glam::*;
use rayon::prelude::*;
use std::fmt::{Debug, Formatter};
use std::num::NonZeroUsize;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
pub trait Primitive<AabbType: Debug + Copy + Send + Sync = i32>: Debug + Send + Sync {
fn center(&self) -> Vec3;
fn aabb(&self) -> Aabb<AabbType>;
}
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)]
pub enum BuildType {
None,
LocallyOrderedClustered,
BinnedSAH,
Spatial,
}
#[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)]
pub enum BuildError {
NoPrimitives,
InequalAabbsAndPrimitives(usize, usize),
}
impl std::fmt::Display for BuildError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{}",
match self {
BuildError::NoPrimitives => "No primitives supplied".to_owned(),
BuildError::InequalAabbsAndPrimitives(bbs, prims) => {
format!("#Aabbs({}) != #Primitives({})", *bbs, *prims)
}
}
)
}
}
impl std::error::Error for BuildError {}
#[derive(Debug, Copy, Clone)]
pub struct Builder<'a, T: Primitive<i32>> {
pub aabbs: Option<&'a [Aabb<i32>]>,
pub primitives: &'a [T],
pub primitives_per_leaf: Option<NonZeroUsize>,
}
impl<'a, T: Primitive<i32>> Builder<'a, T> {
pub fn construct_spatial_sah(self) -> Result<Bvh, BuildError>
where
T: SpatialTriangle,
{
if self.primitives.is_empty() {
return Err(BuildError::NoPrimitives);
}
if let Some(aabbs) = self.aabbs {
if aabbs.len() != self.primitives.len() {
return Err(BuildError::InequalAabbsAndPrimitives(
aabbs.len(),
self.primitives.len(),
));
}
Ok(SpatialSahBuilder::new(aabbs, self.primitives, self.primitives_per_leaf).build())
} else {
let aabbs = self
.primitives
.iter()
.par_bridge()
.map(|p| p.aabb())
.collect::<Vec<_>>();
Ok(SpatialSahBuilder::new(&aabbs, self.primitives, self.primitives_per_leaf).build())
}
}
pub fn construct_binned_sah(self) -> Result<Bvh, BuildError> {
if self.primitives.is_empty() {
return Err(BuildError::NoPrimitives);
}
if let Some(aabbs) = self.aabbs {
if aabbs.len() != self.primitives.len() {
return Err(BuildError::InequalAabbsAndPrimitives(
aabbs.len(),
self.primitives.len(),
));
}
Ok(BinnedSahBuilder::new(&aabbs, self.primitives, self.primitives_per_leaf).build())
} else {
let aabbs = self
.primitives
.iter()
.par_bridge()
.map(|p| p.aabb())
.collect::<Vec<_>>();
Ok(BinnedSahBuilder::new(&aabbs, self.primitives, self.primitives_per_leaf).build())
}
}
pub fn construct_locally_ordered_clustered(self) -> Result<Bvh, BuildError> {
if self.primitives.is_empty() {
return Err(BuildError::NoPrimitives);
}
if let Some(aabbs) = self.aabbs {
if aabbs.len() != self.primitives.len() {
return Err(BuildError::InequalAabbsAndPrimitives(
aabbs.len(),
self.primitives.len(),
));
}
Ok(LocallyOrderedClusteringBuilder::new(aabbs, self.primitives).build())
} else {
let aabbs = self
.primitives
.iter()
.par_bridge()
.map(|p| p.aabb())
.collect::<Vec<_>>();
Ok(LocallyOrderedClusteringBuilder::new(&aabbs, self.primitives).build())
}
}
}
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug, Clone)]
pub struct Bvh {
pub(crate) nodes: Vec<BvhNode>,
pub(crate) prim_indices: Vec<u32>,
pub(crate) build_type: BuildType,
}
impl Default for Bvh {
fn default() -> Self {
Self {
nodes: Default::default(),
prim_indices: Default::default(),
build_type: BuildType::None,
}
}
}
impl Bvh {
pub fn nodes(&self) -> &[BvhNode] {
self.nodes.as_slice()
}
pub fn indices(&self) -> &[u32] {
self.prim_indices.as_slice()
}
pub fn build_type(&self) -> BuildType {
self.build_type
}
pub fn prim_count(&self) -> usize {
self.prim_indices.len()
}
pub fn refit(&mut self, new_aabbs: &[Aabb]) {
for i in (0..self.nodes.len()).rev() {
let mut aabb = Aabb::new();
if let Some(left_first) = self.nodes[i].get_left_first() {
let count = self.nodes[i].get_count_unchecked();
if self.nodes[i].is_leaf() {
for i in 0..count {
let prim_id = self.prim_indices[(left_first + i as u32) as usize] as usize;
aabb.grow_bb(&new_aabbs[prim_id]);
}
} else {
aabb.grow_bb(&self.nodes[left_first as usize].bounds);
aabb.grow_bb(&self.nodes[(left_first + 1) as usize].bounds);
}
aabb.offset_by(0.0001);
}
self.nodes[i].bounds = aabb;
}
}
fn traverse_check(&self, cur_node: &BvhNode, checks: &mut [u8]) {
if let Some(left_first) = cur_node.get_left_first() {
if cur_node.is_leaf() {
for i in 0..cur_node.get_count_unchecked() {
let prim_id = self.prim_indices[(left_first as i32 + i) as usize] as usize;
debug_assert!(
prim_id < checks.len(),
"prim_id: {}, max_prim_id: {}",
prim_id,
checks.len()
);
checks[prim_id] = 1;
}
} else {
let left_node = &self.nodes[left_first as usize];
let right_node = &self.nodes[left_first as usize + 1];
self.traverse_check(left_node, checks);
self.traverse_check(right_node, checks);
}
}
}
pub fn validate(&self, prim_count: usize) -> bool {
let mut found_indices = vec![0_u8; prim_count];
self.traverse_check(&self.nodes[0], found_indices.as_mut());
let mut valid = true;
for (i, index) in found_indices.iter().copied().enumerate() {
if index == 0 {
eprintln!("prim ({}) not referenced", i);
valid = false;
}
}
valid
}
pub fn into_raw_indices(self) -> Vec<u32> {
self.prim_indices
}
pub fn into_raw_nodes(self) -> Vec<BvhNode> {
self.nodes
}
pub fn into_raw(self) -> (Vec<BvhNode>, Vec<u32>) {
(self.nodes, self.prim_indices)
}
pub fn traverse_iter<'a, 'b, T: 'a + Primitive>(
&'a self,
ray: &'b mut Ray,
primitives: &'a [T],
) -> BvhIterator<'a, 'b, T> {
BvhIterator::new(ray, self, primitives)
}
pub fn traverse_iter_indices<'a>(&'a self, ray: &'a mut Ray) -> BvhIndexIterator<'a> {
BvhIndexIterator::new(ray, self)
}
pub fn traverse_iter_packet<'a, 'b, T: 'a + Primitive>(
&'a self,
ray: &'b mut RayPacket4,
primitives: &'a [T],
) -> BvhPacketIterator<'a, 'b, T> {
BvhPacketIterator::new(ray, self, primitives)
}
pub fn traverse_iter_indices_packet<'a>(
&'a self,
ray: &'a mut RayPacket4,
) -> BvhPacketIndexIterator<'a> {
BvhPacketIndexIterator::new(ray, self)
}
}
impl<'a, 'b, T: 'a + Primitive> IntoRayIterator<'a, 'b, T> for Bvh {
type RIterator = BvhIterator<'a, 'b, T>;
fn iter(&'a self, ray: &'b mut Ray, primitives: &'a [T]) -> Self::RIterator {
BvhIterator::new(ray, self, primitives)
}
}
impl<'a> IntoRayIndexIterator<'a> for Bvh {
type RIterator = BvhIndexIterator<'a>;
fn iter(&'a self, ray: &'a mut Ray) -> Self::RIterator {
BvhIndexIterator::new(ray, self)
}
}
impl<'a, 'b, T: 'a + Primitive> IntoPacketIterator<'a, 'b, T> for Bvh {
type RIterator = BvhPacketIterator<'a, 'b, T>;
fn iter(&'a self, packet: &'b mut RayPacket4, primitives: &'a [T]) -> Self::RIterator {
BvhPacketIterator::new(packet, self, primitives)
}
}
impl<'a> IntoPacketIndexIterator<'a> for Bvh {
type RIterator = BvhPacketIndexIterator<'a>;
fn iter(&'a self, packet: &'a mut RayPacket4) -> Self::RIterator {
BvhPacketIndexIterator::new(packet, self)
}
}
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug, Clone)]
pub struct Mbvh {
pub(crate) nodes: Vec<BvhNode>,
pub(crate) m_nodes: Vec<MbvhNode>,
pub(crate) prim_indices: Vec<u32>,
}
impl Default for Mbvh {
fn default() -> Self {
Self {
nodes: Default::default(),
m_nodes: Default::default(),
prim_indices: Default::default(),
}
}
}
impl Mbvh {
pub fn nodes(&self) -> &[BvhNode] {
self.nodes.as_slice()
}
pub fn quad_nodes(&self) -> &[MbvhNode] {
self.m_nodes.as_slice()
}
pub fn indices(&self) -> &[u32] {
self.prim_indices.as_slice()
}
pub fn prim_count(&self) -> usize {
self.prim_indices.len()
}
pub fn construct_from_raw(nodes: &[BvhNode], prim_indices: &[u32]) -> Self {
debug_assert!(!nodes.is_empty());
let mut m_nodes = vec![MbvhNode::new(); nodes.len()];
let mut pool_ptr = 1;
if nodes.len() <= 4 {
for (i, node) in nodes.iter().enumerate() {
m_nodes[0].set_bounds_bb(i, &node.bounds);
m_nodes[0].children[i] = node.get_left_first_unchecked();
m_nodes[0].counts[i] = node.get_count_unchecked();
}
return Mbvh {
nodes: nodes.to_vec(),
m_nodes,
prim_indices: prim_indices.to_vec(),
};
}
MbvhNode::merge_nodes(0, 0, nodes, m_nodes.as_mut_slice(), &mut pool_ptr);
Mbvh {
nodes: nodes.to_vec(),
m_nodes,
prim_indices: prim_indices.to_vec(),
}
}
pub fn construct(bvh: &Bvh) -> Self {
if bvh.nodes.is_empty() {
return Mbvh::default();
}
let mut m_nodes = vec![MbvhNode::new(); bvh.nodes.len()];
let mut pool_ptr = 1;
MbvhNode::merge_nodes(
0,
0,
bvh.nodes.as_slice(),
m_nodes.as_mut_slice(),
&mut pool_ptr,
);
m_nodes.resize(pool_ptr, Default::default());
Mbvh {
nodes: bvh.nodes.clone(),
m_nodes,
prim_indices: bvh.prim_indices.clone(),
}
}
pub fn into_raw_indices(self) -> Vec<u32> {
self.prim_indices
}
pub fn into_raw_nodes(self) -> Vec<MbvhNode> {
self.m_nodes
}
pub fn into_raw(self) -> (Vec<MbvhNode>, Vec<u32>) {
(self.m_nodes, self.prim_indices)
}
pub fn traverse_iter<'a, 'b, T: Primitive>(
&'a self,
ray: &'b mut Ray,
primitives: &'a [T],
) -> MbvhIterator<'a, 'b, T> {
MbvhIterator::new(ray, self, primitives)
}
pub fn traverse_iter_indices<'a>(&'a self, ray: &'a mut Ray) -> MbvhIndexIterator<'a> {
MbvhIndexIterator::new(ray, self)
}
pub fn traverse_iter_packet<'a, 'b, T: Primitive>(
&'a self,
ray: &'b mut RayPacket4,
primitives: &'a [T],
) -> MbvhPacketIterator<'a, 'b, T> {
MbvhPacketIterator::new(ray, self, primitives)
}
pub fn traverse_iter_indices_packet<'a>(
&'a self,
ray: &'a mut RayPacket4,
) -> MbvhPacketIndexIterator<'a> {
MbvhPacketIndexIterator::new(ray, self)
}
}
impl From<Bvh> for Mbvh {
fn from(bvh: Bvh) -> Self {
Mbvh::construct(&bvh)
}
}
impl Bounds<i32> for Bvh {
fn bounds(&self) -> Aabb {
self.nodes
.get(0)
.map(|node| node.bounds)
.unwrap_or_default()
}
}
impl Bounds<i32> for Mbvh {
fn bounds(&self) -> Aabb {
self.nodes
.get(0)
.map(|node| node.bounds)
.unwrap_or_default()
}
}
impl<'a, 'b, T: 'a + Primitive> IntoRayIterator<'a, 'b, T> for Mbvh {
type RIterator = MbvhIterator<'a, 'b, T>;
fn iter(&'a self, ray: &'b mut Ray, primitives: &'a [T]) -> Self::RIterator {
MbvhIterator::new(ray, self, primitives)
}
}
impl<'a> IntoRayIndexIterator<'a> for Mbvh {
type RIterator = MbvhIndexIterator<'a>;
fn iter(&'a self, ray: &'a mut Ray) -> Self::RIterator {
MbvhIndexIterator::new(ray, self)
}
}
impl<'a, 'b, T: 'a + Primitive> IntoPacketIterator<'a, 'b, T> for Mbvh {
type RIterator = MbvhPacketIterator<'a, 'b, T>;
fn iter(&'a self, packet: &'b mut RayPacket4, primitives: &'a [T]) -> Self::RIterator {
MbvhPacketIterator::new(packet, self, primitives)
}
}
impl<'a> IntoPacketIndexIterator<'a> for Mbvh {
type RIterator = MbvhPacketIndexIterator<'a>;
fn iter(&'a self, packet: &'a mut RayPacket4) -> Self::RIterator {
MbvhPacketIndexIterator::new(packet, self)
}
}