From 6000d786c554426b875d82c5d894b0a1e97894c6 Mon Sep 17 00:00:00 2001
From: Benjamin Saunders <ben.e.saunders@gmail.com>
Date: Sat, 17 Oct 2020 17:27:03 -0700
Subject: [PATCH] Introduce mutation tracking query helper

---
 src/lib.rs     |   2 +
 src/tracked.rs | 124 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 126 insertions(+)
 create mode 100644 src/tracked.rs

diff --git a/src/lib.rs b/src/lib.rs
index 8fcb2415..47a2ba90 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -69,6 +69,7 @@ mod entities;
 mod entity_builder;
 mod query;
 mod query_one;
+mod tracked;
 mod world;
 
 pub use archetype::Archetype;
@@ -78,6 +79,7 @@ pub use entities::{Entity, NoSuchEntity};
 pub use entity_builder::{BuiltEntity, EntityBuilder};
 pub use query::{Access, BatchedIter, Query, QueryBorrow, QueryIter, With, Without};
 pub use query_one::QueryOne;
+pub use tracked::{Modified, Tracked};
 pub use world::{ArchetypesGeneration, Component, ComponentError, Iter, SpawnBatchIter, World};
 
 // Unstable implementation details needed by the macros
diff --git a/src/tracked.rs b/src/tracked.rs
new file mode 100644
index 00000000..acce2723
--- /dev/null
+++ b/src/tracked.rs
@@ -0,0 +1,124 @@
+use core::{
+    marker::PhantomData,
+    ops::{Deref, DerefMut},
+};
+
+use crate::{Access, Archetype, Component, Fetch, Query};
+
+/// Query that tracks mutable access to a component
+///
+/// Using this in a query is equivalent to `(&mut T, &mut Modified<T>)`, except that it yields a
+/// smart pointer to `T` which sets the flag inside `Modified<T>` to `true` when it's mutably
+/// borrowed.
+///
+/// A `Modified<T>` component must exist on an entity for it to be exposed to this query.
+///
+/// # Example
+/// ```
+/// # use hecs::*;
+/// let mut world = World::new();
+/// let e = world.spawn((123, Modified::<i32>::new()));
+/// for (_id, mut value) in world.query::<Tracked<i32>>().iter() {
+///   assert_eq!(*value, 123);
+/// }
+/// assert!(!world.get::<Modified<i32>>(e).unwrap().is_set());
+/// for (_id, mut value) in world.query::<Tracked<i32>>().iter() {
+///   *value = 42;
+/// }
+/// assert!(world.get::<Modified<i32>>(e).unwrap().is_set());
+/// ```
+pub struct Tracked<'a, T: Component> {
+    value: &'a mut T,
+    modified: &'a mut Modified<T>,
+}
+
+impl<'a, T: Component> Deref for Tracked<'a, T> {
+    type Target = T;
+    #[inline]
+    fn deref(&self) -> &T {
+        self.value
+    }
+}
+
+impl<'a, T: Component> DerefMut for Tracked<'a, T> {
+    #[inline]
+    fn deref_mut(&mut self) -> &mut T {
+        self.modified.0 = true;
+        self.value
+    }
+}
+
+impl<'a, T: Component> Query for Tracked<'a, T> {
+    type Fetch = FetchTracked<T>;
+}
+
+/// A flag indicating whether the `T` component was modified
+///
+/// Must be manually added to components that will be queried with `Tracked`.
+pub struct Modified<T>(bool, PhantomData<T>);
+
+impl<T> Modified<T> {
+    /// Constructs an unset flag
+    #[inline]
+    pub fn new() -> Self {
+        Self(false, PhantomData)
+    }
+
+    /// Returns whether the `T` component was modified since the last `unset` call
+    #[inline]
+    pub fn is_set(&self) -> bool {
+        self.0
+    }
+
+    /// Unsets the flag
+    #[inline]
+    pub fn unset(&mut self) {
+        self.0 = false;
+    }
+}
+
+impl<T> Default for Modified<T> {
+    #[inline]
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[doc(hidden)]
+pub struct FetchTracked<T: Component> {
+    value: <&'static mut T as Query>::Fetch,
+    modified: <&'static mut Modified<T> as Query>::Fetch,
+}
+
+impl<'a, T: Component> Fetch<'a> for FetchTracked<T> {
+    type Item = Tracked<'a, T>;
+
+    fn access(archetype: &Archetype) -> Option<Access> {
+        Some(
+            <&'a mut T as Query>::Fetch::access(archetype)?
+                .max(<&'a mut Modified<T> as Query>::Fetch::access(archetype)?),
+        )
+    }
+
+    fn borrow(archetype: &Archetype) {
+        <&'a mut T as Query>::Fetch::borrow(archetype);
+        <&'a mut Modified<T> as Query>::Fetch::borrow(archetype);
+    }
+
+    unsafe fn get(archetype: &'a Archetype, offset: usize) -> Option<Self> {
+        Some(Self {
+            value: <&'a mut T as Query>::Fetch::get(archetype, offset)?,
+            modified: <&'a mut Modified<T> as Query>::Fetch::get(archetype, offset)?,
+        })
+    }
+
+    fn release(archetype: &Archetype) {
+        <&'a mut T as Query>::Fetch::release(archetype);
+        <&'a mut Modified<T> as Query>::Fetch::release(archetype);
+    }
+
+    unsafe fn next(&mut self) -> Self::Item {
+        let (value, modified) = (self.value.next(), self.modified.next());
+        Tracked { value, modified }
+    }
+}