use parking_lot::RwLock;
use sp_core::traits::{ReadRuntimeVersion, ReadRuntimeVersionExt};
use sp_externalities::{Extension, Extensions};
use sp_runtime::traits::{Block as BlockT, NumberFor};
use std::{marker::PhantomData, sync::Arc};
pub trait ExtensionsFactory<Block: BlockT>: Send + Sync {
fn extensions_for(&self, block_hash: Block::Hash, block_number: NumberFor<Block>)
-> Extensions;
}
impl<Block: BlockT> ExtensionsFactory<Block> for () {
fn extensions_for(&self, _: Block::Hash, _: NumberFor<Block>) -> Extensions {
Extensions::new()
}
}
impl<Block: BlockT, T: ExtensionsFactory<Block>> ExtensionsFactory<Block> for Vec<T> {
fn extensions_for(
&self,
block_hash: Block::Hash,
block_number: NumberFor<Block>,
) -> Extensions {
let mut exts = Extensions::new();
exts.extend(self.iter().map(|e| e.extensions_for(block_hash, block_number)));
exts
}
}
pub struct ExtensionBeforeBlock<Block: BlockT, Ext> {
before: NumberFor<Block>,
_marker: PhantomData<fn(Ext) -> Ext>,
}
impl<Block: BlockT, Ext> ExtensionBeforeBlock<Block, Ext> {
pub fn new(before: NumberFor<Block>) -> Self {
Self { before, _marker: PhantomData }
}
}
impl<Block: BlockT, Ext: Default + Extension> ExtensionsFactory<Block>
for ExtensionBeforeBlock<Block, Ext>
{
fn extensions_for(&self, _: Block::Hash, block_number: NumberFor<Block>) -> Extensions {
let mut exts = Extensions::new();
if block_number < self.before {
exts.register(Ext::default());
}
exts
}
}
pub struct ExecutionExtensions<Block: BlockT> {
extensions_factory: RwLock<Box<dyn ExtensionsFactory<Block>>>,
read_runtime_version: Arc<dyn ReadRuntimeVersion>,
}
impl<Block: BlockT> ExecutionExtensions<Block> {
pub fn new(
extensions_factory: Option<Box<dyn ExtensionsFactory<Block>>>,
read_runtime_version: Arc<dyn ReadRuntimeVersion>,
) -> Self {
Self {
extensions_factory: extensions_factory
.map(RwLock::new)
.unwrap_or_else(|| RwLock::new(Box::new(()))),
read_runtime_version,
}
}
pub fn set_extensions_factory(&self, maker: impl ExtensionsFactory<Block> + 'static) {
*self.extensions_factory.write() = Box::new(maker);
}
pub fn extensions(
&self,
block_hash: Block::Hash,
block_number: NumberFor<Block>,
) -> Extensions {
let mut extensions =
self.extensions_factory.read().extensions_for(block_hash, block_number);
extensions.register(ReadRuntimeVersionExt::new(self.read_runtime_version.clone()));
extensions
}
}