1use std::convert::Infallible;
5
6use serde::{Deserialize, Serialize};
7pub use vm_allocator::AllocPolicy;
8use vm_allocator::{AddressAllocator, IdAllocator};
9
10use crate::arch;
11use crate::snapshot::Persist;
12
13fn allocate_many_ids(
15 id_allocator: &mut IdAllocator,
16 count: u32,
17) -> Result<Vec<u32>, vm_allocator::Error> {
18 let mut ids = Vec::with_capacity(count as usize);
19
20 for _ in 0..count {
21 match id_allocator.allocate_id() {
22 Ok(id) => ids.push(id),
23 Err(err) => {
24 ids.into_iter().for_each(|id| {
26 id_allocator.free_id(id).unwrap();
27 });
28 return Err(err);
29 }
30 }
31 }
32
33 Ok(ids)
34}
35
36#[derive(Debug, Clone, Serialize, Deserialize)]
44pub struct ResourceAllocator {
45 pub gsi_legacy_allocator: IdAllocator,
47 pub gsi_msi_allocator: IdAllocator,
49 pub mmio32_memory: AddressAllocator,
51 pub mmio64_memory: AddressAllocator,
53 pub past_mmio64_memory: AddressAllocator,
55 pub system_memory: AddressAllocator,
57}
58
59impl Default for ResourceAllocator {
60 fn default() -> Self {
61 ResourceAllocator::new()
62 }
63}
64
65impl ResourceAllocator {
66 pub fn new() -> Self {
68 Self {
71 gsi_legacy_allocator: IdAllocator::new(arch::GSI_LEGACY_START, arch::GSI_LEGACY_END)
72 .unwrap(),
73 gsi_msi_allocator: IdAllocator::new(arch::GSI_MSI_START, arch::GSI_MSI_END).unwrap(),
74 mmio32_memory: AddressAllocator::new(
75 arch::MEM_32BIT_DEVICES_START,
76 arch::MEM_32BIT_DEVICES_SIZE,
77 )
78 .unwrap(),
79 mmio64_memory: AddressAllocator::new(
80 arch::MEM_64BIT_DEVICES_START,
81 arch::MEM_64BIT_DEVICES_SIZE,
82 )
83 .unwrap(),
84 past_mmio64_memory: AddressAllocator::new(
85 arch::FIRST_ADDR_PAST_64BITS_MMIO,
86 arch::PAST_64BITS_MMIO_SIZE,
87 )
88 .unwrap(),
89 system_memory: AddressAllocator::new(arch::SYSTEM_MEM_START, arch::SYSTEM_MEM_SIZE)
90 .unwrap(),
91 }
92 }
93
94 pub fn allocate_gsi_legacy(&mut self, gsi_count: u32) -> Result<Vec<u32>, vm_allocator::Error> {
100 allocate_many_ids(&mut self.gsi_legacy_allocator, gsi_count)
101 }
102
103 pub fn allocate_gsi_msi(&mut self, gsi_count: u32) -> Result<Vec<u32>, vm_allocator::Error> {
109 allocate_many_ids(&mut self.gsi_msi_allocator, gsi_count)
110 }
111
112 pub fn allocate_32bit_mmio_memory(
122 &mut self,
123 size: u64,
124 alignment: u64,
125 policy: AllocPolicy,
126 ) -> Result<u64, vm_allocator::Error> {
127 Ok(self
128 .mmio32_memory
129 .allocate(size, alignment, policy)?
130 .start())
131 }
132
133 pub fn allocate_64bit_mmio_memory(
143 &mut self,
144 size: u64,
145 alignment: u64,
146 policy: AllocPolicy,
147 ) -> Result<u64, vm_allocator::Error> {
148 Ok(self
149 .mmio64_memory
150 .allocate(size, alignment, policy)?
151 .start())
152 }
153
154 pub fn allocate_system_memory(
164 &mut self,
165 size: u64,
166 alignment: u64,
167 policy: AllocPolicy,
168 ) -> Result<u64, vm_allocator::Error> {
169 Ok(self
170 .system_memory
171 .allocate(size, alignment, policy)?
172 .start())
173 }
174}
175
176impl<'a> Persist<'a> for ResourceAllocator {
177 type State = ResourceAllocator;
178 type ConstructorArgs = ();
179 type Error = Infallible;
180
181 fn save(&self) -> Self::State {
182 self.clone()
183 }
184
185 fn restore(
186 _constructor_args: Self::ConstructorArgs,
187 state: &Self::State,
188 ) -> Result<Self, Self::Error> {
189 Ok(state.clone())
190 }
191}
192
193#[cfg(test)]
194mod tests {
195 use vm_allocator::AllocPolicy;
196
197 use super::ResourceAllocator;
198 use crate::arch::{self, GSI_LEGACY_NUM, GSI_LEGACY_START, GSI_MSI_NUM, GSI_MSI_START};
199 use crate::snapshot::{Persist, Snapshot};
200
201 #[test]
202 fn test_allocate_irq() {
203 let mut allocator = ResourceAllocator::new();
204 assert_eq!(allocator.allocate_gsi_legacy(0), Ok(vec![]));
206 assert_eq!(
208 allocator.allocate_gsi_legacy(GSI_LEGACY_NUM + 1),
209 Err(vm_allocator::Error::ResourceNotAvailable)
210 );
211 assert_eq!(
213 allocator.allocate_gsi_legacy(GSI_LEGACY_NUM),
214 Ok((arch::GSI_LEGACY_START..=arch::GSI_LEGACY_END).collect::<Vec<_>>())
215 );
216 assert_eq!(
218 allocator.allocate_gsi_legacy(1),
219 Err(vm_allocator::Error::ResourceNotAvailable)
220 );
221 assert_eq!(allocator.allocate_gsi_legacy(0), Ok(vec![]));
223
224 let mut allocator = ResourceAllocator::new();
225 assert_eq!(
227 allocator.allocate_gsi_legacy(1),
228 Ok(vec![arch::GSI_LEGACY_START])
229 );
230 assert_eq!(
232 allocator.allocate_gsi_legacy(GSI_LEGACY_NUM),
233 Err(vm_allocator::Error::ResourceNotAvailable)
234 );
235 assert_eq!(
237 allocator.allocate_gsi_legacy(1),
238 Ok(vec![arch::GSI_LEGACY_START + 1])
239 );
240 for i in arch::GSI_LEGACY_START + 2..=arch::GSI_LEGACY_END {
242 assert_eq!(allocator.allocate_gsi_legacy(1), Ok(vec![i]));
243 }
244 }
245
246 #[test]
247 fn test_allocate_gsi() {
248 let mut allocator = ResourceAllocator::new();
249 assert_eq!(allocator.allocate_gsi_msi(0), Ok(vec![]));
251 assert_eq!(
253 allocator.allocate_gsi_msi(GSI_MSI_NUM + 1),
254 Err(vm_allocator::Error::ResourceNotAvailable)
255 );
256 assert_eq!(
258 allocator.allocate_gsi_msi(GSI_MSI_NUM),
259 Ok((arch::GSI_MSI_START..=arch::GSI_MSI_END).collect::<Vec<_>>())
260 );
261 assert_eq!(
263 allocator.allocate_gsi_msi(1),
264 Err(vm_allocator::Error::ResourceNotAvailable)
265 );
266 assert_eq!(allocator.allocate_gsi_msi(0), Ok(vec![]));
268
269 let mut allocator = ResourceAllocator::new();
270 assert_eq!(allocator.allocate_gsi_msi(1), Ok(vec![arch::GSI_MSI_START]));
272 assert_eq!(
274 allocator.allocate_gsi_msi(GSI_MSI_NUM),
275 Err(vm_allocator::Error::ResourceNotAvailable)
276 );
277 assert_eq!(
279 allocator.allocate_gsi_msi(1),
280 Ok(vec![arch::GSI_MSI_START + 1])
281 );
282 for i in arch::GSI_MSI_START + 2..=arch::GSI_MSI_END {
284 assert_eq!(allocator.allocate_gsi_msi(1), Ok(vec![i]));
285 }
286 }
287
288 fn clone_allocator(allocator: &ResourceAllocator) -> ResourceAllocator {
289 let mut buf = vec![0u8; 1024];
290 Snapshot::new(allocator.save())
291 .save(&mut buf.as_mut_slice())
292 .unwrap();
293 let restored_state: ResourceAllocator = Snapshot::load_without_crc_check(buf.as_slice())
294 .unwrap()
295 .data;
296 ResourceAllocator::restore((), &restored_state).unwrap()
297 }
298
299 #[test]
300 fn test_save_restore() {
301 let mut allocator0 = ResourceAllocator::new();
302 let irq_0 = allocator0.allocate_gsi_legacy(1).unwrap()[0];
303 assert_eq!(irq_0, GSI_LEGACY_START);
304 let gsi_0 = allocator0.allocate_gsi_msi(1).unwrap()[0];
305 assert_eq!(gsi_0, GSI_MSI_START);
306
307 let mut allocator1 = clone_allocator(&allocator0);
308 let irq_1 = allocator1.allocate_gsi_legacy(1).unwrap()[0];
309 assert_eq!(irq_1, GSI_LEGACY_START + 1);
310 let gsi_1 = allocator1.allocate_gsi_msi(1).unwrap()[0];
311 assert_eq!(gsi_1, GSI_MSI_START + 1);
312 let mmio32_mem = allocator1
313 .allocate_32bit_mmio_memory(0x42, 1, AllocPolicy::FirstMatch)
314 .unwrap();
315 assert_eq!(mmio32_mem, arch::MEM_32BIT_DEVICES_START);
316 let mmio64_mem = allocator1
317 .allocate_64bit_mmio_memory(0x42, 1, AllocPolicy::FirstMatch)
318 .unwrap();
319 assert_eq!(mmio64_mem, arch::MEM_64BIT_DEVICES_START);
320 let system_mem = allocator1
321 .allocate_system_memory(0x42, 1, AllocPolicy::FirstMatch)
322 .unwrap();
323 assert_eq!(system_mem, arch::SYSTEM_MEM_START);
324
325 let mut allocator2 = clone_allocator(&allocator1);
326 allocator2
327 .allocate_32bit_mmio_memory(0x42, 1, AllocPolicy::ExactMatch(mmio32_mem))
328 .unwrap_err();
329 allocator2
330 .allocate_64bit_mmio_memory(0x42, 1, AllocPolicy::ExactMatch(mmio64_mem))
331 .unwrap_err();
332 allocator2
333 .allocate_system_memory(0x42, 1, AllocPolicy::ExactMatch(system_mem))
334 .unwrap_err();
335
336 let irq_2 = allocator2.allocate_gsi_legacy(1).unwrap()[0];
337 assert_eq!(irq_2, GSI_LEGACY_START + 2);
338 let gsi_2 = allocator2.allocate_gsi_msi(1).unwrap()[0];
339 assert_eq!(gsi_2, GSI_MSI_START + 2);
340 let mmio32_mem = allocator1
341 .allocate_32bit_mmio_memory(0x42, 1, AllocPolicy::FirstMatch)
342 .unwrap();
343 assert_eq!(mmio32_mem, arch::MEM_32BIT_DEVICES_START + 0x42);
344 let mmio64_mem = allocator1
345 .allocate_64bit_mmio_memory(0x42, 1, AllocPolicy::FirstMatch)
346 .unwrap();
347 assert_eq!(mmio64_mem, arch::MEM_64BIT_DEVICES_START + 0x42);
348 let system_mem = allocator1
349 .allocate_system_memory(0x42, 1, AllocPolicy::FirstMatch)
350 .unwrap();
351 assert_eq!(system_mem, arch::SYSTEM_MEM_START + 0x42);
352 }
353}