acceptor_alloc/btree_kv_router/
btree_kv_router.rs

1use alloc::collections::BTreeMap;
2use core::marker::PhantomData;
3
4use accepts::{Accepts, AsyncAccepts};
5
6/// Routes `(Key, Value)` pairs to acceptors stored in a `BTreeMap`.
7///
8/// If no route exists for the provided key, the fallback acceptor receives the
9/// original pair. Downstream acceptors also receive the original `(Key, Value)`
10/// so they can keep using the key.
11#[must_use = "BTreeKVRouter must be used to route values"]
12#[derive(Debug, Clone)]
13pub struct BTreeKVRouter<Key, Value, RouteAccepts, Fallback = ()> {
14    routes: BTreeMap<Key, RouteAccepts>,
15    fallback: Fallback,
16    _marker: PhantomData<Value>,
17}
18
19impl<Key, Value, RouteAccepts> BTreeKVRouter<Key, Value, RouteAccepts, ()> {
20    /// Creates a `BTreeKVRouter` that drops pairs when no route matches.
21    pub fn new(routes: BTreeMap<Key, RouteAccepts>) -> Self {
22        Self::with_fallback(routes, ())
23    }
24}
25
26impl<Key, Value, RouteAccepts, Fallback> BTreeKVRouter<Key, Value, RouteAccepts, Fallback> {
27    /// Creates a `BTreeKVRouter` with a custom fallback acceptor.
28    pub fn with_fallback(routes: BTreeMap<Key, RouteAccepts>, fallback: Fallback) -> Self {
29        Self {
30            routes,
31            fallback,
32            _marker: PhantomData,
33        }
34    }
35
36    pub fn routes(&self) -> &BTreeMap<Key, RouteAccepts> {
37        &self.routes
38    }
39
40    pub fn routes_mut(&mut self) -> &mut BTreeMap<Key, RouteAccepts> {
41        &mut self.routes
42    }
43
44    pub fn fallback(&self) -> &Fallback {
45        &self.fallback
46    }
47
48    pub fn fallback_mut(&mut self) -> &mut Fallback {
49        &mut self.fallback
50    }
51}
52
53impl<Key, Value, RouteAccepts, Fallback> Accepts<(Key, Value)>
54    for BTreeKVRouter<Key, Value, RouteAccepts, Fallback>
55where
56    Key: Ord,
57    RouteAccepts: Accepts<(Key, Value)>,
58    Fallback: Accepts<(Key, Value)>,
59{
60    fn accept(&self, pair: (Key, Value)) {
61        let (key, value) = pair;
62
63        if let Some(route) = self.routes.get(&key) {
64            route.accept((key, value));
65            return;
66        }
67
68        self.fallback.accept((key, value));
69    }
70}
71
72impl<Key, Value, RouteAccepts, Fallback> AsyncAccepts<(Key, Value)>
73    for BTreeKVRouter<Key, Value, RouteAccepts, Fallback>
74where
75    Key: Ord,
76    RouteAccepts: AsyncAccepts<(Key, Value)>,
77    Fallback: AsyncAccepts<(Key, Value)>,
78{
79    fn accept_async<'a>(&'a self, pair: (Key, Value)) -> impl core::future::Future<Output = ()> + 'a
80    where
81        Key: 'a,
82        Value: 'a,
83    {
84        async move {
85            let (key, value) = pair;
86
87            if let Some(route) = self.routes.get(&key) {
88                route.accept_async((key, value)).await;
89                return;
90            }
91
92            self.fallback.accept_async((key, value)).await;
93        }
94    }
95}
acceptor_alloc/btree_kv_router/btree_kv_router.rs

acceptor_alloc/btree_kv_router/
btree_kv_router.rs
