fix: getting bounding boxes before model is fully loaded now waits until the model and bounding boxes are loaded

Signed-off-by: Schmarni <marnistromer@gmail.com>

src/nodes/drawable/lines.rs

@@ -16,7 +16,7 @@ use bevy::{
 	pbr::{ExtendedMaterial, MaterialExtension},
 	prelude::*,
 	render::{
-		mesh::{Indices, MeshAabb, PrimitiveTopology},
+		mesh::{Indices, PrimitiveTopology, VertexAttributeValues},
 		primitives::Aabb,
 		render_resource::{AsBindGroup, ShaderRef},
 		view::VisibilitySystems,
@@ -28,6 +28,7 @@ use std::sync::{
 	Arc, OnceLock,
 	atomic::{AtomicBool, Ordering},
 };
+use tokio::sync::Notify;
 
 type LineMaterial = ExtendedMaterial<BevyMaterial, LineExtension>;
 const LINE_SHADER_HANDLE: Handle<Shader> = weak_handle!("7d28aa5a-3abd-43bb-b0e9-0de8b81b650d");
@@ -91,10 +92,8 @@ fn build_line_mesh(
 	mut materials: ResMut<Assets<LineMaterial>>,
 	query: Query<(&GlobalTransform, &InheritedVisibility)>,
 ) {
-	for lines in LINES_REGISTRY
-		.get_valid_contents()
-		.into_iter()
-		// .filter(|l| l.gen_mesh.load(Ordering::Relaxed))
+	for lines in LINES_REGISTRY.get_valid_contents().into_iter()
+	// .filter(|l| l.gen_mesh.load(Ordering::Relaxed))
 	{
 		lines.gen_mesh.store(false, Ordering::Relaxed);
 		let mut vertex_positions = Vec::<Vec3>::new();
@@ -107,7 +106,8 @@ fn build_line_mesh(
 			continue;
 		};
 		if lines_data.is_empty() {
-			*lines.bounds.lock() = Aabb::default();
+			*lines.bounds.lock() = Some(Aabb::default());
+			lines.setup_complete.notify_waiters();
 			match lines.entity.get() {
 				Some(e) => cmds.entity(**e),
 				None => {
@@ -261,9 +261,21 @@ fn build_line_mesh(
 				e
 			}
 		};
-		if let Some(aabb) = mesh.compute_aabb() {
-			*lines.bounds.lock() = aabb;
-			entity.insert(aabb);
+		if let Some(VertexAttributeValues::Float32x3(values)) =
+			mesh.attribute(Mesh::ATTRIBUTE_POSITION)
+		{
+			let global_to_local = transform.affine().inverse();
+			let local_aabb = Aabb::enclosing(
+				values
+					.iter()
+					.map(|p| global_to_local.transform_point3(Vec3::from_slice(p))),
+			)
+			.unwrap_or_default();
+			let global_aabb =
+				Aabb::enclosing(values.iter().map(|p| Vec3::from_slice(p))).unwrap_or_default();
+			*lines.bounds.lock() = Some(local_aabb);
+			lines.setup_complete.notify_waiters();
+			entity.insert(global_aabb);
 		}
 		entity
 			.insert(Mesh3d(meshes.add(mesh)))
@@ -313,7 +325,8 @@ pub struct Lines {
 	data: Mutex<Vec<Line>>,
 	gen_mesh: AtomicBool,
 	entity: OnceLock<EntityHandle>,
-	bounds: Mutex<Aabb>,
+	bounds: Mutex<Option<Aabb>>,
+	setup_complete: Notify,
 }
 impl Lines {
 	pub fn add_to(node: &Arc<Node>, lines: Vec<Line>) -> Result<Arc<Lines>> {
@@ -322,10 +335,19 @@ impl Lines {
 			.unwrap()
 			.bounding_box_calc
 			.set(|node| {
-				node.get_aspect::<Lines>()
-					.ok()
-					.map(|v| *v.bounds.lock())
-					.unwrap_or_default()
+				Box::pin(async {
+					let Ok(lines) = node.get_aspect::<Lines>() else {
+						return Default::default();
+					};
+					let bounds = *lines.bounds.lock();
+					match bounds {
+						Some(aabb) => aabb,
+						None => {
+							lines.setup_complete.notified().await;
+							lines.bounds.lock().unwrap_or_default()
+						}
+					}
+				})
 			});
 
 		let lines = LINES_REGISTRY.add(Lines {
@@ -333,7 +355,8 @@ impl Lines {
 			data: Mutex::new(lines),
 			gen_mesh: AtomicBool::new(true),
 			entity: OnceLock::new(),
-			bounds: Mutex::new(Aabb::default()),
+			bounds: Mutex::new(Some(Aabb::default())),
+			setup_complete: Notify::new(),
 		});
 		node.add_aspect_raw(lines.clone());
 

src/nodes/drawable/model.rs

@@ -25,6 +25,7 @@ use std::hash::{Hash, Hasher};
 use std::path::PathBuf;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, OnceLock, Weak};
+use tokio::sync::Notify;
 
 static LOAD_MODEL: BevyChannel<(Arc<Model>, PathBuf)> = BevyChannel::new();
 
@@ -242,10 +243,12 @@ fn gen_model_parts(
 					)
 					.unwrap_or_default();
 					_ = spatial.bounding_box_calc.set(move |n| {
-						n.get_aspect::<ModelPart>()
-							.ok()
-							.and_then(|v| v.bounds.get().copied())
-							.unwrap_or_default()
+						Box::pin(async {
+							n.get_aspect::<ModelPart>()
+								.ok()
+								.and_then(|v| v.bounds.get().copied())
+								.unwrap_or_default()
+						})
 					});
 					let _ = spatial.set_spatial_parent(&parent_spatial);
 					spatial.set_local_transform(transform.compute_matrix());
@@ -271,6 +274,8 @@ fn gen_model_parts(
 		model
 			.spatial
 			.set_entity(model.bevy_scene_entity.get().unwrap().clone());
+		model.setup_complete.store(true, Ordering::Relaxed);
+		model.setup_complete_notify.notify_waiters();
 	}
 }
 
@@ -534,6 +539,8 @@ pub struct Model {
 	bevy_scene_entity: OnceLock<EntityHandle>,
 	parts: OnceLock<Vec<Arc<ModelPart>>>,
 	pre_bound_parts: Mutex<Vec<Arc<ModelPart>>>,
+	setup_complete: AtomicBool,
+	setup_complete_notify: Notify,
 }
 impl Model {
 	pub fn add_to(node: &Arc<Node>, resource_id: ResourceID) -> Result<Arc<Model>> {
@@ -550,6 +557,18 @@ impl Model {
 			bevy_scene_entity: OnceLock::new(),
 			pre_bound_parts: Mutex::default(),
 			parts: OnceLock::new(),
+			setup_complete_notify: Notify::new(),
+			setup_complete: AtomicBool::new(false),
+		});
+		model.spatial.bounding_box_calc.set(|n| {
+			Box::pin(async {
+				if let Ok(model) = n.get_aspect::<Model>()
+					&& !model.setup_complete.load(Ordering::Relaxed)
+				{
+					model.setup_complete_notify.notified().await;
+				}
+				Aabb::default()
+			})
 		});
 		LOAD_MODEL
 			.send((model.clone(), pending_model_path))

src/nodes/spatial/mod.rs

@@ -20,6 +20,7 @@ use mint::Vector3;
 use parking_lot::{Mutex, RwLock};
 use rustc_hash::FxHashMap;
 use std::fmt::Debug;
+use std::pin::Pin;
 use std::sync::atomic::Ordering;
 use std::sync::{Arc, OnceLock, Weak};
 use std::{f32, ptr};
@@ -142,7 +143,8 @@ pub struct Spatial {
 	transform: RwLock<Mat4>,
 	zone: RwLock<Weak<Zone>>,
 	children: Registry<Spatial>,
-	pub bounding_box_calc: OnceLock<fn(&Node) -> Aabb>,
+	pub bounding_box_calc:
+		OnceLock<for<'a> fn(&'a Node) -> Pin<Box<dyn Future<Output = Aabb> + 'a + Send + Sync>>>,
 }
 
 impl Spatial {
@@ -199,18 +201,17 @@ impl Spatial {
 	}
 
 	// the output bounds are probably way bigger than they need to be
-	pub fn get_bounding_box(&self) -> Aabb {
+	pub async fn get_bounding_box(&self) -> Aabb {
 		let Some(node) = self.node() else {
 			return Aabb::default();
 		};
-		let mut bounds = self
-			.bounding_box_calc
-			.get()
-			.map(|b| (b)(&node))
-			.unwrap_or_default();
+		let mut bounds = match self.bounding_box_calc.get() {
+			Some(f) => f(&node).await,
+			None => Aabb::default(),
+		};
 		for child in self.children.get_valid_contents() {
 			let mat = child.local_transform();
-			let child_aabb = child.get_bounding_box();
+			let child_aabb = Box::pin(child.get_bounding_box()).await;
 			bounds = Aabb::enclosing([
 				bounds.min().into(),
 				bounds.max().into(),
@@ -480,7 +481,7 @@ impl SpatialRefAspect for SpatialRef {
 		_calling_client: Arc<Client>,
 	) -> Result<BoundingBox> {
 		let this_spatial = node.get_aspect::<Spatial>()?;
-		let bounds = this_spatial.get_bounding_box();
+		let bounds = this_spatial.get_bounding_box().await;
 
 		Ok(BoundingBox {
 			center: Vec3::from(bounds.center).into(),
@@ -496,7 +497,7 @@ impl SpatialRefAspect for SpatialRef {
 		let this_spatial = node.get_aspect::<Spatial>()?;
 		let relative_spatial = relative_to.get_aspect::<Spatial>()?;
 		let mat = Spatial::space_to_space_matrix(Some(&this_spatial), Some(&relative_spatial));
-		let bb = this_spatial.get_bounding_box();
+		let bb = this_spatial.get_bounding_box().await;
 		let bounds = Aabb::enclosing([
 			mat.transform_point3(bb.min().into()),
 			mat.transform_point3(bb.max().into()),