nodedb_mem/
engine_limits.rs1use crate::engine::EngineId;
10
11#[derive(Debug, Clone)]
13pub struct EngineLimits([usize; EngineId::COUNT]);
14
15impl EngineLimits {
16 pub const fn zeroed() -> Self {
18 Self([0; EngineId::COUNT])
19 }
20
21 pub const fn uniform(bytes: usize) -> Self {
23 Self([bytes; EngineId::COUNT])
24 }
25
26 pub fn with(mut self, engine: EngineId, bytes: usize) -> Self {
28 self.set(engine, bytes);
29 self
30 }
31
32 pub fn set(&mut self, engine: EngineId, bytes: usize) {
34 self.0[engine.index()] = bytes;
35 }
36
37 pub fn get(&self, engine: EngineId) -> usize {
39 self.0[engine.index()]
40 }
41
42 pub fn total(&self) -> usize {
44 self.0.iter().sum()
45 }
46
47 pub fn iter(&self) -> impl Iterator<Item = (EngineId, usize)> + '_ {
49 EngineId::ALL
50 .iter()
51 .map(|&engine| (engine, self.get(engine)))
52 }
53
54 pub(crate) fn as_array(&self) -> &[usize; EngineId::COUNT] {
56 &self.0
57 }
58}
59
60impl FromIterator<(EngineId, usize)> for EngineLimits {
61 fn from_iter<T: IntoIterator<Item = (EngineId, usize)>>(iter: T) -> Self {
64 let mut limits = Self::zeroed();
65 for (engine, bytes) in iter {
66 limits.set(engine, bytes);
67 }
68 limits
69 }
70}
71
72#[cfg(test)]
73mod tests {
74 use super::*;
75
76 #[test]
77 fn zeroed_has_no_limit() {
78 let limits = EngineLimits::zeroed();
79 for &engine in EngineId::ALL {
80 assert_eq!(limits.get(engine), 0);
81 }
82 assert_eq!(limits.total(), 0);
83 }
84
85 #[test]
86 fn uniform_applies_to_every_engine() {
87 let limits = EngineLimits::uniform(1024);
88 for &engine in EngineId::ALL {
89 assert_eq!(limits.get(engine), 1024);
90 }
91 assert_eq!(limits.total(), 1024 * EngineId::COUNT);
92 }
93
94 #[test]
95 fn with_sets_a_single_engine() {
96 let limits = EngineLimits::zeroed()
97 .with(EngineId::Vector, 100)
98 .with(EngineId::Query, 200);
99 assert_eq!(limits.get(EngineId::Vector), 100);
100 assert_eq!(limits.get(EngineId::Query), 200);
101 assert_eq!(limits.get(EngineId::Kv), 0);
102 assert_eq!(limits.total(), 300);
103 }
104
105 #[test]
106 fn set_mutates_in_place() {
107 let mut limits = EngineLimits::zeroed();
108 limits.set(EngineId::Crdt, 42);
109 assert_eq!(limits.get(EngineId::Crdt), 42);
110 }
111
112 #[test]
113 fn iter_covers_every_engine_in_all_order() {
114 let limits = EngineLimits::zeroed().with(EngineId::Fts, 7);
115 let collected: Vec<_> = limits.iter().collect();
116 assert_eq!(collected.len(), EngineId::COUNT);
117 assert_eq!(
118 collected.iter().map(|(e, _)| *e).collect::<Vec<_>>(),
119 EngineId::ALL.to_vec()
120 );
121 assert_eq!(
122 collected
123 .iter()
124 .find(|(e, _)| *e == EngineId::Fts)
125 .unwrap()
126 .1,
127 7
128 );
129 }
130
131 #[test]
132 fn from_iter_leaves_unlisted_engines_zero() {
133 let limits: EngineLimits = [(EngineId::Vector, 10), (EngineId::Kv, 20)]
134 .into_iter()
135 .collect();
136 assert_eq!(limits.get(EngineId::Vector), 10);
137 assert_eq!(limits.get(EngineId::Kv), 20);
138 assert_eq!(limits.get(EngineId::Graph), 0);
139 }
140}