fix EntityRef::is_alive being able to panic, and add EntityRef::exists

2 files changed, 34 insertions, 6 deletions

diff --git a/azalea/src/entity_ref/mod.rs b/azalea/src/entity_ref/mod.rs
index 3e09d336..2f49d975 100644
--- a/azalea/src/entity_ref/mod.rs
+++ b/azalea/src/entity_ref/mod.rs
@@ -7,7 +7,7 @@ use azalea_registry::builtin::EntityKind;
 use bevy_ecs::{
     component::Component,
     entity::Entity,
-    query::{QueryData, QueryItem},
+    query::{QueryData, QueryEntityError, QueryItem},
 };
 use parking_lot::MappedRwLockReadGuard;
 
@@ -20,7 +20,7 @@ use crate::Client;
 ///
 /// Most functions on `EntityRef` that return a value will result in a panic if
 /// the client has despawned, so if your code involves waiting, you should check
-/// [`Self::is_alive`] before calling those functions.
+/// [`Self::is_alive`] or [`Self::exists`] before calling those functions.
 ///
 /// Also, since `EntityRef` stores the [`Client`] alongside the entity, this
 /// means that it supports interactions such as [`Self::attack`].
@@ -83,11 +83,23 @@ impl EntityRef {
     ///
     /// # Panics
     ///
-    /// This will panic if the entity is missing a component required by the
-    /// query.
+    /// This will panic if the entity doesn't exist or is missing a component
+    /// required by the query. Consider using [`Self::try_query_self`] to
+    /// avoid this.
     pub fn query_self<D: QueryData, R>(&self, f: impl FnOnce(QueryItem<D>) -> R) -> R {
         self.client.query_entity(self.entity, f)
     }
+
+    /// Query the ECS for data from the entity, or return an error if the query
+    /// fails.
+    ///
+    /// Also see [`Self::query_self`].
+    pub fn try_query_self<D: QueryData, R>(
+        &self,
+        f: impl FnOnce(QueryItem<D>) -> R,
+    ) -> Result<R, QueryEntityError> {
+        self.client.try_query_entity(self.entity, f)
+    }
 }
 
 impl Debug for EntityRef {
diff --git a/azalea/src/entity_ref/shared_impls.rs b/azalea/src/entity_ref/shared_impls.rs
index 3c24b836..1dafbde2 100644
--- a/azalea/src/entity_ref/shared_impls.rs
+++ b/azalea/src/entity_ref/shared_impls.rs
@@ -164,6 +164,8 @@ impl_entity_functions! {
     /// You should avoid using this if you have auto-respawn enabled (which is
     /// the default), instead consider watching for
     /// [`Event::Death`](crate::Event::Death) instead.
+    ///
+    /// Also see [`Self::exists`].
     EntityRef:
     /// Returns whether the entity is alive and hasn't despawned.
     ///
@@ -171,8 +173,22 @@ impl_entity_functions! {
     /// entity is despawned. Because of this, it may be useful to check `is_alive`
     /// before calling functions that request data from the world.
     ///
-    /// Also see [`Client::is_alive`].
+    /// Also see [`Client::is_alive`] and [`Self::exists`].
     pub fn is_alive(&self) -> bool {
-        self.query_self::<Option<&Dead>, _>(|dead| dead.is_none())
+        self.try_query_self::<Option<&Dead>, _>(|dead| dead.is_none()).unwrap_or(false)
+    }
+
+    Client:
+    /// Returns whether the client is in the world (has been assigned an entity ID).
+    ///
+    /// Like [`Self::is_alive`], this will not panic.
+    EntityRef:
+    /// Returns whether the entity is in the world and hasn't despawned.
+    ///
+    /// Like [`Self::is_alive`], this will not panic.
+    ///
+    /// Also see [`Client::exists`].
+    pub fn exists(&self) -> bool {
+        self.try_query_self::<Option<&MinecraftEntityId>, _>(|entity_id| entity_id.is_some()).unwrap_or(false)
     }
 }