Dynamic Queries

crates/bevy_script_api/src/common/bevy/mod.rs

@@ -1,12 +1,12 @@
-use std::{
-    ops::{Deref, DerefMut},
-    sync::Arc,
-};
-
 use crate::ReflectReference;
 /// Common functionality for all script hosts
 use bevy::{
-    ecs::system::Command,
+    ecs::{
+        component::ComponentId,
+        query::QueryBuilder,
+        system::Command,
+        world::{EntityRef, World},
+    },
     prelude::{
         AppTypeRegistry, BuildWorldChildren, Children, DespawnChildrenRecursive, DespawnRecursive,
         Entity, Parent, ReflectComponent, ReflectDefault, ReflectResource,
@@ -17,6 +17,12 @@ use bevy::{
     },
 };
 use bevy_mod_scripting_core::{prelude::ScriptError, world::WorldPointer};
+use parking_lot::MappedRwLockWriteGuard;
+use std::{
+    any::Any,
+    ops::{Deref, DerefMut},
+    sync::Arc,
+};
 
 /// Helper trait for retrieving a world pointer from a script context.
 pub trait GetWorld {
@@ -65,6 +71,51 @@ impl Deref for ScriptTypeRegistration {
     }
 }
 
+#[derive(Clone)]
+pub struct ScriptQueryBuilder {
+    world: ScriptWorld,
+    components: Vec<ScriptTypeRegistration>,
+    with: Vec<ScriptTypeRegistration>,
+    without: Vec<ScriptTypeRegistration>,
+}
+
+impl ScriptQueryBuilder {
+    pub fn new(world: ScriptWorld) -> Self {
+        Self {
+            world,
+            components: vec![],
+            with: vec![],
+            without: vec![],
+        }
+    }
+
+    pub fn components(&mut self, components: Vec<ScriptTypeRegistration>) -> &mut Self {
+        self.components.extend(components);
+        self
+    }
+
+    pub fn with(&mut self, with: Vec<ScriptTypeRegistration>) -> &mut Self {
+        self.with.extend(with);
+        self
+    }
+
+    pub fn without(&mut self, without: Vec<ScriptTypeRegistration>) -> &mut Self {
+        self.without.extend(without);
+        self
+    }
+
+    pub fn build(&mut self) -> Result<Vec<ScriptQueryResult>, ScriptError> {
+        self.world.query(
+            std::mem::take(&mut self.components),
+            std::mem::take(&mut self.with),
+            std::mem::take(&mut self.without),
+        )
+    }
+}
+
+#[derive(Clone)]
+pub struct ScriptQueryResult(pub Entity, pub Vec<ReflectReference>);
+
 #[derive(Clone, Debug)]
 pub struct ScriptWorld(WorldPointer);
 
@@ -104,6 +155,7 @@ impl ScriptWorld {
     pub fn new(ptr: WorldPointer) -> Self {
         Self(ptr)
     }
+
     pub fn get_children(&self, parent: Entity) -> Vec<Entity> {
         let w = self.read();
         w.get::<Children>(parent)
@@ -292,6 +344,7 @@ impl ScriptWorld {
 
         Ok(resource_data.reflect(&w).is_some())
     }
+
     pub fn remove_resource(&mut self, res_type: ScriptTypeRegistration) -> Result<(), ScriptError> {
         let mut w = self.write();
 
@@ -301,4 +354,90 @@ impl ScriptWorld {
         resource_data.remove(&mut w);
         Ok(())
     }
+
+    pub fn query(
+        &mut self,
+        components: Vec<ScriptTypeRegistration>,
+        with: Vec<ScriptTypeRegistration>,
+        without: Vec<ScriptTypeRegistration>,
+    ) -> Result<Vec<ScriptQueryResult>, ScriptError> {
+        let mut w = self.write();
+
+        let get_id = |component: &ScriptTypeRegistration,
+                      w: &MappedRwLockWriteGuard<World>|
+         -> Result<ComponentId, ScriptError> {
+            w.components()
+                .get_id(component.type_info().type_id())
+                .ok_or_else(|| {
+                    ScriptError::Other(format!("Not a component {}", component.short_name()))
+                })
+        };
+
+        let components: Vec<(ReflectComponent, ComponentId)> = components
+            .into_iter()
+            .map(|component| {
+                let reflect_component = component.data::<ReflectComponent>().ok_or_else(|| {
+                    ScriptError::Other(format!("Not a component {}", component.short_name()))
+                });
+
+                let component_id = get_id(&component, &w);
+                reflect_component.map(|v1| component_id.map(|v2| (v1.clone(), v2)))?
+            })
+            .collect::<Result<Vec<_>, ScriptError>>()?;
+
+        let with_ids: Vec<ComponentId> = with
+            .iter()
+            .map(|component| get_id(component, &w))
+            .collect::<Result<Vec<_>, ScriptError>>()?;
+
+        let without_ids: Vec<ComponentId> = without
+            .iter()
+            .map(|component| get_id(component, &w))
+            .collect::<Result<Vec<_>, ScriptError>>()?;
+
+        let mut q = QueryBuilder::<EntityRef>::new(&mut w);
+
+        for (_, id) in &components {
+            q.ref_id(*id);
+        }
+
+        for with_id in with_ids {
+            q.with_id(with_id);
+        }
+
+        for without_id in without_ids {
+            q.without_id(without_id);
+        }
+
+        let query_result: Vec<EntityRef<'_>> = q.build().iter_mut(&mut w).collect();
+
+        query_result
+            .into_iter()
+            .map(|filtered_entity| {
+                components
+                    .clone()
+                    .into_iter()
+                    .map(|(reflect_component, _)| {
+                        let type_id = reflect_component.type_id();
+                        reflect_component
+                            .reflect(filtered_entity)
+                            .map(|_component| {
+                                ReflectReference::new_component_ref(
+                                    reflect_component,
+                                    filtered_entity.id(),
+                                    self.clone().into(),
+                                )
+                            })
+                            .ok_or_else(|| {
+                                ScriptError::Other(format!(
+                                    "Failed to reflect component during query: {:?}",
+                                    type_id
+                                ))
+                            })
+                    })
+                    .collect::<Result<Vec<_>, ScriptError>>()
+                    .map(|references| ScriptQueryResult(filtered_entity.id(), references))
+            })
+            .collect::<Result<Vec<_>, ScriptError>>()
+    }
 }

crates/bevy_script_api/src/core_providers.rs

@@ -56,6 +56,7 @@ impl bevy_mod_scripting_core::hosts::APIProvider for LuaCoreBevyAPIProvider {
 			.process_type::<bevy_mod_scripting_lua::tealr::mlu::UserDataProxy<crate::lua::bevy::LuaScriptData>>()
 			.process_type::<crate::lua::bevy::LuaTypeRegistration>()
 			.process_type::<crate::lua::std::LuaVec<T>>()
+            .process_type::<crate::lua::bevy::LuaQueryBuilder>()
             },
         ))
     }

crates/bevy_script_api/src/lua/bevy/mod.rs

@@ -1,15 +1,19 @@
-use crate::common::bevy::{ScriptTypeRegistration, ScriptWorld};
+use crate::common::bevy::{
+    ScriptQueryBuilder, ScriptQueryResult, ScriptTypeRegistration, ScriptWorld,
+};
+use crate::lua::{
+    mlua::prelude::{IntoLuaMulti, LuaError, LuaMultiValue},
+    tealr::{mlu::TypedFunction, ToTypename},
+    util::{VariadicComponents, VariadicQueryResult},
+    Lua,
+};
 use crate::providers::bevy_ecs::LuaEntity;
 use crate::{impl_from_lua_with_clone, impl_tealr_type};
-
-use std::sync::Arc;
-
 use bevy::hierarchy::BuildWorldChildren;
-use bevy::prelude::AppTypeRegistry;
-
-use bevy::prelude::ReflectResource;
+use bevy::prelude::{AppTypeRegistry, ReflectResource};
 use bevy_mod_scripting_core::prelude::*;
-use bevy_mod_scripting_lua::tealr;
+use bevy_mod_scripting_lua::{prelude::IntoLua, tealr};
+use std::sync::Arc;
 
 use tealr::mlu::{
     mlua::{self},
@@ -63,6 +67,84 @@ impl TealData for LuaScriptData {
     }
 }
 
+pub type LuaQueryResult = ScriptQueryResult;
+
+impl_from_lua_with_clone!(LuaQueryResult);
+
+impl IntoLuaMulti<'_> for LuaQueryResult {
+    fn into_lua_multi(self, lua: &Lua) -> Result<LuaMultiValue<'_>, LuaError> {
+        let mut values = LuaMultiValue::from_vec(
+            self.1
+                .into_iter()
+                .map(|v| v.into_lua(lua))
+                .collect::<Result<Vec<_>, LuaError>>()?,
+        );
+        values.push_front(LuaEntity::new(self.0).into_lua(lua)?);
+        Ok(values)
+    }
+}
+
+impl ToTypename for LuaQueryResult {
+    fn to_typename() -> bevy_mod_scripting_lua::tealr::Type {
+        bevy_mod_scripting_lua::tealr::Type::new_single(
+            stringify!(QueryResult),
+            bevy_mod_scripting_lua::tealr::KindOfType::External,
+        )
+    }
+}
+
+pub type LuaQueryBuilder = ScriptQueryBuilder;
+
+impl_tealr_type!(LuaQueryBuilder);
+impl_from_lua_with_clone!(LuaQueryBuilder);
+
+impl TealData for LuaQueryBuilder {
+    fn add_fields<'lua, F: tealr::mlu::TealDataFields<'lua, Self>>(fields: &mut F) {
+        fields.document("A Builder object which allows for filtering and iterating over components and entities in the world.");
+    }
+
+    fn add_methods<'lua, T: TealDataMethods<'lua, Self>>(methods: &mut T) {
+        methods.document("Filters out entities without any of the components passed");
+        methods.add_method_mut("with", |_, s, components: VariadicComponents| {
+            s.with(components.0);
+            Ok(s.clone())
+        });
+
+        methods.document("Filters out entities with any components passed");
+        methods.add_method_mut("without", |_, s, components: VariadicComponents| {
+            s.without(components.0);
+            Ok(s.clone())
+        });
+
+        methods
+            .document("Queries the world and returns an iterator over the entity and components.");
+        methods.add_method_mut("iter", |ctx, s, _: ()| {
+            let query_result = s
+                .build()
+                .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
+
+            let len = query_result.len();
+            let mut curr_idx = 0;
+            TypedFunction::from_rust_mut(
+                move |_, ()| {
+                    let o = if curr_idx < len {
+                        let query_result = query_result.get(curr_idx).unwrap();
+                        VariadicQueryResult::Some(
+                            LuaEntity::new(query_result.0),
+                            query_result.1.clone(),
+                        )
+                    } else {
+                        VariadicQueryResult::None
+                    };
+                    curr_idx += 1;
+                    Ok(o)
+                },
+                ctx,
+            )
+        });
+    }
+}
+
 pub type LuaWorld = ScriptWorld;
 
 impl_tealr_type!(LuaWorld);
@@ -118,6 +200,14 @@ impl TealData for LuaWorld {
             },
         );
 
+        methods.document("Creates a LuaQueryBuilder, querying for the passed components types.");
+        methods.document("Can be iterated over using `LuaQueryBuilder:iter()`");
+        methods.add_method_mut("query", |_, world, components: VariadicComponents| {
+            Ok(LuaQueryBuilder::new(world.clone())
+                .components(components.0)
+                .clone())
+        });
+
         methods
             .document("Returns `true` if the given entity contains a component of the given type.");
         methods.add_method(