xref: /DragonOS/kernel/crates/ida/src/lib.rs (revision 4dd4856f933be0b4624c7f7ffa9e3d0c8c218873)
1 #![no_std]
2 #![feature(core_intrinsics)]
3 #![allow(internal_features)]
4 #![allow(clippy::needless_return)]
5 
6 #[cfg(test)]
7 #[macro_use]
8 extern crate std;
9 
10 use core::cmp::min;
11 use core::intrinsics::unlikely;
12 use core::marker::PhantomData;
13 use core::ops::Deref;
14 
15 struct EmptyIdaItemRef<'a> {
16     _marker: PhantomData<&'a EmptyIdaItem>,
17 }
18 
19 impl<'a> Deref for EmptyIdaItemRef<'a> {
20     type Target = EmptyIdaItem;
21 
22     fn deref(&self) -> &Self::Target {
23         &EmptyIdaItem
24     }
25 }
26 
27 struct EmptyIdaItem;
28 
29 unsafe impl kdepends::xarray::ItemEntry for EmptyIdaItem {
30     type Ref<'a> = EmptyIdaItemRef<'a> where Self: 'a;
31 
32     fn into_raw(self) -> *const () {
33         core::ptr::null()
34     }
35 
36     unsafe fn from_raw(_raw: *const ()) -> Self {
37         EmptyIdaItem
38     }
39 
40     unsafe fn raw_as_ref<'a>(_raw: *const ()) -> Self::Ref<'a> {
41         EmptyIdaItemRef {
42             _marker: PhantomData,
43         }
44     }
45 }
46 /// id分配器
47 pub struct IdAllocator {
48     current_id: usize,
49     min_id: usize,
50     max_id: usize,
51     used: usize,
52     xarray: kdepends::xarray::XArray<EmptyIdaItem>,
53 }
54 
55 impl IdAllocator {
56     /// 创建一个新的id分配器
57     pub const fn new(initial_id: usize, max_id: usize) -> Option<Self> {
58         if initial_id >= max_id {
59             return None;
60         }
61         Some(Self {
62             current_id: initial_id,
63             min_id: initial_id,
64             max_id,
65             used: 0,
66             xarray: kdepends::xarray::XArray::new(),
67         })
68     }
69 
70     /// 可用的id数量
71     #[inline]
72     pub fn available(&self) -> usize {
73         self.max_id - self.min_id - self.used
74     }
75 
76     /// 分配一个新的id
77     ///
78     /// ## 返回
79     ///
80     /// 如果分配成功,返回Some(id),否则返回None
81     pub fn alloc(&mut self) -> Option<usize> {
82         if unlikely(self.available() == 0) {
83             return None;
84         }
85 
86         if let Some(try1) = self.do_find_first_free_index(self.current_id, self.max_id) {
87             self.current_id = try1;
88             self.xarray.store(try1 as u64, EmptyIdaItem);
89             self.used += 1;
90             return Some(try1);
91         }
92 
93         // 从头开始找
94         if let Some(try2) =
95             self.do_find_first_free_index(self.min_id, min(self.current_id, self.max_id))
96         {
97             self.current_id = try2;
98             self.xarray.store(try2 as u64, EmptyIdaItem);
99             self.used += 1;
100             return Some(try2);
101         }
102         return None;
103     }
104 
105     /// 检查id是否存在
106     ///
107     /// ## 参数
108     ///
109     /// - `id`:要检查的id
110     ///
111     /// ## 返回
112     ///
113     /// 如果id存在,返回true,否则返回false
114     pub fn exists(&self, id: usize) -> bool {
115         if id < self.min_id || id >= self.max_id {
116             return false;
117         }
118         self.xarray.load(id as u64).is_some()
119     }
120 
121     fn do_find_first_free_index(&self, start_id: usize, end: usize) -> Option<usize> {
122         (start_id..end).find(|&i| !self.exists(i))
123     }
124 
125     /// 释放一个id
126     ///
127     /// ## 参数
128     ///
129     /// - `id`:要释放的id
130     pub fn free(&mut self, id: usize) {
131         if id < self.min_id || id >= self.max_id {
132             return;
133         }
134         if self.xarray.remove(id as u64).is_some() {
135             self.used -= 1;
136         }
137     }
138 
139     /// 返回已经使用的id数量
140     pub fn used(&self) -> usize {
141         self.used
142     }
143 }
144 
145 impl core::fmt::Debug for IdAllocator {
146     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
147         f.debug_struct("IdAllocator")
148             .field("current_id", &self.current_id)
149             .field("min_id", &self.min_id)
150             .field("max_id", &self.max_id)
151             .field("used", &self.used)
152             .field("xarray", &"xarray<()>")
153             .finish()
154     }
155 }
156 
157 #[cfg(test)]
158 mod tests {
159     use super::*;
160     #[test]
161     fn test_new_fail() {
162         assert_eq!(IdAllocator::new(10, 10).is_none(), true);
163         assert_eq!(IdAllocator::new(11, 10).is_none(), true);
164     }
165     #[test]
166     fn test_new_success() {
167         assert_eq!(IdAllocator::new(9, 10).is_some(), true);
168         assert_eq!(IdAllocator::new(0, 10).is_some(), true);
169     }
170 
171     #[test]
172     fn test_id_allocator() {
173         let mut ida = IdAllocator::new(0, 10).unwrap();
174         assert_eq!(ida.alloc(), Some(0));
175         assert_eq!(ida.alloc(), Some(1));
176         assert_eq!(ida.alloc(), Some(2));
177         assert_eq!(ida.alloc(), Some(3));
178         assert_eq!(ida.alloc(), Some(4));
179         assert_eq!(ida.alloc(), Some(5));
180         assert_eq!(ida.alloc(), Some(6));
181         assert_eq!(ida.alloc(), Some(7));
182         assert_eq!(ida.alloc(), Some(8));
183         assert_eq!(ida.alloc(), Some(9));
184         assert_eq!(ida.alloc(), None);
185 
186         for i in 0..10 {
187             assert_eq!(ida.exists(i), true);
188         }
189 
190         ida.free(5);
191 
192         for i in 0..10 {
193             if i == 5 {
194                 assert_eq!(ida.exists(i), false);
195             } else {
196                 assert_eq!(ida.exists(i), true);
197             }
198         }
199         assert_eq!(ida.used(), 9);
200         assert_eq!(ida.alloc(), Some(5));
201         assert_eq!(ida.alloc(), None);
202 
203         assert_eq!(ida.used(), 10);
204         for i in 0..10 {
205             ida.free(i);
206         }
207 
208         assert_eq!(ida.used(), 0);
209     }
210 }
211