1  use core::ops::{Deref, DerefMut};
 2  
 3  use alloc::vec::Vec;
 4  use spin::{Mutex, MutexGuard, RwLock};
 5  
 6  pub struct InstancedAlloc<T: Default + 'static> {
 7      inner: RwLock<InstancedAllocInner<T>>,
 8  }
 9  
10  pub struct InstancedAllocInner<T: Default + 'static> {
11      data: Vec<(bool, Mutex<T>)>,
12      next: usize,
13  }
14  
15  impl<T: Default> InstancedAlloc<T> {
16      pub fn with_size(size: usize) -> Self {
17          let mut data = Vec::with_capacity(size);
18  
19          for _ in 0..size {
20              data.push((false, Mutex::new(T::default())));
21          }
22  
23          Self {
24              inner: RwLock::new(InstancedAllocInner { data, next: 0 }),
25          }
26      }
27  
28      pub fn get(&'static self) -> InstanceHandle<T> {
29          let mut inner = self.inner.write();
30          let start = inner.next;
31          let mut looped = false;
32          let mut found = None;
33          loop {
34              if !looped && inner.next > inner.data.len() {
35                  inner.next = 0;
36                  looped = true;
37              } else if looped && inner.next > start {
38                  break;
39              }
40  
41              if !inner.data[inner.next].0 {
42                  found = Some(inner.next);
43              }
44  
45              inner.next += 1;
46  
47              if found.is_some() {
48                  break;
49              }
50          }
51  
52          inner.data[found.unwrap()].0 = true;
53  
54          InstanceHandle {
55              instances: &self.inner,
56              index: found.unwrap(),
57          }
58      }
59  }
60  
61  pub struct InstanceHandle<T: Default + 'static> {
62      instances: &'static RwLock<InstancedAllocInner<T>>,
63      index: usize,
64  }
65  
66  impl<T: Default + 'static> InstanceHandle<T> {
67      pub fn lock(&mut self) -> InstanceLock<T> {
68          // THIS SHOULD BE SAFE BUT I HATE IT
69          let instance = self.instances.read();
70          let bypassed_mutex = unsafe {
71              let mutex = &instance.data[self.index].1;
72              (mutex as *const Mutex<T> as *mut Mutex<T>)
73                  .as_mut()
74                  .unwrap()
75          };
76  
77          InstanceLock {
78              _parent: self,
79              lock: bypassed_mutex.lock(),
80          }
81      }
82  }
83  
84  pub struct InstanceLock<'a, T: Default + 'static> {
85      _parent: &'a InstanceHandle<T>,
86      lock: MutexGuard<'a, T>,
87  }
88  
89  impl<'a, T: Default + 'static> Deref for InstanceLock<'a, T> {
90      type Target = T;
91      fn deref(&self) -> &Self::Target {
92          &self.lock
93      }
94  }
95  impl<'a, T: Default + 'static> DerefMut for InstanceLock<'a, T> {
96      fn deref_mut(&mut self) -> &mut Self::Target {
97          &mut self.lock
98      }
99  }