@@ -121,7 +121,9 @@ impl<T: Idx> IdxSetBuf<T> {
121121
122122 /// Creates set holding every element whose index falls in range 0..universe_size.
123123 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
125127 }
126128
127129 /// Creates set holding no elements.
@@ -168,6 +170,36 @@ impl<T: Idx> IdxSet<T> {
168170 }
169171 }
170172
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+
171203 /// Removes `elem` from the set `self`; returns true iff this changed `self`.
172204 pub fn remove ( & mut self , elem : & T ) -> bool {
173205 self . bits . clear_bit ( elem. index ( ) )
@@ -252,3 +284,43 @@ impl<'a, T: Idx> Iterator for Iter<'a, T> {
252284 }
253285 }
254286}
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