refactor: node id type

Author: technobaboo

codegen/src/lib.rs

@@ -55,7 +55,7 @@ fn codegen_protocol(protocol: &'static str) -> proc_macro::TokenStream {
 		.map(|p| {
 			let node_id = p.node_id;
 			let node_id = quote! {
-				const INTERFACE_NODE_ID: u64 = #node_id;
+				const INTERFACE_NODE_ID: crate::nodes::Id = crate::nodes::Id(#node_id);
 			};
 			let aspect = generate_aspect(&Aspect {
 				name: "interface".to_string(),
@@ -698,7 +698,7 @@ fn generate_argument_type(
 				quote!(stardust_xr::values::Map<String, #t>)
 			}
 		}
-		ArgumentType::NodeID => quote!(u64),
+		ArgumentType::NodeID => quote!(crate::nodes::Id),
 		ArgumentType::Datamap => {
 			if !owned {
 				quote!(&stardust_xr::values::Datamap)

src/core/client.rs

@@ -3,7 +3,7 @@ use super::{
 	scenegraph::Scenegraph,
 };
 use crate::{
-	core::{registry::OwnedRegistry, task},
+	core::{Id, registry::OwnedRegistry, task},
 	nodes::{
 		Node, audio, drawable, fields, input, items,
 		root::{ClientState, Root},
@@ -202,12 +202,12 @@ impl Client {
 		Some(ClientStateParsed::from_deserialized(self, internal))
 	}
 
-	pub fn generate_id(&self) -> u64 {
-		self.id_counter.inc() as u64
+	pub fn generate_id(&self) -> Id {
+		Id(self.id_counter.inc() as u64)
 	}
 
 	#[inline]
-	pub fn get_node(&self, name: &'static str, id: u64) -> Result<Arc<Node>> {
+	pub fn get_node(&self, name: &'static str, id: Id) -> Result<Arc<Node>> {
 		self.scenegraph
 			.get_node(id)
 			.ok_or_else(|| eyre!("{} not found", name))

src/core/client_state.rs

@@ -1,5 +1,8 @@
 use super::client::{Client, get_env};
-use crate::nodes::{Node, root::ClientState, spatial::Spatial};
+use crate::{
+	core::Id,
+	nodes::{Node, root::ClientState, spatial::Spatial},
+};
 use glam::Mat4;
 use parking_lot::Mutex;
 use rustc_hash::FxHashMap;
@@ -51,7 +54,7 @@ impl ClientStateParsed {
 				.collect(),
 		}
 	}
-	fn spatial_transform(client: &Client, id: u64) -> Option<Mat4> {
+	fn spatial_transform(client: &Client, id: Id) -> Option<Mat4> {
 		let node = client.scenegraph.get_node(id)?;
 		let spatial = node.get_aspect::<Spatial>().ok()?;
 		Some(spatial.global_transform())
@@ -90,7 +93,7 @@ impl ClientStateParsed {
 		}
 		ClientState {
 			data: self.data.clone(),
-			root: 0,
+			root: Id(0),
 			spatial_anchors: self
 				.spatial_anchors
 				.iter()

src/core/id.rs

@@ -0,0 +1,28 @@
+use serde::{Deserialize, Serialize, Serializer};
+use std::fmt::Display;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct Id(pub u64);
+impl From<u64> for Id {
+	fn from(id: u64) -> Self {
+		Self(id)
+	}
+}
+impl Display for Id {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		write!(f, "{}", self.0)
+	}
+}
+impl<'de> Deserialize<'de> for Id {
+	fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+	where
+		D: serde::Deserializer<'de>,
+	{
+		u64::deserialize(deserializer).map(Id)
+	}
+}
+impl Serialize for Id {
+	fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
+		self.0.serialize(serializer)
+	}
+}

src/core/mod.rs

@@ -5,7 +5,10 @@ pub mod color;
 pub mod delta;
 pub mod entity_handle;
 pub mod error;
+mod id;
 pub mod registry;
 pub mod resource;
 pub mod scenegraph;
 pub mod task;
+
+pub use id::Id;

src/core/scenegraph.rs

@@ -1,5 +1,6 @@
 use crate::{
 	core::{
+		Id,
 		client::Client,
 		error::{Result, ServerError},
 	},
@@ -78,7 +79,7 @@ impl std::fmt::Debug for MethodResponseSender {
 #[derive(Default)]
 pub struct Scenegraph {
 	pub(super) client: OnceLock<Weak<Client>>,
-	nodes: Mutex<FxHashMap<u64, Arc<Node>>>,
+	nodes: Mutex<FxHashMap<Id, Arc<Node>>>,
 }
 
 impl Scenegraph {
@@ -96,13 +97,13 @@ impl Scenegraph {
 		self.nodes.lock().insert(node.get_id(), node);
 	}
 
-	pub fn get_node(&self, node: u64) -> Option<Arc<Node>> {
+	pub fn get_node(&self, node: Id) -> Option<Arc<Node>> {
 		let node = self.nodes.lock().get(&node)?.clone();
 		get_original(node, true)
 	}
 
-	pub fn remove_node(&self, node: u64) -> Option<Arc<Node>> {
-		debug!(node, "Remove node");
+	pub fn remove_node(&self, node: Id) -> Option<Arc<Node>> {
+		debug!(node = node.0, "Remove node");
 		self.nodes.lock().remove(&node)
 	}
 }
@@ -119,7 +120,7 @@ impl scenegraph::Scenegraph for Scenegraph {
 			return Err(ScenegraphError::NodeNotFound);
 		};
 		debug_span!("Handle signal", aspect_id, node_id, method).in_scope(|| {
-			self.get_node(node_id)
+			self.get_node(Id(node_id))
 				.ok_or(ScenegraphError::NodeNotFound)?
 				.send_local_signal(
 					client,
@@ -146,7 +147,7 @@ impl scenegraph::Scenegraph for Scenegraph {
 			return;
 		};
 		debug!(aspect_id, node_id, method, "Handle method");
-		let Some(node) = self.get_node(node_id) else {
+		let Some(node) = self.get_node(Id(node_id)) else {
 			response.send(Err(ScenegraphError::NodeNotFound));
 			return;
 		};

src/nodes/alias.rs

@@ -1,6 +1,6 @@
 use super::{Aspect, AspectIdentifier, Node};
 use crate::core::{
-	client::Client, error::Result, registry::Registry, scenegraph::MethodResponseSender,
+	Id, client::Client, error::Result, registry::Registry, scenegraph::MethodResponseSender,
 };
 use std::{
 	ops::Add,
@@ -46,7 +46,7 @@ impl Alias {
 	pub fn create_with_id(
 		original: &Arc<Node>,
 		client: &Arc<Client>,
-		new_id: u64,
+		new_id: Id,
 		info: AliasInfo,
 		list: Option<&AliasList>,
 	) -> Result<Arc<Node>> {

src/nodes/audio.rs

@@ -1,5 +1,6 @@
 use super::spatial::SpatialNode;
 use super::{Aspect, AspectIdentifier, Node};
+use crate::core::Id;
 use crate::core::client::Client;
 use crate::core::entity_handle::EntityHandle;
 use crate::core::error::Result;
@@ -150,7 +151,7 @@ impl InterfaceAspect for Interface {
 	fn create_sound(
 		_node: Arc<Node>,
 		calling_client: Arc<Client>,
-		id: u64,
+		id: Id,
 		parent: Arc<Node>,
 		transform: Transform,
 		resource: ResourceID,

src/nodes/drawable/mod.rs

@@ -8,7 +8,7 @@ use super::{
 	Aspect, AspectIdentifier, Node,
 	spatial::{Spatial, Transform},
 };
-use crate::core::{client::Client, error::Result, resource::get_resource_file};
+use crate::core::{Id, client::Client, error::Result, resource::get_resource_file};
 use crate::nodes::spatial::SPATIAL_ASPECT_ALIAS_INFO;
 use color_eyre::eyre::eyre;
 use model::ModelPart;
@@ -88,7 +88,7 @@ impl InterfaceAspect for Interface {
 	fn create_lines(
 		_node: Arc<Node>,
 		calling_client: Arc<Client>,
-		id: u64,
+		id: Id,
 		parent: Arc<Node>,
 		transform: Transform,
 		lines: Vec<Line>,
@@ -106,7 +106,7 @@ impl InterfaceAspect for Interface {
 	fn load_model(
 		_node: Arc<Node>,
 		calling_client: Arc<Client>,
-		id: u64,
+		id: Id,
 		parent: Arc<Node>,
 		transform: Transform,
 		model: ResourceID,
@@ -123,7 +123,7 @@ impl InterfaceAspect for Interface {
 	fn create_text(
 		_node: Arc<Node>,
 		calling_client: Arc<Client>,
-		id: u64,
+		id: Id,
 		parent: Arc<Node>,
 		transform: Transform,
 		text: String,

src/nodes/drawable/model.rs

@@ -1,4 +1,5 @@
 use super::{MODEL_PART_ASPECT_ALIAS_INFO, MaterialParameter, ModelAspect, ModelPartAspect};
+use crate::core::Id;
 use crate::core::bevy_channel::{BevyChannel, BevyChannelReader};
 use crate::core::client::Client;
 use crate::core::color::ColorConvert as _;
@@ -624,7 +625,7 @@ impl ModelAspect for Model {
 	fn bind_model_part(
 		node: Arc<Node>,
 		calling_client: Arc<Client>,
-		id: u64,
+		id: Id,
 		part_path: String,
 	) -> Result<()> {
 		let model = node.get_aspect::<Model>()?;