diff --git a/substrate/primitives/externalities/src/extensions.rs b/substrate/primitives/externalities/src/extensions.rs index 8b0bbd2c5921..282e6ea914a8 100644 --- a/substrate/primitives/externalities/src/extensions.rs +++ b/substrate/primitives/externalities/src/extensions.rs @@ -35,17 +35,24 @@ use sp_std::{ /// /// As extensions are stored as `Box`, this trait should give more confidence that the correct /// type is registered and requested. -pub trait Extension: Send + Any { +pub trait Extension: Send + 'static { /// Return the extension as `&mut dyn Any`. /// /// This is a trick to make the trait type castable into an `Any`. fn as_mut_any(&mut self) -> &mut dyn Any; + + /// Get the [`TypeId`] of this `Extension`. + fn type_id(&self) -> TypeId; } impl Extension for Box { fn as_mut_any(&mut self) -> &mut dyn Any { (**self).as_mut_any() } + + fn type_id(&self) -> TypeId { + (**self).type_id() + } } /// Macro for declaring an extension that usable with [`Extensions`]. @@ -74,6 +81,10 @@ macro_rules! decl_extension { fn as_mut_any(&mut self) -> &mut dyn std::any::Any { self } + + fn type_id(&self) -> std::any::TypeId { + std::any::Any::type_id(self) + } } impl std::ops::Deref for $ext_name { @@ -107,6 +118,10 @@ macro_rules! decl_extension { fn as_mut_any(&mut self) -> &mut dyn std::any::Any { self } + + fn type_id(&self) -> std::any::TypeId { + std::any::Any::type_id(self) + } } } } @@ -235,4 +250,25 @@ mod tests { assert_eq!(ext_ty.0, 1); } + + #[test] + fn register_box_extension() { + let mut exts = Extensions::new(); + let box1: Box = Box::new(DummyExt(1)); + let box2: Box = Box::new(DummyExt2(2)); + exts.register(box1); + exts.register(box2); + + { + let ext = exts.get_mut(TypeId::of::()).expect("Extension 1 is registered"); + let ext_ty = ext.downcast_mut::().expect("Downcasting works for Extension 1"); + assert_eq!(ext_ty.0, 1); + } + { + let ext2 = exts.get_mut(TypeId::of::()).expect("Extension 2 is registered"); + let ext_ty2 = + ext2.downcast_mut::().expect("Downcasting works for Extension 2"); + assert_eq!(ext_ty2.0, 2); + } + } }