From 7b84235a9be5bdc7c05873467ad8310b57448d79 Mon Sep 17 00:00:00 2001
From: mat <git@matdoes.dev>
Date: Mon, 29 Dec 2025 13:53:54 -1345
Subject: fix EntityRef::is_alive being able to panic, and add
 EntityRef::exists

---
 azalea/src/client_impl/entity_query.rs | 22 ++++++++++++++++------
 azalea/src/entity_ref/mod.rs           | 20 ++++++++++++++++----
 azalea/src/entity_ref/shared_impls.rs  | 20 ++++++++++++++++++--
 3 files changed, 50 insertions(+), 12 deletions(-)

diff --git a/azalea/src/client_impl/entity_query.rs b/azalea/src/client_impl/entity_query.rs
index 8de48478..268eaaf1 100644
--- a/azalea/src/client_impl/entity_query.rs
+++ b/azalea/src/client_impl/entity_query.rs
@@ -84,17 +84,27 @@ impl Client {
     /// # Panics
     ///
     /// This will panic if the client is missing a component required by the
-    /// query.
+    /// 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 {
-        let mut ecs = self.ecs.write();
-        let mut qs = ecs.query::<D>();
-        let res = qs.get_mut(&mut ecs, self.entity).unwrap_or_else(|_| {
+        self.try_query_self::<D, R>(f).unwrap_or_else(|_| {
             panic!(
                 "`Client::query_self` failed when querying for {:?}",
                 any::type_name::<D>()
             )
-        });
-        f(res)
+        })
+    }
+
+    /// Query the ECS for data from our client entity, or return `None` if the
+    /// query failed.
+    ///
+    /// Also see [`Self::query_self`].
+    pub fn try_query_self<D: QueryData, R>(
+        &self,
+        f: impl FnOnce(QueryItem<D>) -> R,
+    ) -> Result<R, QueryEntityError> {
+        let mut ecs = self.ecs.write();
+        let mut qs = ecs.query::<D>();
+        qs.get_mut(&mut ecs, self.entity).map(f)
     }
 
     /// Query the ECS for data from an entity.
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)
     }
 }
-- 
cgit v1.2.3