From c1602ef2cfcf83ed1cae0d1e1e919b72123bcf84 Mon Sep 17 00:00:00 2001 From: Lily Date: Fri, 26 Apr 2024 03:07:02 -0400 Subject: [PATCH 1/2] Initial dynamic query implementation --- crates/bevy_script_api/src/common/bevy/mod.rs | 151 +++++++++++++++++- crates/bevy_script_api/src/core_providers.rs | 1 + crates/bevy_script_api/src/lua/bevy/mod.rs | 106 +++++++++++- crates/bevy_script_api/src/lua/util.rs | 90 ++++++++++- crates/bevy_script_api/src/rhai/bevy/mod.rs | 93 ++++++++++- 5 files changed, 422 insertions(+), 19 deletions(-) diff --git a/crates/bevy_script_api/src/common/bevy/mod.rs b/crates/bevy_script_api/src/common/bevy/mod.rs index 40cf4ab1..bb5bb7d2 100644 --- a/crates/bevy_script_api/src/common/bevy/mod.rs +++ b/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, + with: Vec, + without: Vec, +} + +impl ScriptQueryBuilder { + pub fn new(world: ScriptWorld) -> Self { + Self { + world, + components: vec![], + with: vec![], + without: vec![], + } + } + + pub fn components(&mut self, components: Vec) -> &mut Self { + self.components.extend(components); + self + } + + pub fn with(&mut self, with: Vec) -> &mut Self { + self.with.extend(with); + self + } + + pub fn without(&mut self, without: Vec) -> &mut Self { + self.without.extend(without); + self + } + + pub fn build(&mut self) -> Result, 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); + #[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 { let w = self.read(); w.get::(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, + with: Vec, + without: Vec, + ) -> Result, ScriptError> { + let mut w = self.write(); + + let get_id = |component: &ScriptTypeRegistration, + w: &MappedRwLockWriteGuard| + -> Result { + 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::().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::, ScriptError>>()?; + + let with_ids: Vec = with + .iter() + .map(|component| get_id(component, &w)) + .collect::, ScriptError>>()?; + + let without_ids: Vec = without + .iter() + .map(|component| get_id(component, &w)) + .collect::, ScriptError>>()?; + + let mut q = QueryBuilder::::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> = 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::, ScriptError>>() + .map(|references| ScriptQueryResult(filtered_entity.id(), references)) + }) + .collect::, ScriptError>>() + } } diff --git a/crates/bevy_script_api/src/core_providers.rs b/crates/bevy_script_api/src/core_providers.rs index 0e1196ea..814a499c 100644 --- a/crates/bevy_script_api/src/core_providers.rs +++ b/crates/bevy_script_api/src/core_providers.rs @@ -56,6 +56,7 @@ impl bevy_mod_scripting_core::hosts::APIProvider for LuaCoreBevyAPIProvider { .process_type::>() .process_type::() .process_type::>() + .process_type::() }, )) } diff --git a/crates/bevy_script_api/src/lua/bevy/mod.rs b/crates/bevy_script_api/src/lua/bevy/mod.rs index 69896455..687787d6 100644 --- a/crates/bevy_script_api/src/lua/bevy/mod.rs +++ b/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::{ComponentTuple, QueryResultTuple}, + 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, LuaError> { + let mut values = LuaMultiValue::from_vec( + self.1 + .into_iter() + .map(|v| v.into_lua(lua)) + .collect::, 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: ComponentTuple| { + s.with(components.0); + Ok(s.clone()) + }); + + methods.document("Filters out entities with any components passed"); + methods.add_method_mut("without", |_, s, components: ComponentTuple| { + 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(); + QueryResultTuple::Some( + LuaEntity::new(query_result.0), + query_result.1.clone(), + ) + } else { + QueryResultTuple::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: ComponentTuple| { + 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( diff --git a/crates/bevy_script_api/src/lua/util.rs b/crates/bevy_script_api/src/lua/util.rs index 09ecadb6..0f90d8ab 100644 --- a/crates/bevy_script_api/src/lua/util.rs +++ b/crates/bevy_script_api/src/lua/util.rs @@ -1,6 +1,9 @@ +use crate::{lua::bevy::LuaTypeRegistration, providers::bevy_ecs::LuaEntity, ReflectReference}; use bevy_mod_scripting_lua::{ - prelude::{FromLua, IntoLua, Lua, LuaError, LuaValue}, - tealr::{self, ToTypename}, + prelude::{ + FromLua, FromLuaMulti, IntoLua, IntoLuaMulti, Lua, LuaError, LuaMultiValue, LuaValue, + }, + tealr::{self, FunctionParam, KindOfType, Name, SingleType, TealMultiValue, ToTypename, Type}, }; use std::{ marker::PhantomData, @@ -86,6 +89,89 @@ impl ToTypename for DummyTypeName { } } +/// A utility type that allows us to accept any number of components as a parameter into a function. +#[derive(Clone)] +pub struct ComponentTuple(pub Vec); + +impl IntoLuaMulti<'_> for ComponentTuple { + fn into_lua_multi(self, lua: &Lua) -> Result, LuaError> { + let values = LuaMultiValue::from_vec( + self.0 + .into_iter() + .map(|v| v.into_lua(lua).unwrap()) + .collect(), + ); + + Ok(values) + } +} + +impl FromLuaMulti<'_> for ComponentTuple { + fn from_lua_multi(value: LuaMultiValue<'_>, lua: &Lua) -> Result { + Ok(ComponentTuple( + value + .into_vec() + .into_iter() + .map(|v| LuaTypeRegistration::from_lua(v, lua).unwrap()) + .collect(), + )) + } +} + +impl TealMultiValue for ComponentTuple { + fn get_types_as_params() -> Vec { + vec![FunctionParam { + // `...:T` will be a variadic parameter + param_name: Some(Name("...".into())), + ty: LuaTypeRegistration::to_typename(), + }] + } +} + +/// A utility enum that allows us to return an entity and any number of components from a function. +#[derive(Clone)] +pub enum QueryResultTuple { + Some(LuaEntity, Vec), + None, +} + +impl IntoLuaMulti<'_> for QueryResultTuple { + fn into_lua_multi(self, lua: &Lua) -> Result, LuaError> { + match self { + QueryResultTuple::Some(entity, vec) => { + let mut values = LuaMultiValue::from_vec( + vec.into_iter() + .map(|v| v.into_lua(lua)) + .collect::, LuaError>>()?, + ); + + values.push_front(entity.into_lua(lua)?); + Ok(values) + } + QueryResultTuple::None => Ok(().into_lua_multi(lua)?), + } + } +} + +impl TealMultiValue for QueryResultTuple { + fn get_types_as_params() -> Vec { + vec![ + FunctionParam { + param_name: None, + ty: LuaEntity::to_typename(), + }, + FunctionParam { + param_name: None, + ty: Type::Single(SingleType { + kind: KindOfType::External, + // tealr doesn't have a way to add variadic return types, so it's inserted into the type name instead + name: Name(format!("{}...", stringify!(ReflectedValue)).into()), + }), + }, + ] + } +} + #[macro_export] macro_rules! impl_from_lua_with_clone { ($v:ty) => { diff --git a/crates/bevy_script_api/src/rhai/bevy/mod.rs b/crates/bevy_script_api/src/rhai/bevy/mod.rs index 4ab540e1..3f4ea2a8 100644 --- a/crates/bevy_script_api/src/rhai/bevy/mod.rs +++ b/crates/bevy_script_api/src/rhai/bevy/mod.rs @@ -9,7 +9,7 @@ use bevy_mod_scripting_rhai::{ use rhai::plugin::*; use crate::{ - common::bevy::{ScriptTypeRegistration, ScriptWorld}, + common::bevy::{ScriptQueryBuilder, ScriptTypeRegistration, ScriptWorld}, ReflectedValue, }; @@ -29,6 +29,69 @@ impl CustomType for ScriptTypeRegistration { } } +impl CustomType for ScriptQueryBuilder { + fn build(mut builder: rhai::TypeBuilder) { + builder + .with_name("QueryBuilder") + // `with` is a reserved keyword, so we add _components on the end + .with_fn("with_components", |self_: &mut Self, with: Vec| { + self_.with( + with.into_iter() + .map(Dynamic::cast::) + .collect(), + ); + + Dynamic::from(self_.clone()) + }) + .with_fn( + "without_components", + |self_: &mut Self, without: Vec| { + self_.without( + without + .into_iter() + .map(Dynamic::cast::) + .collect(), + ); + + Dynamic::from(self_.clone()) + }, + ); + } +} + +impl IntoIterator for ScriptQueryBuilder { + type Item = rhai::Map; + type IntoIter = std::vec::IntoIter; + + fn into_iter(mut self) -> Self::IntoIter { + self.build() + .expect("Query failed!") + .into_iter() + .map(|result| { + let mut map = rhai::Map::new(); + map.insert("Entity".into(), Dynamic::from(result.0)); + + for component in result.1.into_iter() { + let name = component + .get(|value| value.get_represented_type_info()?.type_path_table().ident()) + .unwrap() + .unwrap(); + + map.insert( + name.into(), + component.to_dynamic().unwrap_or_else(|_| { + panic!("Converting component {} to dynamic failed!", &name) + }), + ); + } + + map + }) + .collect::>() + .into_iter() + } +} + #[allow(deprecated)] impl CustomType for ScriptWorld { fn build(mut builder: rhai::TypeBuilder) { @@ -71,7 +134,7 @@ impl CustomType for ScriptWorld { }, ) .with_fn( - "has_compoennt", + "has_component", |self_: ScriptWorld, entity: Entity, comp_type: ScriptTypeRegistration| { self_.has_component(entity, comp_type).map_err(|e| { Box::new(EvalAltResult::ErrorRuntime( @@ -210,7 +273,28 @@ impl CustomType for ScriptWorld { w.despawn(entity) }) .with_fn("to_string", |self_: &mut ScriptWorld| self_.to_string()) - .with_fn("to_debug", |self_: &mut ScriptWorld| format!("{:?}", self_)); + .with_fn("to_debug", |self_: &mut ScriptWorld| format!("{:?}", self_)) + .with_fn( + "query", + |self_: &mut ScriptWorld, component: ScriptTypeRegistration| { + ScriptQueryBuilder::new(self_.clone()) + .components(vec![component]) + .clone() + }, + ) + .with_fn( + "query", + |self_: &mut ScriptWorld, components: Vec| { + ScriptQueryBuilder::new(self_.clone()) + .components( + components + .into_iter() + .map(Dynamic::cast::) + .collect::>(), + ) + .clone() + }, + ); } } @@ -225,6 +309,9 @@ impl APIProvider for RhaiBevyAPIProvider { engine.build_type::(); engine.build_type::(); engine.build_type::(); + engine.build_type::(); + engine.register_iterator::>(); + engine.register_iterator::(); Ok(()) } From 1c1a55421ce9c070bf300ec00d94dc9c5f7e9653 Mon Sep 17 00:00:00 2001 From: Lily Date: Fri, 26 Apr 2024 03:34:35 -0400 Subject: [PATCH 2/2] Rename Tuple -> Variadic --- crates/bevy_script_api/src/lua/bevy/mod.rs | 12 +++++------ crates/bevy_script_api/src/lua/util.rs | 24 +++++++++++----------- 2 files changed, 18 insertions(+), 18 deletions(-) diff --git a/crates/bevy_script_api/src/lua/bevy/mod.rs b/crates/bevy_script_api/src/lua/bevy/mod.rs index 687787d6..8b2068b2 100644 --- a/crates/bevy_script_api/src/lua/bevy/mod.rs +++ b/crates/bevy_script_api/src/lua/bevy/mod.rs @@ -4,7 +4,7 @@ use crate::common::bevy::{ use crate::lua::{ mlua::prelude::{IntoLuaMulti, LuaError, LuaMultiValue}, tealr::{mlu::TypedFunction, ToTypename}, - util::{ComponentTuple, QueryResultTuple}, + util::{VariadicComponents, VariadicQueryResult}, Lua, }; use crate::providers::bevy_ecs::LuaEntity; @@ -105,13 +105,13 @@ impl TealData for LuaQueryBuilder { 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: ComponentTuple| { + 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: ComponentTuple| { + methods.add_method_mut("without", |_, s, components: VariadicComponents| { s.without(components.0); Ok(s.clone()) }); @@ -129,12 +129,12 @@ impl TealData for LuaQueryBuilder { move |_, ()| { let o = if curr_idx < len { let query_result = query_result.get(curr_idx).unwrap(); - QueryResultTuple::Some( + VariadicQueryResult::Some( LuaEntity::new(query_result.0), query_result.1.clone(), ) } else { - QueryResultTuple::None + VariadicQueryResult::None }; curr_idx += 1; Ok(o) @@ -202,7 +202,7 @@ 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: ComponentTuple| { + methods.add_method_mut("query", |_, world, components: VariadicComponents| { Ok(LuaQueryBuilder::new(world.clone()) .components(components.0) .clone()) diff --git a/crates/bevy_script_api/src/lua/util.rs b/crates/bevy_script_api/src/lua/util.rs index 0f90d8ab..1221e2e9 100644 --- a/crates/bevy_script_api/src/lua/util.rs +++ b/crates/bevy_script_api/src/lua/util.rs @@ -91,9 +91,9 @@ impl ToTypename for DummyTypeName { /// A utility type that allows us to accept any number of components as a parameter into a function. #[derive(Clone)] -pub struct ComponentTuple(pub Vec); +pub struct VariadicComponents(pub Vec); -impl IntoLuaMulti<'_> for ComponentTuple { +impl IntoLuaMulti<'_> for VariadicComponents { fn into_lua_multi(self, lua: &Lua) -> Result, LuaError> { let values = LuaMultiValue::from_vec( self.0 @@ -106,9 +106,9 @@ impl IntoLuaMulti<'_> for ComponentTuple { } } -impl FromLuaMulti<'_> for ComponentTuple { - fn from_lua_multi(value: LuaMultiValue<'_>, lua: &Lua) -> Result { - Ok(ComponentTuple( +impl FromLuaMulti<'_> for VariadicComponents { + fn from_lua_multi(value: LuaMultiValue<'_>, lua: &Lua) -> Result { + Ok(VariadicComponents( value .into_vec() .into_iter() @@ -118,10 +118,10 @@ impl FromLuaMulti<'_> for ComponentTuple { } } -impl TealMultiValue for ComponentTuple { +impl TealMultiValue for VariadicComponents { fn get_types_as_params() -> Vec { vec![FunctionParam { - // `...:T` will be a variadic parameter + // `...:T` will be a variadic type param_name: Some(Name("...".into())), ty: LuaTypeRegistration::to_typename(), }] @@ -130,15 +130,15 @@ impl TealMultiValue for ComponentTuple { /// A utility enum that allows us to return an entity and any number of components from a function. #[derive(Clone)] -pub enum QueryResultTuple { +pub enum VariadicQueryResult { Some(LuaEntity, Vec), None, } -impl IntoLuaMulti<'_> for QueryResultTuple { +impl IntoLuaMulti<'_> for VariadicQueryResult { fn into_lua_multi(self, lua: &Lua) -> Result, LuaError> { match self { - QueryResultTuple::Some(entity, vec) => { + VariadicQueryResult::Some(entity, vec) => { let mut values = LuaMultiValue::from_vec( vec.into_iter() .map(|v| v.into_lua(lua)) @@ -148,12 +148,12 @@ impl IntoLuaMulti<'_> for QueryResultTuple { values.push_front(entity.into_lua(lua)?); Ok(values) } - QueryResultTuple::None => Ok(().into_lua_multi(lua)?), + VariadicQueryResult::None => Ok(().into_lua_multi(lua)?), } } } -impl TealMultiValue for QueryResultTuple { +impl TealMultiValue for VariadicQueryResult { fn get_types_as_params() -> Vec { vec![ FunctionParam {