From 6c7156f70d940a5bf0ee7bcb29a768b7864aea0d Mon Sep 17 00:00:00 2001
From: ErrorNoInternet <errornointernet@envs.net>
Date: Fri, 28 Feb 2025 20:51:19 -0500
Subject: [PATCH] refactor: make some more things async

---
 src/lua/block.rs        |  8 +++---
 src/lua/client/mod.rs   |  2 +-
 src/lua/client/world.rs | 59 ++++++++++++++++++++++++-----------------
 3 files changed, 39 insertions(+), 30 deletions(-)

diff --git a/src/lua/block.rs b/src/lua/block.rs
index a43c73b..1e54bd9 100644
--- a/src/lua/block.rs
+++ b/src/lua/block.rs
@@ -9,7 +9,7 @@ pub fn register_functions(lua: &Lua, globals: &Table) -> Result<()> {
         "get_block_from_state",
         lua.create_function(get_block_from_state)?,
     )?;
-    globals.set("get_block_states", lua.create_function(get_block_states)?)?;
+    globals.set("get_block_states", lua.create_async_function(get_block_states)?)?;
 
     Ok(())
 }
@@ -34,8 +34,8 @@ pub fn get_block_from_state(lua: &Lua, state: u32) -> Result<Option<Table>> {
     Ok(Some(block))
 }
 
-pub fn get_block_states(
-    lua: &Lua,
+pub async fn get_block_states(
+    lua: Lua,
     (block_names, filter_fn): (Vec<String>, Option<Function>),
 ) -> Result<Vec<u16>> {
     let mut matched = Vec::new();
@@ -49,7 +49,7 @@ pub fn get_block_states(
                     p.set("chest_type", b.property::<ChestType>().map(|v| v as u8))?;
                     p.set("facing", b.property::<Facing>().map(|v| v as u8))?;
                     p.set("light_level", b.property::<LightLevel>().map(|v| v as u8))?;
-                    filter_fn.call::<bool>(p.clone())?
+                    filter_fn.call_async::<bool>(p.clone()).await?
                 } else {
                     true
                 })
diff --git a/src/lua/client/mod.rs b/src/lua/client/mod.rs
index dc2ed69..f5bf063 100644
--- a/src/lua/client/mod.rs
+++ b/src/lua/client/mod.rs
@@ -59,6 +59,7 @@ impl UserData for Client {
 
     fn add_methods<M: UserDataMethods<Self>>(m: &mut M) {
         m.add_async_method("attack", interaction::attack);
+        m.add_async_method("find_entities", world::find_entities);
         m.add_async_method("go_to", movement::go_to);
         m.add_async_method("look_at", movement::look_at);
         m.add_async_method("mine", interaction::mine);
@@ -68,7 +69,6 @@ impl UserData for Client {
         m.add_method("chat", chat);
         m.add_method("disconnect", disconnect);
         m.add_method("find_blocks", world::find_blocks);
-        m.add_method("find_entities", world::find_entities);
         m.add_method("get_block_state", world::get_block_state);
         m.add_method("get_fluid_state", world::get_fluid_state);
         m.add_method("set_held_slot", container::set_held_slot);
diff --git a/src/lua/client/world.rs b/src/lua/client/world.rs
index e677a52..bedb0c3 100644
--- a/src/lua/client/world.rs
+++ b/src/lua/client/world.rs
@@ -10,7 +10,7 @@ use azalea::{
     },
     world::{InstanceName, MinecraftEntityId},
 };
-use mlua::{Function, Lua, Result, Table};
+use mlua::{Function, Lua, Result, Table, UserDataRef};
 
 pub fn best_tool_for_block(lua: &Lua, client: &Client, block_state: u16) -> Result<Table> {
     let tr = client.best_tool_in_hotbar_for_block(BlockState { id: block_state });
@@ -48,35 +48,44 @@ pub fn find_blocks(
         .collect())
 }
 
-pub fn find_entities(lua: &Lua, client: &Client, filter_fn: Function) -> Result<Vec<Table>> {
-    let mut matched = Vec::new();
+pub async fn find_entities(
+    lua: Lua,
+    client: UserDataRef<Client>,
+    filter_fn: Function,
+) -> Result<Vec<Table>> {
+    let mut entities = Vec::new();
 
-    let mut ecs = client.ecs.lock();
-    let mut query = ecs.query_filtered::<(
-        &MinecraftEntityId,
-        &EntityUuid,
-        &EntityKind,
-        &CustomName,
-        &AzaleaPosition,
-        &LookDirection,
-        &Pose,
-    ), Without<Dead>>();
+    {
+        let mut ecs = client.ecs.lock();
+        let mut query = ecs.query_filtered::<(
+            &MinecraftEntityId,
+            &EntityUuid,
+            &EntityKind,
+            &CustomName,
+            &AzaleaPosition,
+            &LookDirection,
+            &Pose,
+        ), Without<Dead>>();
 
-    for (&id, uuid, kind, custom_name, position, direction, pose) in query.iter(&ecs) {
-        let entity = lua.create_table()?;
-        entity.set("id", id.0)?;
-        entity.set("uuid", uuid.to_string())?;
-        entity.set("kind", kind.to_string())?;
-        entity.set("custom_name", custom_name.as_ref().map(ToString::to_string))?;
-        entity.set("position", Vec3::from(position))?;
-        entity.set("direction", Direction::from(direction))?;
-        entity.set("pose", *pose as u8)?;
-
-        if filter_fn.call::<bool>(&entity)? {
-            matched.push(entity);
+        for (id, uuid, kind, custom_name, position, direction, pose) in query.iter(&ecs) {
+            let entity = lua.create_table()?;
+            entity.set("id", id.0)?;
+            entity.set("uuid", uuid.to_string())?;
+            entity.set("kind", kind.to_string())?;
+            entity.set("custom_name", custom_name.as_ref().map(ToString::to_string))?;
+            entity.set("position", Vec3::from(position))?;
+            entity.set("direction", Direction::from(direction))?;
+            entity.set("pose", *pose as u8)?;
+            entities.push(entity);
         }
     }
 
+    let mut matched = Vec::new();
+    for entity in entities {
+        if filter_fn.call_async::<bool>(&entity).await? {
+            matched.push(entity)
+        }
+    }
     Ok(matched)
 }