fedimint_core/module/
registry.rs

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

fedimint_core/module/
registry.rs
