use crate::effects::extension::{ExtensionEffect, ExtensionError};
use crate::effects::{
contract::{ExtensionDispatchContract, ExtensionDispatchMode},
Endpoint, RoleId,
};
use std::any::TypeId;
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
pub type ExtensionHandler<E, R> = Box<
dyn for<'a> Fn(
&'a mut E,
&'a dyn ExtensionEffect<R>,
) -> BoxFuture<'a, Result<(), ExtensionError>>
+ Send
+ Sync,
>;
pub struct ExtensionRegistry<E: Endpoint, R: RoleId> {
handlers: HashMap<TypeId, (ExtensionHandler<E, R>, &'static str)>,
}
impl<E: Endpoint, R: RoleId> ExtensionRegistry<E, R> {
#[must_use]
pub fn new() -> Self {
Self {
handlers: HashMap::new(),
}
}
pub fn register<Ext, F>(&mut self, handler: F) -> Result<(), ExtensionError>
where
Ext: ExtensionEffect<R> + 'static,
F: for<'a> Fn(
&'a mut E,
&'a dyn ExtensionEffect<R>,
) -> BoxFuture<'a, Result<(), ExtensionError>>
+ Send
+ Sync
+ 'static,
{
let type_id = TypeId::of::<Ext>();
let type_name = std::any::type_name::<Ext>();
if self.handlers.contains_key(&type_id) {
return Err(ExtensionError::DuplicateHandler { type_name });
}
self.handlers
.insert(type_id, (Box::new(handler), type_name));
Ok(())
}
#[must_use]
pub fn dispatch_contract() -> ExtensionDispatchContract {
ExtensionDispatchContract {
mode: ExtensionDispatchMode::RegisteredOnlyTypeExact,
fail_closed_when_unregistered: true,
type_exact_before_side_effects: true,
}
}
pub async fn handle(
&self,
endpoint: &mut E,
extension: &dyn ExtensionEffect<R>,
) -> Result<(), ExtensionError> {
let type_id = extension.type_id();
match self.handlers.get(&type_id) {
Some((handler, _type_name)) => handler(endpoint, extension).await,
None => Err(ExtensionError::HandlerNotRegistered {
type_name: extension.type_name(),
}),
}
}
#[must_use]
pub fn is_registered<Ext: ExtensionEffect<R> + 'static>(&self) -> bool {
self.handlers.contains_key(&TypeId::of::<Ext>())
}
#[must_use]
pub fn registered_handler_count(&self) -> usize {
self.handlers.len()
}
pub fn merge(&mut self, other: ExtensionRegistry<E, R>) -> Result<(), ExtensionError> {
for (type_id, (handler, type_name)) in other.handlers {
if self.handlers.contains_key(&type_id) {
return Err(ExtensionError::MergeConflict { type_name });
}
self.handlers.insert(type_id, (handler, type_name));
}
Ok(())
}
}
impl<E: Endpoint, R: RoleId> Default for ExtensionRegistry<E, R> {
fn default() -> Self {
Self::new()
}
}
pub trait ExtensibleHandler: crate::effects::ChoreoHandler {
fn extension_registry(&self) -> &ExtensionRegistry<Self::Endpoint, Self::Role>;
}