use std::{collections::HashMap, sync::Arc};
use crate::{
algorithm::{Algorithm, AlgorithmConfig},
feed::events::MarketEventHandler,
graph::EdgeWeightUpdaterWithDerived,
worker_pool::{
pool::WorkerPoolBuilder,
registry::{UnknownAlgorithmError, AVAILABLE_ALGORITHMS},
},
};
type Configure = Arc<dyn Fn(WorkerPoolBuilder) -> WorkerPoolBuilder + Send + Sync>;
#[derive(Debug, Clone, thiserror::Error, PartialEq, Eq)]
pub enum RegisterAlgorithmError {
#[error("'{name}' is a built-in algorithm; registering it would replace the shipped one")]
ShadowsBuiltIn {
name: String,
},
#[error("'{name}' is already registered")]
AlreadyRegistered {
name: String,
},
}
#[derive(Default, Clone)]
pub struct AlgorithmRegistry {
by_name: HashMap<String, Configure>,
}
impl AlgorithmRegistry {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn with_algorithm<A, F>(
mut self,
name: impl Into<String>,
factory: F,
) -> Result<Self, RegisterAlgorithmError>
where
A: Algorithm + 'static,
A::GraphManager: MarketEventHandler + EdgeWeightUpdaterWithDerived + 'static,
F: Fn(AlgorithmConfig) -> A + Clone + Send + Sync + 'static,
{
let name = name.into();
if AVAILABLE_ALGORITHMS.contains(&name.as_str()) {
return Err(RegisterAlgorithmError::ShadowsBuiltIn { name });
}
if self.by_name.contains_key(&name) {
return Err(RegisterAlgorithmError::AlreadyRegistered { name });
}
let registered = name.clone();
self.by_name.insert(
name,
Arc::new(move |builder: WorkerPoolBuilder| {
builder.with_algorithm(registered.clone(), factory.clone())
}),
);
Ok(self)
}
pub fn names(&self) -> impl Iterator<Item = &str> {
self.by_name.keys().map(String::as_str)
}
pub(crate) fn configure(
&self,
name: &str,
builder: WorkerPoolBuilder,
) -> Result<WorkerPoolBuilder, UnknownAlgorithmError> {
if let Some(configure) = self.by_name.get(name) {
return Ok(configure(builder));
}
if AVAILABLE_ALGORITHMS.contains(&name) {
return Ok(builder.algorithm(name));
}
Err(UnknownAlgorithmError::of(
name,
self.names()
.map(str::to_string)
.collect(),
))
}
}
impl std::fmt::Debug for AlgorithmRegistry {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AlgorithmRegistry")
.field("names", &self.by_name.keys().collect::<Vec<_>>())
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::algorithm::{most_liquid::MostLiquidAlgorithm, AlgorithmConfig};
fn most_liquid(config: AlgorithmConfig) -> MostLiquidAlgorithm {
MostLiquidAlgorithm::with_config(config).expect("the default config is valid")
}
fn registry_with(name: &str) -> AlgorithmRegistry {
AlgorithmRegistry::new()
.with_algorithm(name, most_liquid)
.expect("the name is neither built in nor taken")
}
#[test]
fn test_registry_is_empty_by_default() {
assert_eq!(AlgorithmRegistry::new().names().count(), 0);
}
#[test]
fn test_registry_refuses_a_name_already_registered() {
let error = registry_with("mine")
.with_algorithm("mine", most_liquid)
.expect_err("the name is taken");
assert_eq!(error, RegisterAlgorithmError::AlreadyRegistered { name: "mine".to_string() });
}
#[test]
fn test_registry_refuses_a_built_in_name() {
let error = AlgorithmRegistry::new()
.with_algorithm("most_liquid", most_liquid)
.expect_err("the name ships with this crate");
assert_eq!(
error,
RegisterAlgorithmError::ShadowsBuiltIn { name: "most_liquid".to_string() }
);
}
#[test]
fn test_configure_serves_a_registered_name() {
let Ok(builder) = registry_with("brought_from_outside")
.configure("brought_from_outside", WorkerPoolBuilder::new())
else {
panic!("a registered name is served");
};
assert!(builder.serves_custom_algorithm(), "a registered name is served by its factory");
}
#[test]
fn test_configure_leaves_a_built_in_name_to_the_built_in() {
let Ok(builder) = registry_with("mine").configure("water_fill", WorkerPoolBuilder::new())
else {
panic!("a built-in name is served");
};
assert!(
!builder.serves_custom_algorithm(),
"a built-in name is not served from the registry"
);
}
#[test]
fn test_configure_rejects_a_name_neither_side_holds() {
let Err(error) = registry_with("brought_from_outside")
.configure("brought_from_outsid", WorkerPoolBuilder::new())
else {
panic!("a name neither side holds must be refused");
};
let message = error.to_string();
assert!(message.contains("brought_from_outside"), "lists the registered name: {message}");
assert!(message.contains("water_fill"), "lists the built-ins too: {message}");
}
}