Merge pull request #81 from ltratt/double_ended

Double ended iteration for `iter_set_bits`.

Author: vext01

src/lib.rs

@@ -7,7 +7,7 @@ use std::{
     cmp::{min, PartialEq},
     fmt::{self, Debug},
     hash::{Hash, Hasher},
-    iter::FromIterator,
+    iter::{DoubleEndedIterator, FromIterator},
     mem::{replace, size_of},
     ops::{
         Bound::{Excluded, Included, Unbounded},
@@ -1101,6 +1101,10 @@ impl<'a, T: Debug + PrimInt> IntoIterator for &'a Vob<T> {
     }
 }
 
+// In all of the iterators below, we try to optimise what we expect to be the common cases (all
+// bits set or all bits unset) with the general case (random bit patterns). We assume that
+// `leading_zeros` and `trailing_zeros` have efficient implementations on most CPUs.
+
 #[derive(Clone)]
 pub struct IterSetBits<'a, T: 'a> {
     vob: &'a Vob<T>,
@@ -1113,24 +1117,15 @@ impl<T: Debug + PrimInt> Iterator for IterSetBits<'_, T> {
     fn next(&mut self) -> Option<usize> {
         debug_assert!(self.range.end <= self.vob.len);
         if let Some(i) = self.range.next() {
-            // This is a long, but fairly fast, way of finding out what the next set bit is. The
-            // basic problem is that we have no idea where the next set bit is -- but different
-            // patterns of set bits are most efficiently handled by different code paths. This code
-            // is thus a compromise: we prioritise two special cases (all bits set; all bits unset)
-            // for efficiency, and try and make the other possible cases reasonably fast.
             let mut b = block_offset::<T>(i);
             let mut v = self.vob.vec[b];
-            // If all bits are set, we don't need to do any complicated checks.
             if v == T::max_value() {
                 return Some(i);
             }
-            // At this point we've got a block which might or might not have some bits set. We now
-            // fall back to the general case.
             let mut i_off = i % bits_per_block::<T>();
             loop {
                 let tz = (v >> i_off).trailing_zeros() as usize;
                 if tz < bits_per_block::<T>() {
-                    // There is a bit set after i_off in the block.
                     let bs = b * bits_per_block::<T>() + i_off + tz;
                     self.range.start = bs + 1;
                     if bs >= self.range.end {
@@ -1140,7 +1135,6 @@ impl<T: Debug + PrimInt> Iterator for IterSetBits<'_, T> {
                 }
                 b += 1;
                 if b == blocks_required::<T>(self.range.end) {
-                    // We've exhausted the iterator.
                     self.range.start = self.range.end;
                     break;
                 }
@@ -1156,6 +1150,38 @@ impl<T: Debug + PrimInt> Iterator for IterSetBits<'_, T> {
     }
 }
 
+impl<T: Debug + PrimInt> DoubleEndedIterator for IterSetBits<'_, T> {
+    fn next_back(&mut self) -> Option<usize> {
+        if let Some(i) = self.range.next_back() {
+            let mut b = block_offset::<T>(i);
+            let mut v = self.vob.vec[b];
+            if v == T::max_value() {
+                return Some(i);
+            }
+            let mut i_off = i % bits_per_block::<T>();
+            loop {
+                let lz = (v << (bits_per_block::<T>() - 1 - i_off)).leading_zeros() as usize;
+                if lz < bits_per_block::<T>() {
+                    let bs = b * bits_per_block::<T>() + i_off - lz;
+                    self.range.end = bs;
+                    if bs < self.range.start {
+                        break;
+                    }
+                    return Some(bs);
+                }
+                if b == block_offset::<T>(self.range.start) {
+                    self.range.start = self.range.end;
+                    break;
+                }
+                b -= 1;
+                v = self.vob.vec[b];
+                i_off = bits_per_block::<T>() - 1;
+            }
+        }
+        None
+    }
+}
+
 #[derive(Clone)]
 pub struct IterUnsetBits<'a, T: 'a> {
     vob: &'a Vob<T>,
@@ -1168,19 +1194,11 @@ impl<T: Debug + PrimInt> Iterator for IterUnsetBits<'_, T> {
     fn next(&mut self) -> Option<usize> {
         debug_assert!(self.range.end <= self.vob.len);
         if let Some(i) = self.range.next() {
-            // This is a long, but fairly fast, way of finding out what the next usset bit is. The
-            // basic problem is that we have no idea where the next set bit is -- but different
-            // patterns of set bits are most efficiently handled by different code paths. This code
-            // is thus a compromise: we prioritise two special cases (all bits set; all bits unset)
-            // for efficiency, and try and make the other possible cases reasonably fast.
             let mut b = block_offset::<T>(i);
             let mut v = self.vob.vec[b];
-            // If no bits are set, we don't need to do any complicated checks.
             if v == T::zero() {
                 return Some(i);
             }
-            // At this point we've got a block which might or might not have some bits unset. We
-            // now fall back to the general case.
             let mut i_off = i % bits_per_block::<T>();
             loop {
                 let tz = (!v >> i_off).trailing_zeros() as usize;
@@ -1189,15 +1207,12 @@ impl<T: Debug + PrimInt> Iterator for IterUnsetBits<'_, T> {
                     let bs = b * bits_per_block::<T>() + i_off + tz;
                     self.range.start = bs + 1;
                     if bs >= self.range.end {
-                        // The unset bit is after the range we're looking for, so we've reached
-                        // the end of the iterator.
                         break;
                     }
                     return Some(bs);
                 }
                 b += 1;
                 if b == blocks_required::<T>(self.range.end) {
-                    // We've exhausted the iterator.
                     self.range.start = self.range.end;
                     break;
                 }
@@ -1542,6 +1557,7 @@ mod tests {
         hash::{Hash, Hasher},
         iter::FromIterator,
         mem::size_of,
+        ops::RangeBounds,
     };
 
     #[test]
@@ -1791,28 +1807,38 @@ mod tests {
 
     #[test]
     fn test_iter_set_bits() {
+        fn t<R>(v: &Vob, range: R, expected: Vec<usize>)
+        where
+            R: Clone + RangeBounds<usize>,
+        {
+            let rev = expected.iter().cloned().rev().collect::<Vec<usize>>();
+            assert_eq!(
+                v.iter_set_bits(range.clone()).collect::<Vec<usize>>(),
+                expected
+            );
+            assert_eq!(v.iter_set_bits(range).rev().collect::<Vec<usize>>(), rev);
+        }
+
+        t(&vob![], .., vec![]);
+        t(&Vob::from_elem(true, 131), .., (0..131).collect::<Vec<_>>());
+
         let mut v1 = vob![false, true, false, true];
-        assert_eq!(v1.iter_set_bits(..).collect::<Vec<usize>>(), vec![1, 3]);
+        t(&v1, .., vec![1, 3]);
         v1.resize(127, false);
         v1.push(true);
         v1.push(false);
         v1.push(true);
         v1.push(true);
         v1.resize(256, false);
         v1.push(true);
-        assert_eq!(
-            v1.iter_set_bits(..).collect::<Vec<usize>>(),
-            vec![1, 3, 127, 129, 130, 256]
-        );
-        assert_eq!(
-            v1.iter_set_bits(2..256).collect::<Vec<usize>>(),
-            vec![3, 127, 129, 130]
-        );
-        assert_eq!(
-            v1.iter_set_bits(2..).collect::<Vec<usize>>(),
-            vec![3, 127, 129, 130, 256]
-        );
-        assert_eq!(v1.iter_set_bits(..3).collect::<Vec<usize>>(), vec![1]);
+        assert_eq!(v1.len(), 257);
+        t(&v1, .., vec![1, 3, 127, 129, 130, 256]);
+        t(&v1, 2..257, vec![3, 127, 129, 130, 256]);
+        t(&v1, 2..256, vec![3, 127, 129, 130]);
+        t(&v1, 2.., vec![3, 127, 129, 130, 256]);
+        t(&v1, 1..255, vec![1, 3, 127, 129, 130]);
+        t(&v1, ..3, vec![1]);
+        t(&v1, 128.., vec![129, 130, 256]);
     }
 
     #[test]