fedimint_core/module/
registry.rs

1use std::collections::BTreeMap;
2
3pub use crate::core::ModuleInstanceId;
4use crate::core::{Decoder, ModuleKind};
5
6/// Module Registry hold module-specific data `M` by the `ModuleInstanceId`
7#[derive(Debug)]
8pub struct ModuleRegistry<M, State = ()> {
9    inner: BTreeMap<ModuleInstanceId, (ModuleKind, M)>,
10    // It is sometimes useful for registries to have some state to modify
11    // their behavior.
12    state: State,
13}
14
15impl<M, State> Clone for ModuleRegistry<M, State>
16where
17    State: Clone,
18    M: Clone,
19{
20    fn clone(&self) -> Self {
21        Self {
22            inner: self.inner.clone(),
23            state: self.state.clone(),
24        }
25    }
26}
27
28impl<M, State> Default for ModuleRegistry<M, State>
29where
30    State: Default,
31{
32    fn default() -> Self {
33        Self {
34            inner: BTreeMap::new(),
35            state: State::default(),
36        }
37    }
38}
39
40impl<M, State> From<BTreeMap<ModuleInstanceId, (ModuleKind, M)>> for ModuleRegistry<M, State>
41where
42    State: Default,
43{
44    fn from(value: BTreeMap<ModuleInstanceId, (ModuleKind, M)>) -> Self {
45        Self {
46            inner: value,
47            state: Default::default(),
48        }
49    }
50}
51
52impl<M, State> FromIterator<(ModuleInstanceId, ModuleKind, M)> for ModuleRegistry<M, State>
53where
54    State: Default,
55{
56    fn from_iter<T: IntoIterator<Item = (ModuleInstanceId, ModuleKind, M)>>(iter: T) -> Self {
57        Self::new(iter)
58    }
59}
60
61impl<M, State> ModuleRegistry<M, State> {
62    /// Create [`Self`] from an iterator of pairs
63    pub fn new(iter: impl IntoIterator<Item = (ModuleInstanceId, ModuleKind, M)>) -> Self
64    where
65        State: Default,
66    {
67        Self {
68            inner: iter
69                .into_iter()
70                .map(|(id, kind, module)| (id, (kind, module)))
71                .collect(),
72            state: Default::default(),
73        }
74    }
75
76    /// Is registry empty?
77    pub fn is_empty(&self) -> bool {
78        self.inner.is_empty()
79    }
80
81    /// Return an iterator over all module data
82    pub fn iter_modules(&self) -> impl Iterator<Item = (ModuleInstanceId, &ModuleKind, &M)> {
83        self.inner.iter().map(|(id, (kind, m))| (*id, kind, m))
84    }
85
86    /// Return an iterator over module ids an kinds
87    pub fn iter_modules_id_kind(&self) -> impl Iterator<Item = (ModuleInstanceId, &ModuleKind)> {
88        self.inner.iter().map(|(id, (kind, _))| (*id, kind))
89    }
90
91    /// Return an iterator over all module data
92    pub fn iter_modules_mut(
93        &mut self,
94    ) -> impl Iterator<Item = (ModuleInstanceId, &ModuleKind, &mut M)> {
95        self.inner
96            .iter_mut()
97            .map(|(id, (kind, m))| (*id, &*kind, m))
98    }
99
100    /// Return an iterator over all module data
101    pub fn into_iter_modules(self) -> impl Iterator<Item = (ModuleInstanceId, ModuleKind, M)> {
102        self.inner.into_iter().map(|(id, (kind, m))| (id, kind, m))
103    }
104
105    /// Get module data by instance id
106    pub fn get(&self, id: ModuleInstanceId) -> Option<&M> {
107        self.inner.get(&id).map(|m| &m.1)
108    }
109
110    /// Get module data by instance id, including [`ModuleKind`]
111    pub fn get_with_kind(&self, id: ModuleInstanceId) -> Option<&(ModuleKind, M)> {
112        self.inner.get(&id)
113    }
114}
115
116impl<M: std::fmt::Debug, State> ModuleRegistry<M, State> {
117    /// Return the module data belonging to the module identified by the
118    /// supplied `module_id`
119    ///
120    /// # Panics
121    /// If the module isn't in the registry
122    pub fn get_expect(&self, id: ModuleInstanceId) -> &M {
123        let Some((_, module)) = self.inner.get(&id) else {
124            panic!(
125                "Only existing instance should be fetched: Instance ID not found: got {id}, \
126                 expected one of {:?}",
127                self.inner.keys().collect::<Vec<_>>()
128            );
129        };
130        module
131    }
132
133    /// Add a module to the registry
134    pub fn register_module(&mut self, id: ModuleInstanceId, kind: ModuleKind, module: M) {
135        // FIXME: return result
136        assert!(
137            self.inner.insert(id, (kind, module)).is_none(),
138            "Module was already registered!"
139        );
140    }
141
142    pub fn append_module(&mut self, kind: ModuleKind, module: M) {
143        let last_id = self
144            .inner
145            .last_key_value()
146            .map(|id| id.0.checked_add(1).expect("Module id overflow"))
147            .unwrap_or_default();
148        assert!(
149            self.inner.insert(last_id, (kind, module)).is_none(),
150            "Module was already registered?!"
151        );
152    }
153}
154
155#[derive(Default, Clone, Copy, PartialEq, Eq, Hash, Debug)]
156pub enum DecodingMode {
157    /// Reject unknown module instance ids
158    #[default]
159    Reject,
160    /// Fallback to decoding unknown module instance ids as
161    /// [`crate::core::DynUnknown`]
162    Fallback,
163}
164
165/// Collection of decoders belonging to modules, typically obtained from a
166/// `ModuleRegistry`
167pub type ModuleDecoderRegistry = ModuleRegistry<Decoder, DecodingMode>;
168
169impl ModuleDecoderRegistry {
170    pub fn with_fallback(self) -> Self {
171        Self {
172            state: DecodingMode::Fallback,
173            ..self
174        }
175    }
176
177    pub fn decoding_mode(&self) -> DecodingMode {
178        self.state
179    }
180
181    /// Panic if the [`Self::decoding_mode`] is not `Reject`
182    pub fn assert_reject_mode(&self) {
183        assert_eq!(self.state, DecodingMode::Reject);
184    }
185}
fedimint_core/module/registry.rs

fedimint_core/module/
registry.rs
