use std::collections::HashMap;
use std::sync::RwLock;
use once_cell::sync::Lazy;
#[derive(Debug, Clone)]
pub struct Pass {
pub closure_id: u32,
pub args: Vec<u8>,
}
pub type PassResult = Result<Vec<u8>, PassError>;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PassError {
UnknownClosureId(u32),
ExecutionError(String),
}
pub type PassHandler = Box<dyn Fn(&[u8]) -> PassResult + Send + Sync + 'static>;
struct Registry {
handlers: HashMap<u32, PassHandler>,
}
static REGISTRY: Lazy<RwLock<Registry>> = Lazy::new(|| RwLock::new(Registry {
handlers: HashMap::new(),
}));
pub fn register<F>(id: u32, f: F) -> Option<PassHandler>
where F: Fn(&[u8]) -> PassResult + Send + Sync + 'static,
{
let mut g = REGISTRY.write().expect("registry write lock poisoned");
g.handlers.insert(id, Box::new(f))
}
pub fn unregister(id: u32) -> Option<PassHandler> {
let mut g = REGISTRY.write().expect("registry write lock poisoned");
g.handlers.remove(&id)
}
pub fn is_registered(id: u32) -> bool {
let g = REGISTRY.read().expect("registry read lock poisoned");
g.handlers.contains_key(&id)
}
pub fn registered_count() -> usize {
let g = REGISTRY.read().expect("registry read lock poisoned");
g.handlers.len()
}
pub fn execute(pass: &Pass) -> PassResult {
let g = REGISTRY.read().expect("registry read lock poisoned");
match g.handlers.get(&pass.closure_id) {
Some(handler) => handler(&pass.args),
None => Err(PassError::UnknownClosureId(pass.closure_id)),
}
}
#[macro_export]
macro_rules! register_pass {
($id:expr, $name:expr, $handler:expr) => {{
drop($crate::pass_registry::register($id, $handler));
let _name = $name; }};
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn register_then_execute_round_trips_bytes() {
let id = 0x1000_0001;
register(id, |args| {
let mut out = args.to_vec();
out.reverse();
Ok(out)
});
let pass = Pass { closure_id: id, args: b"hello".to_vec() };
let r = execute(&pass).unwrap();
assert_eq!(r, b"olleh");
unregister(id);
}
#[test]
fn unknown_closure_id_returns_error() {
let pass = Pass { closure_id: 0xDEAD_BEEF, args: vec![] };
assert_eq!(execute(&pass), Err(PassError::UnknownClosureId(0xDEAD_BEEF)));
}
#[test]
fn re_register_overwrites_previous() {
let id = 0x1000_0002;
register(id, |_| Ok(b"v1".to_vec()));
register(id, |_| Ok(b"v2".to_vec()));
let pass = Pass { closure_id: id, args: vec![] };
assert_eq!(execute(&pass).unwrap(), b"v2");
unregister(id);
}
#[test]
fn handler_can_return_execution_error() {
let id = 0x1000_0003;
register(id, |_| Err(PassError::ExecutionError("nope".to_string())));
let pass = Pass { closure_id: id, args: vec![] };
match execute(&pass) {
Err(PassError::ExecutionError(msg)) => assert_eq!(msg, "nope"),
other => panic!("expected ExecutionError, got {other:?}"),
}
unregister(id);
}
#[test]
fn is_registered_and_count_accurate() {
let id_a = 0x1000_0004;
let id_b = 0x1000_0005;
register(id_a, |_| Ok(vec![]));
register(id_b, |_| Ok(vec![]));
assert!(is_registered(id_a));
assert!(is_registered(id_b));
assert!(registered_count() >= 2);
unregister(id_a);
unregister(id_b);
assert!(!is_registered(id_a));
}
#[test]
fn macro_registration_works() {
let id = 0x1000_0006;
register_pass!(id, "test_pass", |args: &[u8]| {
Ok(args.iter().map(|b| b.wrapping_add(1)).collect())
});
assert!(is_registered(id));
let pass = Pass { closure_id: id, args: vec![1, 2, 3] };
assert_eq!(execute(&pass).unwrap(), vec![2, 3, 4]);
unregister(id);
}
}