Refactor buffer keys

Signed-off-by: Michael X. Grey <greyxmike@gmail.com>

Author: mxgrey

src/buffer.rs

@@ -123,7 +123,7 @@ impl<T> Clone for Buffer<T> {
 
 impl<T> Copy for Buffer<T> {}
 
-/// Get the general identifying information for a buffer to locate it within the
+/// The general identifying information for a buffer to locate it within the
 /// world. This does not indicate anything about the type of messages that the
 /// buffer can contain.
 #[derive(Clone, Copy, Debug)]
@@ -227,38 +227,36 @@ impl Default for RetentionPolicy {
 /// [1]: crate::Chain::with_access
 /// [2]: crate::Bufferable::listen
 pub struct BufferKey<T> {
-    buffer: Entity,
-    session: Entity,
-    accessor: Entity,
-    lifecycle: Option<Arc<BufferAccessLifecycle>>,
+    tag: BufferKeyTag,
     _ignore: std::marker::PhantomData<fn(T)>,
 }
 
 impl<T> Clone for BufferKey<T> {
     fn clone(&self) -> Self {
         Self {
-            buffer: self.buffer,
-            session: self.session,
-            accessor: self.accessor,
-            lifecycle: self.lifecycle.as_ref().map(Arc::clone),
+            tag: self.tag.clone(),
             _ignore: Default::default(),
         }
     }
 }
 
 impl<T> BufferKey<T> {
     /// The buffer ID of this key.
-    pub fn id(&self) -> Entity {
-        self.buffer
+    pub fn buffer(&self) -> Entity {
+        self.tag.buffer
     }
 
     /// The session that this key belongs to.
     pub fn session(&self) -> Entity {
-        self.session
+        self.tag.session
+    }
+
+    pub(crate) fn tag(&self) -> &BufferKeyTag {
+        &self.tag
     }
 
     pub(crate) fn is_in_use(&self) -> bool {
-        self.lifecycle.as_ref().is_some_and(|l| l.is_in_use())
+        self.tag.is_in_use()
     }
 
     // We do a deep clone of the key when distributing it to decouple the
@@ -270,6 +268,41 @@ impl<T> BufferKey<T> {
     // need to have their own independent lifecycles or else we won't detect
     // when the workflow has dropped them.
     pub(crate) fn deep_clone(&self) -> Self {
+        Self {
+            tag: self.tag.deep_clone(),
+            _ignore: Default::default(),
+        }
+    }
+}
+
+impl<T> std::fmt::Debug for BufferKey<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("BufferKey")
+            .field("message_type_name", &std::any::type_name::<T>())
+            .field("tag", &self.tag)
+            .finish()
+    }
+}
+
+/// The identifying information for a buffer key. This does not indicate
+/// anything about the type of messages that the buffer can contain.
+///
+/// This struct will be internal to the crate until we decide to make
+/// [`BufferAccessLifecycle`] a public struct.
+#[derive(Clone)]
+pub(crate) struct BufferKeyTag {
+    pub buffer: Entity,
+    pub session: Entity,
+    pub accessor: Entity,
+    pub lifecycle: Option<Arc<BufferAccessLifecycle>>,
+}
+
+impl BufferKeyTag {
+    pub fn is_in_use(&self) -> bool {
+        self.lifecycle.as_ref().is_some_and(|l| l.is_in_use())
+    }
+
+    pub fn deep_clone(&self) -> Self {
         let mut deep = self.clone();
         deep.lifecycle = self
             .lifecycle
@@ -279,6 +312,17 @@ impl<T> BufferKey<T> {
     }
 }
 
+impl std::fmt::Debug for BufferKeyTag {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("BufferKeyTag")
+            .field("buffer", &self.buffer)
+            .field("session", &self.session)
+            .field("accessor", &self.accessor)
+            .field("in_use", &self.is_in_use())
+            .finish()
+    }
+}
+
 /// This system parameter lets you get read-only access to a buffer that exists
 /// within a workflow. Use a [`BufferKey`] to unlock the access.
 ///
@@ -293,9 +337,9 @@ where
 
 impl<'w, 's, T: 'static + Send + Sync> BufferAccess<'w, 's, T> {
     pub fn get<'a>(&'a self, key: &BufferKey<T>) -> Result<BufferView<'a, T>, QueryEntityError> {
-        let session = key.session;
+        let session = key.session();
         self.query
-            .get(key.buffer)
+            .get(key.buffer())
             .map(|(storage, gate)| BufferView {
                 storage,
                 gate,
@@ -322,9 +366,9 @@ where
     T: 'static + Send + Sync,
 {
     pub fn get<'a>(&'a self, key: &BufferKey<T>) -> Result<BufferView<'a, T>, QueryEntityError> {
-        let session = key.session;
+        let session = key.session();
         self.query
-            .get(key.buffer)
+            .get(key.buffer())
             .map(|(storage, gate)| BufferView {
                 storage,
                 gate,
@@ -336,10 +380,10 @@ where
         &'a mut self,
         key: &BufferKey<T>,
     ) -> Result<BufferMut<'w, 's, 'a, T>, QueryEntityError> {
-        let buffer = key.buffer;
-        let session = key.session;
-        let accessor = key.accessor;
-        self.query.get_mut(key.buffer).map(|(storage, gate)| {
+        let buffer = key.buffer();
+        let session = key.session();
+        let accessor = key.tag.accessor;
+        self.query.get_mut(key.buffer()).map(|(storage, gate)| {
             BufferMut::new(storage, gate, buffer, session, accessor, &mut self.commands)
         })
     }

src/buffer/any_buffer.rs

@@ -35,9 +35,9 @@ use smallvec::SmallVec;
 
 use crate::{
     add_listener_to_source, Accessed, Buffer, BufferAccessLifecycle, BufferAccessMut,
-    BufferAccessors, BufferError, BufferKey, BufferLocation, BufferStorage, Bufferable, Buffered,
-    Builder, DrainBuffer, Gate, GateState, InspectBuffer, Joined, ManageBuffer, NotifyBufferUpdate,
-    OperationError, OperationResult, OperationRoster, OrBroken,
+    BufferAccessors, BufferError, BufferKey, BufferKeyTag, BufferLocation, BufferStorage,
+    Bufferable, Buffered, Builder, DrainBuffer, Gate, GateState, InspectBuffer, Joined,
+    ManageBuffer, NotifyBufferUpdate, OperationError, OperationResult, OperationRoster, OrBroken,
 };
 
 /// A [`Buffer`] whose message type has been anonymized. Joining with this buffer
@@ -138,10 +138,7 @@ impl<T: 'static + Send + Sync + Any> From<Buffer<T>> for AnyBuffer {
 /// [1]: bevy_ecs::prelude::World
 #[derive(Clone)]
 pub struct AnyBufferKey {
-    pub(crate) buffer: Entity,
-    pub(crate) session: Entity,
-    pub(crate) accessor: Entity,
-    pub(crate) lifecycle: Option<Arc<BufferAccessLifecycle>>,
+    pub(crate) tag: BufferKeyTag,
     pub(crate) interface: &'static (dyn AnyBufferAccessInterface + Send + Sync),
 }
 
@@ -150,10 +147,7 @@ impl AnyBufferKey {
     pub fn downcast<T: 'static>(&self) -> Option<BufferKey<T>> {
         if TypeId::of::<T>() == self.interface.message_type_id() {
             Some(BufferKey {
-                buffer: self.buffer,
-                session: self.session,
-                accessor: self.accessor,
-                lifecycle: self.lifecycle.clone(),
+                tag: self.tag.clone(),
                 _ignore: Default::default(),
             })
         } else {
@@ -163,39 +157,31 @@ impl AnyBufferKey {
 
     /// The buffer ID of this key.
     pub fn id(&self) -> Entity {
-        self.buffer
+        self.tag.buffer
     }
 
     /// The session that this key belongs to.
     pub fn session(&self) -> Entity {
-        self.session
+        self.tag.session
     }
 
-    fn deep_clone(&self) -> Self {
-        let mut deep = self.clone();
-        deep.lifecycle = self
-            .lifecycle
-            .as_ref()
-            .map(|l| Arc::new(l.as_ref().clone()));
-        deep
+    fn is_in_use(&self) -> bool {
+        self.tag.is_in_use()
     }
 
-    fn is_in_use(&self) -> bool {
-        self.lifecycle.as_ref().is_some_and(|l| l.is_in_use())
+    fn deep_clone(&self) -> Self {
+        Self {
+            tag: self.tag.deep_clone(),
+            interface: self.interface,
+        }
     }
 }
 
 impl std::fmt::Debug for AnyBufferKey {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("AnyBufferKey")
-            .field("buffer", &self.buffer)
-            .field("session", &self.session)
-            .field("accessor", &self.accessor)
-            .field(
-                "in_use",
-                &self.lifecycle.as_ref().is_some_and(|l| l.is_in_use()),
-            )
             .field("message_type_name", &self.interface.message_type_name())
+            .field("tag", &self.tag)
             .finish()
     }
 }
@@ -204,10 +190,7 @@ impl<T: 'static + Send + Sync + Any> From<BufferKey<T>> for AnyBufferKey {
     fn from(value: BufferKey<T>) -> Self {
         let interface = AnyBuffer::interface_for::<T>();
         AnyBufferKey {
-            buffer: value.buffer,
-            session: value.session,
-            accessor: value.accessor,
-            lifecycle: value.lifecycle.clone(),
+            tag: value.tag,
             interface,
         }
     }
@@ -770,14 +753,14 @@ impl<'w, 's, T: 'static + Send + Sync + Any> AnyBufferAccessMut<'w, 's>
     ) -> Result<AnyBufferMut<'w, 's, 'a>, BufferError> {
         let BufferAccessMut { query, commands } = self;
         let (storage, gate) = query
-            .get_mut(key.buffer)
+            .get_mut(key.tag.buffer)
             .map_err(|_| BufferError::BufferMissing)?;
         Ok(AnyBufferMut {
             storage: Box::new(storage),
             gate,
-            buffer: key.buffer,
-            session: key.session,
-            accessor: Some(key.accessor),
+            buffer: key.tag.buffer,
+            session: key.tag.session,
+            accessor: Some(key.tag.accessor),
             commands,
             modified: false,
         })
@@ -895,7 +878,7 @@ impl<T: 'static + Send + Sync + Any> AnyBufferAccessInterface for AnyBufferAcces
         world: &'a World,
     ) -> Result<AnyBufferView<'a>, BufferError> {
         let buffer_ref = world
-            .get_entity(key.buffer)
+            .get_entity(key.tag.buffer)
             .ok_or(BufferError::BufferMissing)?;
         let storage = buffer_ref
             .get::<BufferStorage<T>>()
@@ -906,7 +889,7 @@ impl<T: 'static + Send + Sync + Any> AnyBufferAccessInterface for AnyBufferAcces
         Ok(AnyBufferView {
             storage: Box::new(storage),
             gate,
-            session: key.session,
+            session: key.tag.session,
         })
     }
 
@@ -1001,12 +984,8 @@ impl Accessed for AnyBuffer {
     }
 
     fn create_key(&self, builder: &super::BufferKeyBuilder) -> Self::Key {
-        let components = builder.as_components(self.id());
         AnyBufferKey {
-            buffer: components.buffer,
-            session: components.session,
-            accessor: components.accessor,
-            lifecycle: components.lifecycle,
+            tag: builder.make_tag(self.id()),
             interface: self.interface,
         }
     }

src/buffer/buffer_key_builder.rs

@@ -19,7 +19,7 @@ use bevy_ecs::prelude::Entity;
 
 use std::sync::Arc;
 
-use crate::{BufferAccessLifecycle, BufferKey, ChannelSender};
+use crate::{BufferAccessLifecycle, BufferKey, BufferKeyTag, ChannelSender};
 
 pub struct BufferKeyBuilder {
     scope: Entity,
@@ -28,37 +28,18 @@ pub struct BufferKeyBuilder {
     lifecycle: Option<(ChannelSender, Arc<()>)>,
 }
 
-pub struct BufferKeyComponents {
-    pub buffer: Entity,
-    pub session: Entity,
-    pub accessor: Entity,
-    pub lifecycle: Option<Arc<BufferAccessLifecycle>>,
-}
-
 impl BufferKeyBuilder {
     pub(crate) fn build<T>(&self, buffer: Entity) -> BufferKey<T> {
         BufferKey {
-            buffer,
-            session: self.session,
-            accessor: self.accessor,
-            lifecycle: self.lifecycle.as_ref().map(|(sender, tracker)| {
-                Arc::new(BufferAccessLifecycle::new(
-                    self.scope,
-                    buffer,
-                    self.session,
-                    self.accessor,
-                    sender.clone(),
-                    tracker.clone(),
-                ))
-            }),
+            tag: self.make_tag(buffer),
             _ignore: Default::default(),
         }
     }
 
     // TODO(@mxgrey): Consider refactoring all the buffer key structs to use a
     // single inner struct like BufferKeyComponents
-    pub(crate) fn as_components(&self, buffer: Entity) -> BufferKeyComponents {
-        BufferKeyComponents {
+    pub(crate) fn make_tag(&self, buffer: Entity) -> BufferKeyTag {
+        BufferKeyTag {
             buffer,
             session: self.session,
             accessor: self.accessor,