Skip to content

Commit

Permalink
Merge pull request #23 from ten3roberts/relations_like_mut
Browse files Browse the repository at this point in the history
Add relations_like_mut
  • Loading branch information
ten3roberts authored May 1, 2024
2 parents 606a9c5 + e9d1bdc commit 2e2aac6
Show file tree
Hide file tree
Showing 8 changed files with 260 additions and 18 deletions.
16 changes: 8 additions & 8 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -18,18 +18,18 @@ members = ["flax-derive", "asteroids"]
[dependencies]
flax-derive = { path = "./flax-derive", version = "0.7.0", optional = true }

flume = { version = "0.11.0", default_features = false, optional = true }
flume = { version = "0.11.0", default-features = false, optional = true }


atomic_refcell = { version = "0.1.13", default_features = false }
bitflags = { version = "2.4.1", default_features = false }
anyhow = { version = "1.0.75", default_features = false }
atomic_refcell = { version = "0.1.13", default-features = false }
bitflags = { version = "2.4.1", default-features = false }
anyhow = { version = "1.0.75", default-features = false }
itertools = "0.12.1"
rayon = { version = "1.8.0", default_features = false, optional = true }
tokio = { version = "1.33.0", default_features = false, features = [
rayon = { version = "1.8.0", default-features = false, optional = true }
tokio = { version = "1.33.0", default-features = false, features = [
"sync",
], optional = true }
smallvec = { version = "1.11.1", default_features = false }
smallvec = { version = "1.11.1", default-features = false }
tracing = { version = "0.1.40", optional = true }
tynm = "0.1.9"
serde = { version = "1.0.190", features = ["derive"], optional = true }
Expand All @@ -40,7 +40,7 @@ puffin = { version = "0.19", optional = true }
[dev-dependencies]
tokio = { version = "1.33.0", features = ["test-util", "macros"] }
futures = "0.3.29"
itertools = { version = "0.12.1", default_features = false }
itertools = { version = "0.12.1", default-features = false }
tracing-subscriber = { version = "0.3.17", features = ["env-filter"] }
serde_json = "1.0.107"
glam = { version = "0.27.0", features = ["rand", "serde", "scalar-math"] }
Expand Down
4 changes: 2 additions & 2 deletions asteroids/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,8 @@ flax = { path = "../", default-features = false, features = [ "flume", "derive",
macroquad = "0.4.4"
# sapp-wasm = "=0.1.26"

flume = { version = "0.11.0", default_features = false }
rand = { version = "0.8.5", default_features = false, features = ["std_rng"] }
flume = { version = "0.11.0", default-features = false }
rand = { version = "0.8.5", default-features = false, features = ["std_rng"] }
tracing = "0.1.40"
tracing-subscriber = "0.3.17"
tracing-tree = "0.2.5"
Expand Down
17 changes: 17 additions & 0 deletions src/events.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,23 @@ pub struct Event {
pub kind: EventKind,
}

impl Event {
/// Construct a new event
pub fn new(id: Entity, key: ComponentKey, kind: EventKind) -> Self {
Self { id, key, kind }
}

/// Construct a modified event
pub fn modified(id: Entity, key: ComponentKey) -> Self {
Self::new(id, key, EventKind::Modified)
}

/// Construct an added event
pub fn added(id: Entity, key: ComponentKey) -> Self {
Self::new(id, key, EventKind::Added)
}
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
/// The type of ECS event
pub enum EventKind {
Expand Down
2 changes: 1 addition & 1 deletion src/fetch/component_mut.rs
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ pub struct WriteComponent<'a, T> {
impl<'w, 'q, T: 'q + ComponentValue> PreparedFetch<'q> for WriteComponent<'w, T> {
type Item = &'q mut T;
type Chunk = PtrMut<'q, T>;

const HAS_FILTER: bool = false;

unsafe fn create_chunk(&'q mut self, slots: Slice) -> Self::Chunk {
Expand Down
2 changes: 2 additions & 0 deletions src/fetch/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ mod maybe_mut;
mod opt;
mod read_only;
mod relations;
mod relations_mut;
mod satisfied;
mod source;
mod transform;
Expand All @@ -36,6 +37,7 @@ pub use maybe_mut::{MaybeMut, MutGuard};
pub use opt::*;
pub use read_only::*;
pub use relations::{nth_relation, relations_like, NthRelation, Relations, RelationsIter};
pub use relations_mut::{relations_like_mut, RelationsIterMut, RelationsMut};
pub use satisfied::Satisfied;
pub use source::Source;
pub use transform::{Added, Modified, TransformFetch};
Expand Down
13 changes: 6 additions & 7 deletions src/fetch/relations.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@ use core::{
};

use alloc::vec::Vec;
use smallvec::SmallVec;

use crate::{
archetype::{CellGuard, Slot},
Expand All @@ -16,7 +15,7 @@ use crate::{

use super::{FetchAccessData, FetchPrepareData, PreparedFetch, RandomFetch};

/// Returns a list of relations of a specified type
/// Returns an iterator of all relations of the specified type
#[derive(Debug, Clone)]
pub struct Relations<T: ComponentValue> {
relation: Relation<T>,
Expand All @@ -31,7 +30,7 @@ where
type Prepared = PreparedRelations<'w, T>;

fn prepare(&self, data: FetchPrepareData<'w>) -> Option<Self::Prepared> {
let borrows: SmallVec<[_; 4]> = {
let borrows = {
data.arch
.relations_like(self.relation.id())
.map(|(desc, &cell_index)| {
Expand Down Expand Up @@ -71,7 +70,7 @@ impl<'q, T: ComponentValue> FetchItem<'q> for Relations<T> {

#[doc(hidden)]
pub struct PreparedRelations<'a, T> {
borrows: SmallVec<[(Entity, CellGuard<'a, [T]>); 4]>,
borrows: Vec<(Entity, CellGuard<'a, [T]>)>,
}

pub struct Batch<'a, T> {
Expand Down Expand Up @@ -123,16 +122,16 @@ impl<'a, T> Iterator for RelationsIter<'a, T> {
}
}

/// Query all relations of the specified kind.
/// Access all relations of the specified type on the entity.
///
/// **Note**: This still matches if there are `0` relations.
/// **Note**: This still matches if there are no relations on the entity
pub fn relations_like<T: ComponentValue>(relation: impl RelationExt<T>) -> Relations<T> {
Relations {
relation: relation.as_relation(),
}
}

/// Query the nth relation of the specified kind.
/// Access the nth relation of the specified kind.
///
/// This is useful for [`Exclusive`](crate::metadata::Exclusive) relations where there is only one parent
///
Expand Down
150 changes: 150 additions & 0 deletions src/fetch/relations_mut.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,150 @@
use core::{
fmt::{self, Formatter},
slice,
};

use alloc::vec::Vec;
use smallvec::SmallVec;

use crate::{
archetype::{Archetype, CellMutGuard, Slice, Slot},
component::ComponentValue,
relation::Relation,
system::{Access, AccessKind},
util::PtrMut,
Entity, Fetch, FetchItem, RelationExt,
};

use super::{FetchAccessData, FetchPrepareData, PreparedFetch};

/// Returns an iterator of all relations of the specified type
#[derive(Debug, Clone)]
pub struct RelationsMut<T: ComponentValue> {
relation: Relation<T>,
}

impl<'w, T> Fetch<'w> for RelationsMut<T>
where
T: ComponentValue,
{
const MUTABLE: bool = false;

type Prepared = PreparedRelationsMut<'w, T>;

fn prepare(&self, data: FetchPrepareData<'w>) -> Option<Self::Prepared> {
let borrows: SmallVec<[_; 4]> = {
data.arch
.relations_like(self.relation.id())
.map(|(desc, &cell_index)| {
(
desc.target.unwrap(),
data.arch.cells()[cell_index].borrow_mut(),
)
})
.collect()
};

Some(PreparedRelationsMut {
borrows,
arch: data.arch,
tick: data.new_tick,
})
}

fn filter_arch(&self, _: FetchAccessData) -> bool {
true
}

fn access(&self, data: FetchAccessData, dst: &mut Vec<Access>) {
let relation = self.relation.id();
let val = data.arch.relations_like(relation).map(|v| Access {
kind: AccessKind::Archetype {
id: data.arch_id,
component: *v.0,
},
mutable: true,
});

dst.extend(val);
}

fn describe(&self, f: &mut Formatter) -> fmt::Result {
write!(f, "relations({})", self.relation)
}
}

impl<'q, T: ComponentValue> FetchItem<'q> for RelationsMut<T> {
type Item = RelationsIterMut<'q, T>;
}

#[doc(hidden)]
pub struct PreparedRelationsMut<'a, T> {
borrows: SmallVec<[(Entity, CellMutGuard<'a, [T]>); 4]>,
arch: &'a Archetype,
tick: u32,
}

pub struct Batch<'a, T> {
borrows: PtrMut<'a, (Entity, CellMutGuard<'a, [T]>)>,
borrow_count: usize,
slot: Slot,
}

impl<'w, 'q, T> PreparedFetch<'q> for PreparedRelationsMut<'w, T>
where
T: 'q + ComponentValue,
{
type Item = RelationsIterMut<'q, T>;

type Chunk = Batch<'q, T>;

const HAS_FILTER: bool = false;

unsafe fn create_chunk(&'q mut self, slots: Slice) -> Self::Chunk {
for (_target, borrow) in &mut self.borrows {
borrow.set_modified(&self.arch.entities[slots.as_range()], slots, self.tick)
}

Batch {
borrow_count: self.borrows.len(),
borrows: PtrMut::new(self.borrows.as_mut_ptr() as _),
slot: slots.start,
}
}

unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
let slot = chunk.slot;
chunk.slot += 1;

RelationsIterMut {
borrows: slice::from_raw_parts_mut(chunk.borrows.as_ptr() as _, chunk.borrow_count)
.iter_mut(),
slot,
}
}
}

/// Iterates the relation targets and data for the yielded query item
pub struct RelationsIterMut<'a, T> {
borrows: slice::IterMut<'a, (Entity, CellMutGuard<'a, [T]>)>,
slot: Slot,
}

impl<'a, T> Iterator for RelationsIterMut<'a, T> {
type Item = (Entity, &'a mut T);

fn next(&mut self) -> Option<Self::Item> {
let (id, borrow) = self.borrows.next()?;
let borrow = &mut borrow.get_mut()[self.slot];
Some((*id, borrow))
}
}

/// Access all relations of the specified type on the entity.
///
/// See: [`relations`](crate::fetch::relations::relations_like)
pub fn relations_like_mut<T: ComponentValue>(relation: impl RelationExt<T>) -> RelationsMut<T> {
RelationsMut {
relation: relation.as_relation(),
}
}
74 changes: 74 additions & 0 deletions tests/relations.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
use flax::{
components::{child_of, name},
events::{Event, EventSubscriber},
fetch::relations_like_mut,
filter::All,
relation::RelationExt,
*,
Expand Down Expand Up @@ -352,3 +354,75 @@ fn exclusive() {
let entity = world.entity_mut(id3).unwrap();
assert_eq!(entity.relations(child_of).map(|v| v.0).collect_vec(), [id2])
}

#[test]
#[cfg(feature = "flume")]
fn relations_mut() {
component! {
relationship(id): f32,
}

let mut world = World::new();

let (changed_tx, changed_rx) = flume::unbounded();

world.subscribe(changed_tx.filter(|_kind, data| data.key.id() == relationship.id()));

let id1 = Entity::builder()
.set(name(), "id1".into())
.spawn(&mut world);

let id2 = Entity::builder()
.set(name(), "id2".into())
.set(relationship(id1), 1.0)
.spawn(&mut world);

let id3 = Entity::builder()
.set(name(), "id3".into())
.set(relationship(id1), 2.0)
.spawn(&mut world);

let id4 = Entity::builder()
.set(name(), "id4".into())
.set(relationship(id2), 3.0)
.set(relationship(id1), 4.0)
.spawn(&mut world);

assert_eq!(
changed_rx.drain().collect_vec(),
[
Event::added(id2, relationship(id1).key()),
Event::added(id3, relationship(id1).key()),
Event::added(id4, relationship(id1).key()),
Event::added(id4, relationship(id2).key()),
]
);

Query::new(relations_like_mut(relationship))
.borrow(&world)
.for_each(|v| v.for_each(|v| *v.1 *= -1.0));

assert_eq!(
Query::new((entity_ids(), relations_like(relationship)))
.borrow(&world)
.iter()
.flat_map(|(id, v)| v.map(move |(target, value)| (id, target, *value)))
.collect_vec(),
[
(id2, id1, -1.0),
(id3, id1, -2.0),
(id4, id1, -4.0),
(id4, id2, -3.0),
]
);

assert_eq!(
changed_rx.drain().collect_vec(),
[
Event::modified(id2, relationship(id1).key()),
Event::modified(id3, relationship(id1).key()),
Event::modified(id4, relationship(id1).key()),
Event::modified(id4, relationship(id2).key()),
]
);
}

0 comments on commit 2e2aac6

Please sign in to comment.