@@ -121,7 +121,9 @@ impl<T: Idx> IdxSetBuf<T> {
121
121
122
122
/// Creates set holding every element whose index falls in range 0..universe_size.
123
123
pub fn new_filled ( universe_size : usize ) -> Self {
124
- Self :: new ( !0 , universe_size)
124
+ let mut result = Self :: new ( !0 , universe_size) ;
125
+ result. trim_to ( universe_size) ;
126
+ result
125
127
}
126
128
127
129
/// Creates set holding no elements.
@@ -168,6 +170,36 @@ impl<T: Idx> IdxSet<T> {
168
170
}
169
171
}
170
172
173
+ /// Sets all elements up to `universe_size`
174
+ pub fn set_up_to ( & mut self , universe_size : usize ) {
175
+ for b in & mut self . bits {
176
+ * b = !0 ;
177
+ }
178
+ self . trim_to ( universe_size) ;
179
+ }
180
+
181
+ /// Clear all elements above `universe_size`.
182
+ fn trim_to ( & mut self , universe_size : usize ) {
183
+ let word_bits = mem:: size_of :: < Word > ( ) * 8 ;
184
+
185
+ // `trim_block` is the first block where some bits have
186
+ // to be cleared.
187
+ let trim_block = universe_size / word_bits;
188
+
189
+ // all the blocks above it have to be completely cleared.
190
+ if trim_block < self . bits . len ( ) {
191
+ for b in & mut self . bits [ trim_block+1 ..] {
192
+ * b = 0 ;
193
+ }
194
+
195
+ // at that block, the `universe_size % word_bits` lsbs
196
+ // should remain.
197
+ let remaining_bits = universe_size % word_bits;
198
+ let mask = ( 1 <<remaining_bits) -1 ;
199
+ self . bits [ trim_block] &= mask;
200
+ }
201
+ }
202
+
171
203
/// Removes `elem` from the set `self`; returns true iff this changed `self`.
172
204
pub fn remove ( & mut self , elem : & T ) -> bool {
173
205
self . bits . clear_bit ( elem. index ( ) )
@@ -252,3 +284,43 @@ impl<'a, T: Idx> Iterator for Iter<'a, T> {
252
284
}
253
285
}
254
286
}
287
+
288
+ #[ test]
289
+ fn test_trim_to ( ) {
290
+ use std:: cmp;
291
+
292
+ for i in 0 ..256 {
293
+ let mut idx_buf: IdxSetBuf < usize > = IdxSetBuf :: new_filled ( 128 ) ;
294
+ idx_buf. trim_to ( i) ;
295
+
296
+ let elems: Vec < usize > = idx_buf. iter ( ) . collect ( ) ;
297
+ let expected: Vec < usize > = ( 0 ..cmp:: min ( i, 128 ) ) . collect ( ) ;
298
+ assert_eq ! ( elems, expected) ;
299
+ }
300
+ }
301
+
302
+ #[ test]
303
+ fn test_set_up_to ( ) {
304
+ for i in 0 ..128 {
305
+ for mut idx_buf in
306
+ vec ! [ IdxSetBuf :: new_empty( 128 ) , IdxSetBuf :: new_filled( 128 ) ]
307
+ . into_iter ( )
308
+ {
309
+ idx_buf. set_up_to ( i) ;
310
+
311
+ let elems: Vec < usize > = idx_buf. iter ( ) . collect ( ) ;
312
+ let expected: Vec < usize > = ( 0 ..i) . collect ( ) ;
313
+ assert_eq ! ( elems, expected) ;
314
+ }
315
+ }
316
+ }
317
+
318
+ #[ test]
319
+ fn test_new_filled ( ) {
320
+ for i in 0 ..128 {
321
+ let mut idx_buf = IdxSetBuf :: new_filled ( i) ;
322
+ let elems: Vec < usize > = idx_buf. iter ( ) . collect ( ) ;
323
+ let expected: Vec < usize > = ( 0 ..i) . collect ( ) ;
324
+ assert_eq ! ( elems, expected) ;
325
+ }
326
+ }
0 commit comments