Implement more general buffer key downcasting

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

Author: mxgrey

src/buffer.rs

@@ -251,7 +251,7 @@ impl<T> BufferKey<T> {
         self.tag.session
     }
 
-    pub(crate) fn tag(&self) -> &BufferKeyTag {
+    pub fn tag(&self) -> &BufferKeyTag {
         &self.tag
     }
 
@@ -286,11 +286,8 @@ impl<T> std::fmt::Debug for BufferKey<T> {
 
 /// 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 struct BufferKeyTag {
     pub buffer: Entity,
     pub session: Entity,
     pub accessor: Entity,

src/buffer/any_buffer.rs

@@ -21,7 +21,7 @@ use std::{
     any::{Any, TypeId},
     collections::{hash_map::Entry, HashMap},
     ops::RangeBounds,
-    sync::{Arc, Mutex, OnceLock},
+    sync::{Mutex, OnceLock},
 };
 
 use bevy_ecs::{
@@ -34,10 +34,10 @@ use thiserror::Error as ThisError;
 use smallvec::SmallVec;
 
 use crate::{
-    add_listener_to_source, Accessed, Buffer, BufferAccessLifecycle, BufferAccessMut,
-    BufferAccessors, BufferError, BufferKey, BufferKeyTag, BufferLocation, BufferStorage,
-    Bufferable, Buffered, Builder, DrainBuffer, Gate, GateState, InspectBuffer, Joined,
-    ManageBuffer, NotifyBufferUpdate, OperationError, OperationResult, OperationRoster, OrBroken,
+    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,
 };
 
 /// A [`Buffer`] whose message type has been anonymized. Joining with this buffer
@@ -143,18 +143,29 @@ pub struct AnyBufferKey {
 }
 
 impl AnyBufferKey {
-    /// Downcast this into a concrete [`BufferKey`] type.
-    pub fn downcast<T: 'static>(&self) -> Option<BufferKey<T>> {
-        if TypeId::of::<T>() == self.interface.message_type_id() {
+    /// Downcast this into a concrete [`BufferKey`] for the specified message type.
+    ///
+    /// To downcast to a specialized kind of key, use [`Self::downcast_buffer_key`] instead.
+    pub fn downcast_for_message<Message: 'static>(self) -> Option<BufferKey<Message>> {
+        if TypeId::of::<Message>() == self.interface.message_type_id() {
             Some(BufferKey {
-                tag: self.tag.clone(),
+                tag: self.tag,
                 _ignore: Default::default(),
             })
         } else {
             None
         }
     }
 
+    /// Downcast this into a different special buffer key representation, such
+    /// as a `JsonBufferKey`.
+    pub fn downcast_buffer_key<KeyType: 'static>(self) -> Option<KeyType> {
+        self.interface.key_downcast(TypeId::of::<KeyType>())?(self.tag)
+            .downcast::<KeyType>()
+            .ok()
+            .map(|x| *x)
+    }
+
     /// The buffer ID of this key.
     pub fn id(&self) -> Entity {
         self.tag.buffer
@@ -780,6 +791,10 @@ pub trait AnyBufferAccessInterface {
 
     fn buffer_downcast(&self, buffer_type: TypeId) -> Option<BufferDowncastRef>;
 
+    fn register_key_downcast(&self, key_type: TypeId, f: KeyDowncastBox);
+
+    fn key_downcast(&self, key_type: TypeId) -> Option<KeyDowncastRef>;
+
     fn pull(
         &self,
         entity_mut: &mut EntityWorldMut,
@@ -800,16 +815,23 @@ pub trait AnyBufferAccessInterface {
 
 pub type BufferDowncastBox = Box<dyn Fn(BufferLocation) -> Box<dyn Any> + Send + Sync>;
 pub type BufferDowncastRef = &'static (dyn Fn(BufferLocation) -> Box<dyn Any> + Send + Sync);
+pub type KeyDowncastBox = Box<dyn Fn(BufferKeyTag) -> Box<dyn Any> + Send + Sync>;
+pub type KeyDowncastRef = &'static (dyn Fn(BufferKeyTag) -> Box<dyn Any> + Send + Sync);
 
 struct AnyBufferAccessImpl<T> {
     buffer_downcasts: Mutex<HashMap<TypeId, BufferDowncastRef>>,
+    key_downcasts: Mutex<HashMap<TypeId, KeyDowncastRef>>,
     _ignore: std::marker::PhantomData<fn(T)>,
 }
 
 impl<T: 'static + Send + Sync> AnyBufferAccessImpl<T> {
     fn new() -> Self {
         let mut buffer_downcasts: HashMap<_, BufferDowncastRef> = HashMap::new();
 
+        // SAFETY: These leaks are okay because we will only ever instantiate
+        // AnyBufferAccessImpl once per generic argument T, which puts a firm
+        // ceiling on how many of these callbacks will get leaked.
+
         // Automatically register a downcast into AnyBuffer
         buffer_downcasts.insert(
             TypeId::of::<AnyBuffer>(),
@@ -821,8 +843,22 @@ impl<T: 'static + Send + Sync> AnyBufferAccessImpl<T> {
             })),
         );
 
+        let mut key_downcasts: HashMap<_, KeyDowncastRef> = HashMap::new();
+
+        // Automatically register a downcast to AnyBufferKey
+        key_downcasts.insert(
+            TypeId::of::<AnyBufferKey>(),
+            Box::leak(Box::new(|tag| -> Box<dyn Any> {
+                Box::new(AnyBufferKey {
+                    tag,
+                    interface: AnyBuffer::interface_for::<T>(),
+                })
+            })),
+        );
+
         Self {
             buffer_downcasts: Mutex::new(buffer_downcasts),
+            key_downcasts: Mutex::new(key_downcasts),
             _ignore: Default::default(),
         }
     }
@@ -862,6 +898,19 @@ impl<T: 'static + Send + Sync + Any> AnyBufferAccessInterface for AnyBufferAcces
             .copied()
     }
 
+    fn register_key_downcast(&self, key_type: TypeId, f: KeyDowncastBox) {
+        let mut downcasts = self.key_downcasts.lock().unwrap();
+
+        if let Entry::Vacant(entry) = downcasts.entry(key_type) {
+            // We should only leak this in to the register once per type
+            entry.insert(Box::leak(f));
+        }
+    }
+
+    fn key_downcast(&self, key_type: TypeId) -> Option<KeyDowncastRef> {
+        self.key_downcasts.lock().unwrap().get(&key_type).copied()
+    }
+
     fn pull(
         &self,
         entity_mut: &mut EntityWorldMut,

src/buffer/buffer_access_lifecycle.rs

@@ -29,7 +29,7 @@ use crate::{emit_disposal, ChannelItem, Disposal, OperationRoster};
 /// we would be needlessly doing a reachability check every time the key gets
 /// cloned.
 #[derive(Clone)]
-pub(crate) struct BufferAccessLifecycle {
+pub struct BufferAccessLifecycle {
     scope: Entity,
     accessor: Entity,
     session: Entity,

src/buffer/json_buffer.rs

@@ -21,7 +21,7 @@ use std::{
     any::TypeId,
     collections::HashMap,
     ops::RangeBounds,
-    sync::{Arc, Mutex, OnceLock},
+    sync::{Mutex, OnceLock},
 };
 
 use bevy_ecs::{
@@ -143,20 +143,19 @@ pub struct JsonBufferKey {
 
 impl JsonBufferKey {
     /// Downcast this into a concrete [`BufferKey`] for the specified message type.
-    pub fn downcast_for_message<T: 'static>(&self) -> Option<BufferKey<T>> {
-        if TypeId::of::<T>() == self.interface.any_access_interface().message_type_id() {
-            Some(BufferKey {
-                tag: self.tag.clone(),
-                _ignore: Default::default(),
-            })
-        } else {
-            None
-        }
+    ///
+    /// To downcast to a specialized kind of key, use [`Self::downcast_buffer_key`] instead.
+    pub fn downcast_for_message<T: 'static>(self) -> Option<BufferKey<T>> {
+        self.as_any_buffer_key().downcast_for_message()
+    }
+
+    pub fn downcast_buffer_key<KeyType: 'static>(self) -> Option<KeyType> {
+        self.as_any_buffer_key().downcast_buffer_key()
     }
 
     /// Cast this into an [`AnyBufferKey`]
-    pub fn as_any_buffer_key(&self) -> AnyBufferKey {
-        self.clone().into()
+    pub fn as_any_buffer_key(self) -> AnyBufferKey {
+        self.into()
     }
 
     fn deep_clone(&self) -> Self {
@@ -816,6 +815,16 @@ impl<T: 'static + Send + Sync + Serialize + DeserializeOwned> JsonBufferAccessIm
             }),
         );
 
+        any_interface.register_key_downcast(
+            TypeId::of::<JsonBufferKey>(),
+            Box::new(|tag| {
+                Box::new(JsonBufferKey {
+                    tag,
+                    interface: Self::get_interface(),
+                })
+            }),
+        );
+
         // Create and cache the json buffer access interface
         static INTERFACE_MAP: OnceLock<
             Mutex<HashMap<TypeId, &'static (dyn JsonBufferAccessInterface + Send + Sync)>>,
@@ -1319,7 +1328,21 @@ mod tests {
             let _original_from_json: Buffer<TestMessage> =
                 json_buffer.downcast_for_message().unwrap();
 
-            builder.connect(scope.input, scope.terminate);
+            scope
+                .input
+                .chain(builder)
+                .with_access(buffer)
+                .map_block(|(data, key)| {
+                    let any_key: AnyBufferKey = key.clone().into();
+                    let json_key: JsonBufferKey = any_key.clone().downcast_buffer_key().unwrap();
+                    let _original_from_any: BufferKey<TestMessage> =
+                        any_key.downcast_for_message().unwrap();
+                    let _original_from_json: BufferKey<TestMessage> =
+                        json_key.downcast_for_message().unwrap();
+
+                    data
+                })
+                .connect(scope.terminate);
         });
 
         let mut promise = context.command(|commands| commands.request(1, workflow).take_response());