Working towards BufferKeyMap implementation

Signed-off-by: Michael X. Grey <[email protected]>

Author: mxgrey

macros/src/buffer.rs

@@ -25,7 +25,7 @@ pub(crate) fn impl_joined_value(input_struct: &ItemStruct) -> Result<TokenStream
         }
     };
 
-    let buffer_clone_impl = if noncopy {
+    let impl_buffer_clone = if noncopy {
         // Clone impl for structs with a buffer that is not copyable
         quote! {
             impl #impl_generics ::std::clone::Clone for #buffer_struct_ident #ty_generics #where_clause {
@@ -61,7 +61,7 @@ pub(crate) fn impl_joined_value(input_struct: &ItemStruct) -> Result<TokenStream
 
         #buffer_struct
 
-        #buffer_clone_impl
+        #impl_buffer_clone
 
         impl #impl_generics #struct_ident #ty_generics #where_clause {
             fn select_buffers(

src/buffer.rs

@@ -34,6 +34,7 @@ mod any_buffer;
 pub use any_buffer::*;
 
 mod buffer_access_lifecycle;
+pub use buffer_access_lifecycle::BufferKeyLifecycle;
 pub(crate) use buffer_access_lifecycle::*;
 
 mod buffer_key_builder;
@@ -134,12 +135,15 @@ pub struct BufferLocation {
     pub source: Entity,
 }
 
-#[derive(Clone, Copy)]
+#[derive(Clone)]
 pub struct CloneFromBuffer<T: Clone> {
     pub(crate) location: BufferLocation,
     pub(crate) _ignore: std::marker::PhantomData<fn(T)>,
 }
 
+//
+impl<T: Clone> Copy for CloneFromBuffer<T> {}
+
 impl<T: Clone> CloneFromBuffer<T> {
     /// Get the entity ID of the buffer.
     pub fn id(&self) -> Entity {
@@ -157,6 +161,15 @@ impl<T: Clone> CloneFromBuffer<T> {
     }
 }
 
+impl<T: Clone> From<CloneFromBuffer<T>> for Buffer<T> {
+    fn from(value: CloneFromBuffer<T>) -> Self {
+        Buffer {
+            location: value.location,
+            _ignore: Default::default(),
+        }
+    }
+}
+
 /// Settings to describe the behavior of a buffer.
 #[derive(Default, Clone, Copy)]
 pub struct BufferSettings {
@@ -254,20 +267,23 @@ impl<T> BufferKey<T> {
     pub fn tag(&self) -> &BufferKeyTag {
         &self.tag
     }
+}
+
+impl<T> BufferKeyLifecycle for BufferKey<T> {
+    type TargetBuffer = Buffer<T>;
+
+    fn create_key(buffer: &Self::TargetBuffer, builder: &BufferKeyBuilder) -> Self {
+        BufferKey {
+            tag: builder.make_tag(buffer.id()),
+            _ignore: Default::default(),
+        }
+    }
 
-    pub(crate) fn is_in_use(&self) -> bool {
+    fn is_in_use(&self) -> bool {
         self.tag.is_in_use()
     }
 
-    // We do a deep clone of the key when distributing it to decouple the
-    // lifecycle of the keys that we send out from the key that's held by the
-    // accessor node.
-    //
-    // The key instance held by the accessor node will never be dropped until
-    // the session is cleaned up, so the keys that we send out into the workflow
-    // 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 {
+    fn deep_clone(&self) -> Self {
         Self {
             tag: self.tag.deep_clone(),
             _ignore: Default::default(),

src/buffer/any_buffer.rs

@@ -35,9 +35,9 @@ use smallvec::SmallVec;
 
 use crate::{
     add_listener_to_source, Accessed, Buffer, BufferAccessMut, BufferAccessors, BufferError,
-    BufferKey, BufferKeyTag, BufferLocation, BufferStorage, Bufferable, Buffered, Builder,
-    DrainBuffer, Gate, GateState, InspectBuffer, Joined, ManageBuffer, NotifyBufferUpdate,
-    OperationError, OperationResult, OperationRoster, OrBroken,
+    BufferKey, BufferKeyBuilder, BufferKeyLifecycle, 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
@@ -198,12 +198,23 @@ impl AnyBufferKey {
     pub fn session(&self) -> Entity {
         self.tag.session
     }
+}
+
+impl BufferKeyLifecycle for AnyBufferKey {
+    type TargetBuffer = AnyBuffer;
+
+    fn create_key(buffer: &AnyBuffer, builder: &BufferKeyBuilder) -> Self {
+        AnyBufferKey {
+            tag: builder.make_tag(buffer.id()),
+            interface: buffer.interface,
+        }
+    }
 
-    pub(crate) fn is_in_use(&self) -> bool {
+    fn is_in_use(&self) -> bool {
         self.tag.is_in_use()
     }
 
-    pub(crate) fn deep_clone(&self) -> Self {
+    fn deep_clone(&self) -> Self {
         Self {
             tag: self.tag.deep_clone(),
             interface: self.interface,

src/buffer/buffer_access_lifecycle.rs

@@ -21,7 +21,7 @@ use tokio::sync::mpsc::UnboundedSender as TokioSender;
 
 use std::sync::Arc;
 
-use crate::{emit_disposal, ChannelItem, Disposal, OperationRoster};
+use crate::{emit_disposal, BufferKeyBuilder, ChannelItem, Disposal, OperationRoster};
 
 /// This is used as a field inside of [`crate::BufferKey`] which keeps track of
 /// when a key that was sent out into the world gets fully dropped from use. We
@@ -87,3 +87,29 @@ impl Drop for BufferAccessLifecycle {
         }
     }
 }
+
+/// This trait is implemented by [`crate::BufferKey`]-like structs so their
+/// lifecycles can be managed.
+pub trait BufferKeyLifecycle {
+    /// What kind of buffer this key can unlock.
+    type TargetBuffer;
+
+    /// Create a new key of this type.
+    fn create_key(buffer: &Self::TargetBuffer, builder: &BufferKeyBuilder) -> Self;
+
+    /// Check if the key is currently in use.
+    fn is_in_use(&self) -> bool;
+
+    /// Create a deep clone of the key. The usage tracking of the clone will
+    /// be unrelated to the usage tracking of the original.
+    ///
+    /// We do a deep clone of the key when distributing it to decouple the
+    /// lifecycle of the keys that we send out from the key that's held by the
+    /// accessor node.
+    //
+    /// The key instance held by the accessor node will never be dropped until
+    /// the session is cleaned up, so the keys that we send out into the workflow
+    /// need to have their own independent lifecycles or else we won't detect
+    /// when the workflow has dropped them.
+    fn deep_clone(&self) -> Self;
+}

src/buffer/buffer_key_builder.rs

@@ -19,7 +19,7 @@ use bevy_ecs::prelude::Entity;
 
 use std::sync::Arc;
 
-use crate::{BufferAccessLifecycle, BufferKey, BufferKeyTag, ChannelSender};
+use crate::{BufferAccessLifecycle, BufferKeyTag, ChannelSender};
 
 pub struct BufferKeyBuilder {
     scope: Entity,
@@ -29,16 +29,8 @@ pub struct BufferKeyBuilder {
 }
 
 impl BufferKeyBuilder {
-    pub(crate) fn build<T>(&self, buffer: Entity) -> BufferKey<T> {
-        BufferKey {
-            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 make_tag(&self, buffer: Entity) -> BufferKeyTag {
+    /// Make a [`BufferKeyTag`] that can be given to a [`crate::BufferKey`]-like struct.
+    pub fn make_tag(&self, buffer: Entity) -> BufferKeyTag {
         BufferKeyTag {
             buffer,
             session: self.session,