re-implement dynamic queries

Co-authored-by: Lily <mistrustfully@gmail.com>

Author: makspll

crates/bevy_mod_scripting_core/src/bindings/query.rs

@@ -1,10 +1,14 @@
-use super::{ReflectReference, WorldCallbackAccess};
-use crate::prelude::ScriptResult;
+use super::{ReflectReference, WorldAccessGuard, WorldCallbackAccess};
+use crate::{
+    bindings::{CONCURRENT_WORLD_ACCESS_MSG, STALE_WORLD_MSG},
+    prelude::ScriptResult,
+};
 use bevy::{
     ecs::{component::ComponentId, entity::Entity},
+    prelude::{EntityRef, QueryBuilder},
     reflect::TypeRegistration,
 };
-use std::{any::TypeId, ops::Deref, sync::Arc};
+use std::{any::TypeId, collections::VecDeque, sync::Arc};
 
 /// A wrapper around a `TypeRegistration` that provides additional information about the type.
 ///
@@ -65,24 +69,14 @@ impl std::fmt::Display for ScriptTypeRegistration {
     }
 }
 
-#[derive(Clone)]
+#[derive(Clone, Default)]
 pub struct ScriptQueryBuilder {
-    world: WorldCallbackAccess,
     components: Vec<ScriptTypeRegistration>,
     with: Vec<ScriptTypeRegistration>,
     without: Vec<ScriptTypeRegistration>,
 }
 
 impl ScriptQueryBuilder {
-    pub fn new(world: WorldCallbackAccess) -> Self {
-        Self {
-            world,
-            components: vec![],
-            with: vec![],
-            without: vec![],
-        }
-    }
-
     pub fn components(&mut self, components: Vec<ScriptTypeRegistration>) -> &mut Self {
         self.components.extend(components);
         self
@@ -97,29 +91,76 @@ impl ScriptQueryBuilder {
         self.without.extend(without);
         self
     }
-
-    pub fn build(&mut self) -> ScriptResult<Vec<ScriptQueryResult>> {
-        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>);
 
 impl WorldCallbackAccess {
-    pub fn query(
-        &mut self,
-        components: Vec<ScriptTypeRegistration>,
-        with: Vec<ScriptTypeRegistration>,
-        without: Vec<ScriptTypeRegistration>,
-    ) -> ScriptResult<Vec<ScriptQueryResult>> {
-        // for c in components {
-
-        // }
-        todo!()
+    pub fn query(&self, query: ScriptQueryBuilder) -> ScriptResult<VecDeque<ScriptQueryResult>> {
+        // find the set of components
+        self.read()
+            .unwrap_or_else(|| panic!("{STALE_WORLD_MSG}"))
+            .query(query)
     }
 }
+
+impl<'w> WorldAccessGuard<'w> {
+    pub fn query(&self, query: ScriptQueryBuilder) -> ScriptResult<VecDeque<ScriptQueryResult>> {
+        let world = self
+            .get_whole_world_access()
+            .unwrap_or_else(|| panic!("{CONCURRENT_WORLD_ACCESS_MSG}"));
+
+        let mut dynamic_query = QueryBuilder::<EntityRef>::new(world);
+
+        // we don't actually want to fetch the data for components now, only figure out
+        // which entities match the query
+        // so we might be being slightly overkill
+        for c in &query.components {
+            dynamic_query.ref_id(c.component_id().unwrap());
+        }
+
+        for w in query.with {
+            dynamic_query.with_id(w.component_id.unwrap());
+        }
+
+        for without_id in query.without {
+            dynamic_query.without_id(without_id.component_id.unwrap());
+        }
+
+        let mut built_query = dynamic_query.build();
+        let query_result = built_query.iter(world);
+
+        Ok(query_result
+            .map(|r| {
+                let references: Vec<_> = query
+                    .components
+                    .iter()
+                    .map(|c| ReflectReference {
+                        base: super::ReflectBaseType {
+                            type_id: c.type_id(),
+                            base_id: super::ReflectBase::Component(r.id(), c.component_id.unwrap()),
+                        },
+                        reflect_path: vec![],
+                    })
+                    .collect();
+                ScriptQueryResult(r.id(), references)
+            })
+            .collect())
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use test_utils::test_data::setup_world;
+
+    use super::*;
+
+    // #[test]
+    // fn test_simple_query() {
+    //     let world = setup_world(|w,r|{
+    //         w.spawn(TestComponent::init())
+    //         w.spawn(Te)
+    //     })
+    // }
+}

crates/bevy_mod_scripting_core/src/bindings/world.rs

@@ -245,11 +245,11 @@ pub const DEFAULT_INTERVAL: Duration = Duration::from_millis(10);
 /// Provides safe access to the world via [`WorldAccess`] permissions, which enforce aliasing rules at runtime in multi-thread environments
 #[derive(Clone)]
 pub struct WorldAccessGuard<'w> {
-    cell: UnsafeWorldCell<'w>,
+    pub cell: UnsafeWorldCell<'w>,
     // TODO: this is fairly hefty, explore other ways to hand out locks on WorldAccess
-    accesses: Arc<LockableHashMap<ReflectAccessId, Option<WorldAccessUnit<'w>>>>,
+    pub accesses: Arc<LockableHashMap<ReflectAccessId, Option<WorldAccessUnit<'w>>>>,
     /// true if anybody has any access to the world
-    accesses_count: Arc<AtomicUsize>,
+    pub accesses_count: Arc<AtomicUsize>,
     // TODO can we track code/stack locations of things holding onto theese locks for debugging?
 }
 

crates/languages/bevy_mod_scripting_lua/src/bindings/mod.rs

@@ -1,7 +1,6 @@
-use bevy_mod_scripting_core::bindings::WorldCallbackAccess;
-
 pub mod providers;
 pub mod proxy;
+pub mod query;
 pub mod reference;
 pub mod std;
 pub mod type_registration;

crates/languages/bevy_mod_scripting_lua/src/bindings/query.rs

@@ -0,0 +1,122 @@
+use std::ops::{Deref, DerefMut};
+
+use bevy::prelude::Entity;
+use bevy_mod_scripting_core::bindings::{
+    ReflectAllocator, ReflectReference, ScriptQueryBuilder, ScriptTypeRegistration, Unproxy,
+    ValProxy,
+};
+use tealr::{
+    mlu::{
+        mlua::{IntoLua, Value},
+        TealData, TypedFunction,
+    },
+    ToTypename,
+};
+
+use crate::{impl_userdata_from_lua, impl_userdata_with_tealdata, util::Variadic};
+
+use super::{
+    proxy::{LuaProxied, LuaValProxy},
+    reference::LuaReflectReference,
+    world::GetWorld,
+};
+
+#[derive(Default, Clone)]
+pub struct LuaQueryBuilder(pub(crate) ScriptQueryBuilder);
+
+impl Deref for LuaQueryBuilder {
+    type Target = ScriptQueryBuilder;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for LuaQueryBuilder {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl LuaQueryBuilder {}
+
+impl LuaProxied for ScriptQueryBuilder {
+    type Proxy = LuaQueryBuilder;
+}
+
+impl_userdata_from_lua!(LuaQueryBuilder);
+impl_userdata_with_tealdata!(LuaQueryBuilder);
+
+impl TealData for LuaQueryBuilder {
+    fn add_methods<'lua, T: tealr::mlu::TealDataMethods<'lua, Self>>(methods: &mut T) {
+        methods.add_function(
+            "with",
+            |_, (mut this, mut query): (Self, Variadic<LuaValProxy<ScriptTypeRegistration>>)| {
+                let registrations: Vec<_> = query
+                    .iter_mut()
+                    .map(|v| v.unproxy())
+                    .collect::<Result<_, _>>()
+                    .map_err(tealr::mlu::mlua::Error::external)?;
+                this.0.with(registrations);
+                Ok(this)
+            },
+        );
+
+        methods.add_function(
+            "without",
+            |_, (mut this, mut query): (Self, Variadic<LuaValProxy<ScriptTypeRegistration>>)| {
+                let registrations: Vec<_> = query
+                    .iter_mut()
+                    .map(|v| v.unproxy())
+                    .collect::<Result<_, _>>()
+                    .map_err(tealr::mlu::mlua::Error::external)?;
+                this.0.without(registrations);
+                Ok(this)
+            },
+        );
+
+        methods.add_function("iter", |l, this: LuaQueryBuilder| {
+            let world = l.get_world()?;
+            let mut result = world
+                .query(this.0)
+                .map_err(tealr::mlu::mlua::Error::external)?;
+            let mut len = result.len();
+            let iterator = TypedFunction::from_rust_mut(
+                move |lua, ()| {
+                    if len > 0 {
+                        let result = result
+                            .pop_front()
+                            .expect("invariant: len of array = len && len > 0");
+                        let entity =
+                            world.with_resource::<ReflectAllocator, _, _>(|_, mut allocator| {
+                                ReflectReference::new_allocated(result.0, &mut allocator)
+                            });
+                        let proxy_entity =
+                            <Entity as LuaProxied>::Proxy::from(entity).into_lua(lua)?;
+                        let component_refs: Vec<_> = result
+                            .1
+                            .into_iter()
+                            .map(|r| LuaReflectReference(r).to_lua_proxy(lua))
+                            .collect::<Result<_, _>>()?;
+
+                        len -= 1;
+
+                        Ok(Variadic::new(
+                            std::iter::once(proxy_entity).chain(component_refs.into_iter()),
+                        ))
+                    } else {
+                        Ok(Variadic::new(vec![Value::Nil, Value::Nil]))
+                    }
+                },
+                l,
+            )?;
+            Ok(iterator)
+        });
+    }
+}
+
+impl ToTypename for LuaQueryBuilder {
+    fn to_typename() -> tealr::Type {
+        tealr::Type::new_single("QueryBuilder", tealr::KindOfType::External)
+    }
+}

crates/languages/bevy_mod_scripting_lua/src/bindings/world.rs

@@ -20,13 +20,15 @@ use tealr::{
 };
 
 use super::proxy::LuaReflectRefProxy;
+use super::query::LuaQueryBuilder;
 use super::{
     providers::bevy_ecs::LuaEntity,
     proxy::{
         ErrorProxy, LuaIdentityProxy, LuaProxied, LuaReflectValProxy, LuaValProxy, TypenameProxy,
     },
     type_registration::LuaTypeRegistration,
 };
+use crate::util::Variadic;
 use crate::{impl_userdata_from_lua, impl_userdata_with_tealdata};
 
 pub struct Nil;
@@ -375,6 +377,25 @@ impl TealData for LuaWorld {
                 Ok(TypenameProxy::<_, Nil>::new(out))
             },
         );
+
+        methods.add_method(
+            "query",
+            |_, this, mut components: Variadic<LuaValProxy<ScriptTypeRegistration>>| {
+                let world = this.0.read().ok_or_else(|| {
+                    mlua::Error::external(ScriptError::new_reflection_error("Stale world access"))
+                })?;
+                let mut builder = LuaQueryBuilder::default();
+                let deque = components.0;
+                builder.components(
+                    deque
+                        .into_iter()
+                        .map(|mut c| c.unproxy())
+                        .collect::<Result<_, _>>()
+                        .map_err(tealr::mlu::mlua::Error::external)?,
+                );
+                Ok(builder)
+            },
+        );
     }
 
     fn add_fields<'lua, F: tealr::mlu::TealDataFields<'lua, Self>>(_fields: &mut F) {}