@@ -29,6 +29,8 @@ pub struct AddressAllocator {
29
29
// tree will represent a memory location and can have two states either
30
30
// `NodeState::Free` or `NodeState::Allocated`.
31
31
interval_tree : IntervalTree ,
32
+ // Available memory space
33
+ available : usize ,
32
34
}
33
35
34
36
impl AddressAllocator {
@@ -43,6 +45,7 @@ impl AddressAllocator {
43
45
Ok ( AddressAllocator {
44
46
address_space : aux_range,
45
47
interval_tree : IntervalTree :: new ( aux_range) ,
48
+ available : aux_range. len ( ) as usize ,
46
49
} )
47
50
}
48
51
@@ -63,13 +66,22 @@ impl AddressAllocator {
63
66
policy : AllocPolicy ,
64
67
) -> Result < RangeInclusive > {
65
68
let constraint = Constraint :: new ( size, alignment, policy) ?;
66
- self . interval_tree . allocate ( constraint)
69
+ let allocated = self . interval_tree . allocate ( constraint) ?;
70
+ self . available -= allocated. len ( ) as usize ;
71
+ Ok ( allocated)
67
72
}
68
73
69
74
/// Deletes the specified memory slot or returns `ResourceNotAvailable` if
70
75
/// the node was not allocated before.
71
76
pub fn free ( & mut self , key : & RangeInclusive ) -> Result < ( ) > {
72
- self . interval_tree . free ( key)
77
+ self . interval_tree . free ( key) ?;
78
+ self . available += key. len ( ) as usize ;
79
+ Ok ( ( ) )
80
+ }
81
+
82
+ /// Returns the available memory size in this allocator.
83
+ pub fn available ( & self ) -> usize {
84
+ self . available
73
85
}
74
86
}
75
87
@@ -158,20 +170,27 @@ mod tests {
158
170
#[ test]
159
171
fn test_allocate_with_alignment_first_ok ( ) {
160
172
let mut pool = AddressAllocator :: new ( 0x1000 , 0x1000 ) . unwrap ( ) ;
173
+ assert_eq ! ( pool. available( ) , 0x1000 ) ;
174
+ // Allocate 0x110
161
175
assert_eq ! (
162
176
pool. allocate( 0x110 , 0x100 , AllocPolicy :: FirstMatch )
163
177
. unwrap( ) ,
164
178
RangeInclusive :: new( 0x1000 , 0x110F ) . unwrap( )
165
179
) ;
180
+ assert_eq ! ( pool. available( ) , 0x1000 - 0x110 ) ;
181
+ // Allocate 0x100
166
182
assert_eq ! (
167
183
pool. allocate( 0x100 , 0x100 , AllocPolicy :: FirstMatch )
168
184
. unwrap( ) ,
169
185
RangeInclusive :: new( 0x1200 , 0x12FF ) . unwrap( )
170
186
) ;
187
+ assert_eq ! ( pool. available( ) , 0x1000 - 0x110 - 0x100 ) ;
188
+ // Allocate 0x10
171
189
assert_eq ! (
172
190
pool. allocate( 0x10 , 0x100 , AllocPolicy :: FirstMatch ) . unwrap( ) ,
173
191
RangeInclusive :: new( 0x1300 , 0x130F ) . unwrap( )
174
192
) ;
193
+ assert_eq ! ( pool. available( ) , 0x1000 - 0x110 - 0x100 - 0x10 ) ;
175
194
}
176
195
177
196
#[ test]
@@ -230,18 +249,24 @@ mod tests {
230
249
#[ test]
231
250
fn test_tree_allocate_address_free_and_realloc ( ) {
232
251
let mut pool = AddressAllocator :: new ( 0x1000 , 0x1000 ) . unwrap ( ) ;
252
+ assert_eq ! ( pool. available( ) , 0x1000 ) ;
253
+ // Allocate 0x800
233
254
assert_eq ! (
234
255
pool. allocate( 0x800 , 0x100 , AllocPolicy :: FirstMatch )
235
256
. unwrap( ) ,
236
257
RangeInclusive :: new( 0x1000 , 0x17FF ) . unwrap( )
237
258
) ;
238
-
259
+ assert_eq ! ( pool. available( ) , 0x1000 - 0x800 ) ;
260
+ // Free 0x800
239
261
let _ = pool. free ( & RangeInclusive :: new ( 0x1000 , 0x17FF ) . unwrap ( ) ) ;
262
+ assert_eq ! ( pool. available( ) , 0x1000 ) ;
263
+ // Allocate 0x800 again
240
264
assert_eq ! (
241
265
pool. allocate( 0x800 , 0x100 , AllocPolicy :: FirstMatch )
242
266
. unwrap( ) ,
243
267
RangeInclusive :: new( 0x1000 , 0x17FF ) . unwrap( )
244
268
) ;
269
+ assert_eq ! ( pool. available( ) , 0x1000 - 0x800 ) ;
245
270
}
246
271
247
272
#[ test]
0 commit comments