use std::{
marker::PhantomData,
os::raw::{c_uint, c_void},
panic::{catch_unwind, AssertUnwindSafe},
ptr::NonNull,
};
use crate::{sys::*, Bounds, BuildPrimitive, BuildQuality, Device, DeviceProperty, Error};
pub struct Bvh {
device: Device,
handle: RTCBVH,
}
impl Drop for Bvh {
fn drop(&mut self) { unsafe { rtcReleaseBVH(self.handle) } }
}
impl Bvh {
pub(crate) fn new(device: &Device) -> Result<Self, Error> {
let handle = unsafe { rtcNewBVH(device.handle) };
if handle.is_null() {
Err(device.get_error())
} else {
Ok(Self {
device: device.clone(),
handle,
})
}
}
}
pub unsafe trait BvhNode: Copy + Send + Sync {}
#[repr(transparent)]
pub struct NodePtr<'id, N> {
ptr: NonNull<N>,
_brand: PhantomData<fn(&'id ()) -> &'id ()>,
}
impl<'id, N> Clone for NodePtr<'id, N> {
fn clone(&self) -> Self { *self }
}
impl<'id, N> Copy for NodePtr<'id, N> {}
unsafe impl<'id, N> Send for NodePtr<'id, N> {}
unsafe impl<'id, N> Sync for NodePtr<'id, N> {}
pub struct Allocator<'id> {
raw: RTCThreadLocalAllocator,
_brand: PhantomData<&'id ()>,
_not_send_sync: PhantomData<*mut ()>,
}
impl<'id> Allocator<'id> {
unsafe fn from_raw(raw: RTCThreadLocalAllocator) -> Self {
Self {
raw,
_brand: PhantomData,
_not_send_sync: PhantomData,
}
}
pub fn alloc<T: Copy>(&self, value: T) -> &'id mut T {
if std::mem::size_of::<T>() == 0 {
std::process::abort();
}
let layout = std::alloc::Layout::new::<T>();
let p = unsafe { rtcThreadLocalAlloc(self.raw, layout.size(), layout.align()) } as *mut T;
if p.is_null() {
std::process::abort();
}
unsafe {
p.write(value);
&mut *p
}
}
}
pub struct Children<'id, N> {
ptr: *const *mut c_void,
len: usize,
_m: PhantomData<NodePtr<'id, N>>,
}
impl<'id, N> Children<'id, N> {
unsafe fn from_raw(ptr: *mut *mut c_void, len: usize) -> Self {
Self {
ptr: ptr as *const *mut c_void,
len,
_m: PhantomData,
}
}
pub fn len(&self) -> usize { self.len }
pub fn is_empty(&self) -> bool { self.len == 0 }
pub fn get(&self, i: usize) -> Option<NodePtr<'id, N>> {
if i >= self.len {
return None;
}
let raw = unsafe { *self.ptr.add(i) };
NonNull::new(raw as *mut N).map(|ptr| NodePtr {
ptr,
_brand: PhantomData,
})
}
}
pub struct ChildBounds<'a> {
ptr: *const *const RTCBounds,
len: usize,
_m: PhantomData<&'a Bounds>,
}
impl<'a> ChildBounds<'a> {
unsafe fn from_raw(ptr: *mut *const RTCBounds, len: usize) -> Self {
Self {
ptr: ptr as *const *const RTCBounds,
len,
_m: PhantomData,
}
}
#[inline]
pub fn len(&self) -> usize { self.len }
#[inline]
pub fn is_empty(&self) -> bool { self.len == 0 }
#[inline]
pub fn get(&self, i: usize) -> Option<&Bounds> {
if i >= self.len {
return None;
}
unsafe { (*self.ptr.add(i)).as_ref() }
}
}
pub trait BvhBuilder: Send + Sync {
type Node<'id>: BvhNode + 'id;
const MAX_CHILDREN: usize;
const SPATIAL_SPLITS: bool = false;
const PROGRESS: bool = false;
fn create_node<'id>(
&self,
alloc: &Allocator<'id>,
child_count: usize,
) -> &'id mut Self::Node<'id>;
fn set_children<'id>(
&self,
node: &mut Self::Node<'id>,
children: Children<'id, Self::Node<'id>>,
);
fn set_bounds<'id>(&self, node: &mut Self::Node<'id>, bounds: ChildBounds<'_>);
fn create_leaf<'id>(
&self,
alloc: &Allocator<'id>,
prims: &[BuildPrimitive],
) -> &'id mut Self::Node<'id>;
fn split(&self, prim: &BuildPrimitive, dim: u32, pos: f32) -> (Bounds, Bounds) {
let lo_full = bounds_of(prim);
let mut lo = lo_full;
let mut hi = lo_full;
set_axis_upper(&mut lo, dim, pos);
set_axis_lower(&mut hi, dim, pos);
(lo, hi)
}
fn progress(&self, _fraction: f64) {}
}
#[derive(Clone, Debug)]
pub struct BuildConfig {
pub quality: BuildQuality,
pub dynamic: bool,
pub max_branching_factor: u32,
pub max_depth: u32,
pub sah_block_size: u32,
pub min_leaf_size: u32,
pub max_leaf_size: u32,
pub traversal_cost: f32,
pub intersection_cost: f32,
}
impl Default for BuildConfig {
fn default() -> Self {
Self {
quality: BuildQuality::MEDIUM,
dynamic: false,
max_branching_factor: 2,
max_depth: 32,
sah_block_size: 1,
min_leaf_size: 1,
max_leaf_size: 32,
traversal_cost: 1.0,
intersection_cost: 1.0,
}
}
}
impl BuildConfig {
fn validate(&self, prim_count: usize, max_children: usize) -> Result<(), Error> {
let bad = Err(Error::INVALID_ARGUMENT);
match self.quality {
BuildQuality::LOW | BuildQuality::MEDIUM | BuildQuality::HIGH => {}
_ => return bad,
}
let branch_cap = if self.quality == BuildQuality::LOW {
8
} else {
16
};
if self.max_branching_factor < 2
|| self.max_branching_factor > branch_cap
|| self.max_branching_factor as usize > max_children
{
return bad;
}
if self.max_leaf_size < 1 || self.max_leaf_size > 32 {
return bad;
}
if self.min_leaf_size < 1 || self.min_leaf_size > self.max_leaf_size {
return bad;
}
if self.sah_block_size < 1 || self.max_depth < 1 {
return bad;
}
if !self.traversal_cost.is_finite()
|| self.traversal_cost <= 0.0
|| !self.intersection_cost.is_finite()
|| self.intersection_cost <= 0.0
{
return bad;
}
if self.quality == BuildQuality::LOW && prim_count > u32::MAX as usize {
return bad;
}
Ok(())
}
}
pub struct BvhResult<'id, B: BvhBuilder> {
root: Option<NodePtr<'id, B::Node<'id>>>,
_brand: PhantomData<fn(&'id ()) -> &'id ()>,
}
impl<'id, B: BvhBuilder> BvhResult<'id, B> {
pub fn root(&self) -> Option<&B::Node<'id>> { self.root.map(|p| unsafe { p.ptr.as_ref() }) }
pub fn root_ptr(&self) -> Option<NodePtr<'id, B::Node<'id>>> { self.root }
pub fn resolve(&self, p: NodePtr<'id, B::Node<'id>>) -> &B::Node<'id> {
unsafe { p.ptr.as_ref() }
}
}
struct BuildState<B> {
builder: *const B,
}
fn catch_abort<R>(f: impl FnOnce() -> R) -> R {
match catch_unwind(AssertUnwindSafe(f)) {
Ok(r) => r,
Err(_) => std::process::abort(),
}
}
#[inline]
unsafe fn builder_of<'x, B: BvhBuilder>(user: *mut c_void) -> &'x B {
&*((*(user as *const BuildState<B>)).builder)
}
unsafe extern "C" fn create_node_tramp<B: BvhBuilder>(
alloc: RTCThreadLocalAllocator,
child_count: c_uint,
user: *mut c_void,
) -> *mut c_void {
catch_abort(|| unsafe {
let builder = builder_of::<B>(user);
let allocator = Allocator::from_raw(alloc);
let node = builder.create_node(&allocator, child_count as usize);
node as *mut _ as *mut c_void
})
}
unsafe extern "C" fn create_leaf_tramp<B: BvhBuilder>(
alloc: RTCThreadLocalAllocator,
prims: *const RTCBuildPrimitive,
prim_count: usize,
user: *mut c_void,
) -> *mut c_void {
catch_abort(|| unsafe {
let builder = builder_of::<B>(user);
let allocator = Allocator::from_raw(alloc);
let slice = std::slice::from_raw_parts(prims, prim_count);
let node = builder.create_leaf(&allocator, slice);
node as *mut _ as *mut c_void
})
}
unsafe extern "C" fn set_children_tramp<B: BvhBuilder>(
node: *mut c_void,
children: *mut *mut c_void,
child_count: c_uint,
user: *mut c_void,
) {
catch_abort(|| unsafe {
let builder = builder_of::<B>(user);
let node = &mut *(node as *mut B::Node<'_>);
let kids = Children::from_raw(children, child_count as usize);
builder.set_children(node, kids);
})
}
unsafe extern "C" fn set_bounds_tramp<B: BvhBuilder>(
node: *mut c_void,
bounds: *mut *const RTCBounds,
child_count: c_uint,
user: *mut c_void,
) {
catch_abort(|| unsafe {
let builder = builder_of::<B>(user);
let node = &mut *(node as *mut B::Node<'_>);
let view = ChildBounds::from_raw(bounds, child_count as usize);
builder.set_bounds(node, view);
})
}
unsafe extern "C" fn split_tramp<B: BvhBuilder>(
prim: *const RTCBuildPrimitive,
dim: c_uint,
pos: f32,
lbounds: *mut RTCBounds,
rbounds: *mut RTCBounds,
user: *mut c_void,
) {
catch_abort(|| unsafe {
let builder = builder_of::<B>(user);
let (lo, hi) = builder.split(&*prim, dim, pos);
*lbounds = lo;
*rbounds = hi;
});
}
unsafe extern "C" fn progress_tramp<B: BvhBuilder>(user: *mut c_void, n: f64) -> bool {
catch_abort(|| unsafe {
let builder = builder_of::<B>(user);
builder.progress(n);
});
true
}
impl Bvh {
pub fn build_scoped<B, F, R>(
&mut self,
config: &BuildConfig,
primitives: &mut Vec<BuildPrimitive>,
builder: &B,
f: F,
) -> Result<R, Error>
where
B: BvhBuilder,
F: for<'id> FnOnce(BvhResult<'id, B>) -> R,
{
self.build_with_brand(config, primitives, builder, f)
}
fn build_with_brand<'id, B, F, R>(
&mut self,
config: &BuildConfig,
primitives: &mut Vec<BuildPrimitive>,
builder: &B,
f: F,
) -> Result<R, Error>
where
B: BvhBuilder,
F: FnOnce(BvhResult<'id, B>) -> R,
{
config.validate(primitives.len(), B::MAX_CHILDREN)?;
if std::mem::size_of::<B::Node<'id>>() == 0 {
return Err(Error::INVALID_ARGUMENT);
}
if self.device.get_property(DeviceProperty::VERSION)? != RTC_VERSION as isize {
return Err(Error::INVALID_OPERATION);
}
if primitives.is_empty() {
return Ok(f(BvhResult {
root: None,
_brand: PhantomData,
}));
}
if B::SPATIAL_SPLITS && config.quality == BuildQuality::HIGH {
let need = primitives
.len()
.checked_mul(2)
.ok_or(Error::INVALID_ARGUMENT)?;
primitives.reserve(need - primitives.len());
}
let capacity = primitives.capacity();
let count = primitives.len();
let state = BuildState::<B> {
builder: builder as *const B,
};
let mut args: RTCBuildArguments = unsafe { std::mem::zeroed() };
args.byteSize = std::mem::size_of::<RTCBuildArguments>();
args.buildQuality = config.quality;
args.buildFlags = if config.dynamic {
RTCBuildFlags::DYNAMIC
} else {
RTCBuildFlags::NONE
};
args.maxBranchingFactor = config.max_branching_factor;
args.maxDepth = config.max_depth;
args.sahBlockSize = config.sah_block_size;
args.minLeafSize = config.min_leaf_size;
args.maxLeafSize = config.max_leaf_size;
args.traversalCost = config.traversal_cost;
args.intersectionCost = config.intersection_cost;
args.bvh = self.handle;
args.primitives = primitives.as_mut_ptr();
args.primitiveCount = count;
args.primitiveArrayCapacity = capacity;
args.createNode = Some(create_node_tramp::<B>);
args.setNodeChildren = Some(set_children_tramp::<B>);
args.setNodeBounds = Some(set_bounds_tramp::<B>);
args.createLeaf = Some(create_leaf_tramp::<B>);
args.splitPrimitive = if B::SPATIAL_SPLITS {
Some(split_tramp::<B>)
} else {
None
};
args.buildProgress = if B::PROGRESS {
Some(progress_tramp::<B>)
} else {
None
};
args.userPtr = &state as *const BuildState<B> as *mut c_void;
let _ = self.device.get_error();
let root = unsafe { rtcBuildBVH(&args) };
if root.is_null() {
return Err(match self.device.get_error() {
Error::NONE => Error::UNKNOWN,
e => e,
});
}
let root = NonNull::new(root as *mut B::Node<'id>).map(|ptr| NodePtr {
ptr,
_brand: PhantomData,
});
Ok(f(BvhResult {
root,
_brand: PhantomData,
}))
}
}
fn bounds_of(p: &BuildPrimitive) -> Bounds {
Bounds {
lower_x: p.lower_x,
lower_y: p.lower_y,
lower_z: p.lower_z,
align0: 0.0,
upper_x: p.upper_x,
upper_y: p.upper_y,
upper_z: p.upper_z,
align1: 0.0,
}
}
fn set_axis_upper(b: &mut Bounds, dim: u32, pos: f32) {
match dim {
0 => b.upper_x = pos,
1 => b.upper_y = pos,
_ => b.upper_z = pos,
}
}
fn set_axis_lower(b: &mut Bounds, dim: u32, pos: f32) {
match dim {
0 => b.lower_x = pos,
1 => b.lower_y = pos,
_ => b.lower_z = pos,
}
}