refactor: probably get models mostly working

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

src/nodes/drawable/model.rs

@@ -1,5 +1,5 @@
-use super::{MaterialParameter, ModelAspect, ModelPartAspect, MODEL_PART_ASPECT_ALIAS_INFO};
 use crate::bail;
+use crate::bevy_plugin::MainWorldEntity;
 use crate::core::client::Client;
 use crate::core::error::Result;
 use crate::core::registry::Registry;
@@ -8,24 +8,35 @@ use crate::nodes::alias::{Alias, AliasList};
 use crate::nodes::spatial::Spatial;
 use crate::nodes::{Aspect, Node};
 use crate::DefaultMaterial;
-use bevy::app::Plugin;
+use bevy::app::{Plugin, PostUpdate};
 use bevy::asset::Handle;
+use bevy::asset::{AssetServer, Assets};
 use bevy::color::{Alpha, Color, LinearRgba, Srgba};
+use bevy::core::Name;
+use bevy::gltf::GltfAssetLabel;
+use bevy::math::bounding::Aabb3d;
+use bevy::pbr::MeshMaterial3d;
 use bevy::prelude::AlphaMode;
+use bevy::prelude::{
+	Children, Commands, Component, Deref, Entity, HierarchyQueryExt, Mesh3d, Parent, Query, Res,
+	ResMut, Resource, Transform, Visibility, With, Without,
+};
 use bevy::reflect::{GetField, PartialReflect, Reflect};
+use bevy::render::primitives::Aabb;
+use bevy::scene::SceneRoot;
 use color_eyre::eyre::eyre;
-use glam::{Mat4, Vec2, Vec3};
-use once_cell::sync::{Lazy, OnceCell};
+use glam::{Vec2, Vec3};
+use once_cell::sync::OnceCell;
 use parking_lot::Mutex;
 use rustc_hash::FxHashMap;
 use stardust_xr::values::ResourceID;
-use tracing::warn;
+use tracing::{error, warn};
 
 use std::ffi::OsStr;
 use std::ops::Deref;
+use std::path::PathBuf;
 use std::sync::{Arc, Weak};
 
-
 static MODEL_REGISTRY: Registry<Model> = Registry::new();
 static HOLDOUT_MATERIAL: OnceCell<Arc<DefaultMaterial>> = OnceCell::new();
 
@@ -114,167 +125,110 @@ impl MaterialParameter {
 				else {
 					return;
 				};
-				// if let Ok(tex) = Tex::from_file(texture_path, true, None) {
-				// 	params.set_texture(parameter_name, &tex);
-				// }
+				error!("TODO: implement texture changing");
 			}
 		}
 	}
 }
 
 pub struct ModelPart {
-	id: i32,
+	entity: MainWorldEntity,
 	path: String,
 	space: Arc<Spatial>,
 	model: Weak<Model>,
 	pending_material_parameters: Mutex<FxHashMap<String, MaterialParameter>>,
-	pending_material_replacement: Mutex<Option<Arc<Handle<DefaultMaterial>>>>,
+	pending_material_replacement: Mutex<Option<Arc<DefaultMaterial>>>,
 	aliases: AliasList,
 }
 
+#[derive(Component, Clone)]
+pub struct StardustModel(Arc<Model>);
+#[derive(Component, Clone)]
+pub struct UnprocessedModel;
 pub struct StardustModelPlugin;
 impl Plugin for StardustModelPlugin {
-	fn build(&self, app: &mut bevy::prelude::App) {}
+	fn build(&self, app: &mut bevy::prelude::App) {
+		let (tx, rx) = crossbeam_channel::unbounded();
+		LOAD_MODEL_SENDER.set(tx);
+		app.insert_resource(LoadModelReader(rx));
+		app.add_systems(PostUpdate, create_model_parts_for_loaded_models);
+	}
+}
+static LOAD_MODEL_SENDER: OnceCell<crossbeam_channel::Sender<(PathBuf, Arc<Model>)>> =
+	OnceCell::new();
+#[derive(Resource, Deref)]
+struct LoadModelReader(crossbeam_channel::Receiver<(PathBuf, Arc<Model>)>);
+
+fn update_models(mut query: Query<(&StardustModel, &mut Visibility, &mut Transform)>) {
+	for (model, mut vis, mut transform) in query.iter_mut() {
+		*transform = Transform::from_matrix(model.0.space.global_transform());
+		if let Some(node) = model.0.space.node() {
+			*vis = match node.enabled() {
+				true => Visibility::Inherited,
+				false => Visibility::Hidden,
+			}
+		}
+	}
 }
 
-impl ModelPart {
-	fn create_for_model(model: &Arc<Model>, sk_model: &SKModel) {
-		// The value isn't even used?!
-		HOLDOUT_MATERIAL.get_or_init(|| {
-			let mut mat = DefaultMaterial {
-				unlit: true,
-				alpha_mode: AlphaMode::Opaque,
-				base_color: Srgba::BLACK.with_alpha(0.0).into(),
-				..Default::default()
-			};
-			Arc::new(mat)
-		});
-
-		let nodes = sk_model.get_nodes();
-		for part in nodes.all() {
-			ModelPart::create(model, &part);
-		}
+fn load_models(rx: Res<LoadModelReader>, mut cmds: Commands, asset_server: Res<AssetServer>) {
+	for (path, model) in rx.try_iter() {
+		let handle = asset_server.load(GltfAssetLabel::Scene(0).from_asset(path));
+		let entity = cmds
+			.spawn((
+				SceneRoot(handle),
+				StardustModel(model.clone()),
+				UnprocessedModel,
+			))
+			.id();
+		model.entity.set(entity);
 	}
+}
 
-	fn create(model: &Arc<Model>, part: &stereokit_rust::model::ModelNode) -> Option<Arc<Self>> {
-		let mut parts = model.parts.lock();
-		let parent_part = part
-			.get_parent()
-			.and_then(|part| parts.iter().find(|p| p.id == *part.get_id()));
-
-		let stardust_model_part = model.space.node()?;
-		let client = stardust_model_part.get_client()?;
-		let mut part_path = parent_part
-			.map(|n| n.path.clone() + "/")
-			.unwrap_or_default();
-		part_path += part.get_name().unwrap();
-
-		let node = client.scenegraph.add_node(Node::generate(&client, false));
-		let spatial_parent = parent_part
-			.map(|n| n.space.clone())
-			.unwrap_or_else(|| model.space.clone());
-
-		let local_transform = unsafe { part.get_local_transform().m };
-		let space = Spatial::add_to(
-			&node,
-			Some(spatial_parent),
-			Mat4::from_cols_array(&local_transform),
-			false,
-		);
-
-		let _ = space.bounding_box_calc.set(|node| {
-			let Ok(model_part) = node.get_aspect::<ModelPart>() else {
-				return Bounds::default();
+fn update_model_parts(
+	models: Query<&StardustModel>,
+	mut mats: ResMut<Assets<DefaultMaterial>>,
+	mut part_query: Query<(
+		&mut Transform,
+		&mut MeshMaterial3d<DefaultMaterial>,
+		&mut Mesh3d,
+	)>,
+) {
+	for model in &models {
+		let model = &model.0;
+		for part in model.parts.lock().iter() {
+			let Ok((mut transform, mut mat, mut _mesh)) = part_query.get_mut(*part.entity) else {
+				continue;
 			};
-			let Some(model) = model_part.model.upgrade() else {
-				return Bounds::default();
-			};
-			let Some(sk_model) = model.sk_model.get() else {
-				return Bounds::default();
-			};
-			let model_nodes = sk_model.get_nodes();
-			let Some(model_node) = model_nodes.get_index(model_part.id) else {
-				return Bounds::default();
-			};
-			let Some(sk_mesh) = model_node.get_mesh() else {
-				return Bounds::default();
-			};
-			sk_mesh.get_bounds()
-		});
+			*transform = Transform::from_matrix(Spatial::space_to_space_matrix(
+				Some(&part.space),
+				Some(&model.space),
+			));
+			if let Some(material_replacement) = part.pending_material_replacement.lock().take() {
+				let material = material_replacement.deref().clone();
+				let mat_handle = mats.add(material);
+				mat.0 = mat_handle;
+			}
+			// todo: find all materials with identical parameters and batch them into 1 material again
+			'mat_params: {
+				let mut material_parameters = part.pending_material_parameters.lock();
+				if !material_parameters.is_empty() {
+					let Some(material) = mats.get(&mat.0) else {
+						break 'mat_params;
+					};
 
-		let model_part = Arc::new(ModelPart {
-			id: *part.get_id(),
-			path: part_path,
-			space,
-			model: Arc::downgrade(model),
-			pending_material_parameters: Mutex::new(FxHashMap::default()),
-			pending_material_replacement: Mutex::new(None),
-			aliases: AliasList::default(),
-		});
-		node.add_aspect_raw(model_part.clone());
-		parts.push(model_part.clone());
-		Some(model_part)
-	}
-
-	pub fn replace_material(&self, replacement: Arc<MaterialWrapper>) {
-		let shared_material = MATERIAL_REGISTRY.add_or_get(replacement);
-		self.pending_material_replacement
-			.lock()
-			.replace(shared_material);
-	}
-	/// only to be run on the main thread
-	pub fn replace_material_now(&self, replacement: &Material) {
-		let Some(model) = self.model.upgrade() else {
-			return;
-		};
-		let Some(sk_model) = model.sk_model.get() else {
-			return;
-		};
-		let nodes = sk_model.get_nodes();
-		let Some(mut part) = nodes.get_index(self.id) else {
-			return;
-		};
-		let shared_material =
-			MATERIAL_REGISTRY.add_or_get(Arc::new(MaterialWrapper(replacement.copy())));
-		part.material(&shared_material.0);
-	}
-
-	fn update(&self) {
-		let Some(model) = self.model.upgrade() else {
-			return;
-		};
-		let Some(sk_model) = model.sk_model.get() else {
-			return;
-		};
-		let Some(node) = model.space.node() else {
-			return;
-		};
-		let nodes = sk_model.get_nodes();
-		let Some(mut part) = nodes.get_index(self.id) else {
-			return;
-		};
-		part.model_transform(Spatial::space_to_space_matrix(
-			Some(&self.space),
-			Some(&model.space),
-		));
-
-		let Some(client) = node.get_client() else {
-			return;
-		};
-
-		if let Some(material_replacement) = self.pending_material_replacement.lock().take() {
-			part.material(&material_replacement.0);
-		}
-
-		'mat_params: {
-			let mut material_parameters = self.pending_material_parameters.lock();
-			if !material_parameters.is_empty() {
-				let Some(material) = part.get_material() else {
-					break 'mat_params;
-				};
-				let new_material = material.copy();
-				for (parameter_name, parameter_value) in material_parameters.drain() {
-					parameter_value.apply_to_material(&client, &new_material, &parameter_name);
+					let mut new_material = material.clone();
+					let Some(client) = part.space.node().and_then(|v| v.get_client()) else {
+						return;
+					};
+					for (parameter_name, parameter_value) in material_parameters.drain() {
+						parameter_value.apply_to_material(
+							&client,
+							&mut new_material,
+							&parameter_name,
+						);
+					}
+					mat.0 = mats.add(new_material);
 				}
 
 				let shared_material =
@@ -284,6 +238,85 @@ impl ModelPart {
 		}
 	}
 }
+
+fn get_path(
+	entity: Entity,
+	query: &Query<(Entity, Option<&Parent>, &Transform, &Name, &Aabb), Without<SceneRoot>>,
+) -> Option<String> {
+	let (_, parent, _, name, _) = query.get(entity).ok()?;
+	let next = parent.and_then(|p| get_path(p.get(), query));
+	match next {
+		Some(next) => Some(format!("{name}/{next}")),
+		None => Some(name.to_string()),
+	}
+}
+
+fn create_model_parts_for_loaded_models(
+	query: Query<(Entity, &StardustModel), With<UnprocessedModel>>,
+	children: Query<&Children>,
+	gltf_model_parts: Query<
+		(Entity, Option<&Parent>, &Transform, &Name, &Aabb),
+		Without<SceneRoot>,
+	>,
+	mut cmds: Commands,
+) {
+	for (entity, model) in &query {
+		let model = &model.0;
+		let mut parts = model.parts.lock();
+		cmds.entity(entity).remove::<UnprocessedModel>();
+		for (entity, parent, transform, name, aabb) in children
+			.iter_descendants(entity)
+			.filter_map(|e| gltf_model_parts.get(e).ok())
+		{
+			let parent_part = parent
+				.and_then(|e| gltf_model_parts.get(e.get()).ok())
+				.and_then(|(e, _, _, _, _)| parts.iter().find(|v| v.entity.0 == e));
+
+			let Some(stardust_model_part) = model.space.node() else {
+				continue;
+			};
+			let Some(client) = stardust_model_part.get_client() else {
+				continue;
+			};
+			let part_path = get_path(entity, &gltf_model_parts).unwrap_or_else(|| name.to_string());
+
+			let node = client.scenegraph.add_node(Node::generate(&client, false));
+			let spatial_parent = parent_part
+				.map(|n| n.space.clone())
+				.unwrap_or_else(|| model.space.clone());
+
+			let space = Spatial::add_to(
+				&node,
+				Some(spatial_parent),
+				transform.compute_matrix(),
+				false,
+			);
+
+			*space.bounding_box_calc.lock() = Aabb3d::new(aabb.center, aabb.half_extents);
+
+			let model_part = Arc::new(ModelPart {
+				entity: MainWorldEntity(entity),
+				path: part_path,
+				space,
+				model: Arc::downgrade(model),
+				pending_material_parameters: Mutex::new(FxHashMap::default()),
+				pending_material_replacement: Mutex::new(None),
+				aliases: AliasList::default(),
+			});
+			<ModelPart as ModelPartAspect>::add_node_members(&node);
+			node.add_aspect_raw(model_part.clone());
+			parts.push(model_part.clone());
+		}
+	}
+}
+
+impl ModelPart {
+	pub fn replace_material(&self, replacement: Arc<DefaultMaterial>) {
+		self.pending_material_replacement
+			.lock()
+			.replace(replacement);
+	}
+}
 impl ModelPartAspect for ModelPart {
 	#[doc = "Set this model part's material to one that cuts a hole in the world. Often used for overlays/passthrough where you want to show the background through an object."]
 	fn apply_holdout_material(node: Arc<Node>, _calling_client: Arc<Client>) -> Result<()> {
@@ -312,7 +345,7 @@ impl ModelPartAspect for ModelPart {
 pub struct Model {
 	space: Arc<Spatial>,
 	_resource_id: ResourceID,
-	sk_model: OnceCell<()>,
+	entity: OnceCell<Entity>,
 	parts: Mutex<Vec<Arc<ModelPart>>>,
 }
 impl Model {
@@ -327,42 +360,17 @@ impl Model {
 		let model = Arc::new(Model {
 			space: node.get_aspect::<Spatial>().unwrap().clone(),
 			_resource_id: resource_id,
-			sk_model: OnceCell::new(),
+			entity: OnceCell::new(),
 			parts: Mutex::new(Vec::default()),
 		});
 		MODEL_REGISTRY.add_raw(&model);
-
-		// technically doing this in anything but the main thread isn't a good idea but dangit we need those model nodes ASAP
-		let sk_model = SKModel::copy(SKModel::from_file(
-			pending_model_path.to_str().unwrap(),
-			None,
-		)?);
-		ModelPart::create_for_model(&model, &sk_model);
-		let _ = model.sk_model.set(sk_model);
+		if let Some(sender) = LOAD_MODEL_SENDER.get() {
+			sender.send((pending_model_path, model.clone()));
+		}
 		node.add_aspect_raw(model.clone());
 		Ok(model)
 	}
-
-	fn draw(&self, token: &MainThreadToken) {
-		let Some(sk_model) = self.sk_model.get() else {
-			return;
-		};
-		let parts = self.parts.lock();
-		for model_node in &*parts {
-			model_node.update();
-		}
-		drop(parts);
-
-		if let Some(node) = self.space.node() {
-			if node.enabled() {
-				sk_model.draw(token, self.space.global_transform(), None, None);
-			}
-		}
-	}
 }
-// TODO: proper hread safety in stereokit_rust (probably just bind stereokit directly)
-unsafe impl Send for Model {}
-unsafe impl Sync for Model {}
 impl ModelAspect for Model {
 	#[doc = "Bind a model part to the node with the ID input."]
 	fn bind_model_part(
@@ -392,9 +400,3 @@ impl Drop for Model {
 		MODEL_REGISTRY.remove(self);
 	}
 }
-
-pub fn draw_all(token: &MainThreadToken) {
-	for model in MODEL_REGISTRY.get_valid_contents() {
-		model.draw(token);
-	}
-}