Implement append and split_off for BitVec (RFC 509)

Author: jooert

src/libcollections/bit.rs

@@ -89,6 +89,7 @@ use core::hash;
 use core::iter::RandomAccessIterator;
 use core::iter::{Chain, Enumerate, Repeat, Skip, Take, repeat, Cloned};
 use core::iter::{self, FromIterator};
+use core::mem::swap;
 use core::ops::Index;
 use core::slice;
 use core::{u8, u32, usize};
@@ -604,6 +605,110 @@ impl BitVec {
         Iter { bit_vec: self, next_idx: 0, end_idx: self.nbits }
     }
 
+    /// Moves all bits from `other` into `Self`, leaving `other` empty.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #![feature(collections)]
+    /// use std::collections::BitVec;
+    ///
+    /// let mut a = BitVec::new();
+    /// a.push(true);
+    /// a.push(false);
+    ///
+    /// let mut b = BitVec::new();
+    /// b.push(false);
+    /// b.push(true);
+    /// b.push(true);
+    ///
+    /// a.append(&mut b);
+    ///
+    /// assert_eq!(a.len(), 5);
+    /// assert_eq!(b.len(), 0);
+    /// assert!(a.eq_vec(&[true, false, false, true, true]));
+    /// ```
+    #[unstable(feature = "collections",
+               reason = "recently added as part of collections reform 2")]
+    pub fn append(&mut self, other: &mut Self) {
+        let b = self.len() % u32::BITS;
+
+        self.nbits += other.len();
+        other.nbits = 0;
+
+        if b == 0 {
+            self.storage.append(&mut other.storage);
+        } else {
+            self.storage.reserve(other.storage.len());
+
+            for block in other.storage.drain() {
+                *(self.storage.last_mut().unwrap()) += block << b;
+                self.storage.push(block >> (u32::BITS - b));
+            }
+        }
+    }
+
+    /// Splits the `BitVec` into two at the given bit.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #![feature(collections)]
+    /// use std::collections::BitVec;
+    /// let mut a = BitVec::new();
+    /// a.push(true);
+    /// a.push(false);
+    /// a.push(false);
+    /// a.push(true);
+    ///
+    /// let b = a.split_off(2);
+    ///
+    /// assert_eq!(a.len(), 2);
+    /// assert_eq!(b.len(), 2);
+    /// assert!(a.eq_vec(&[true, false]));
+    /// assert!(b.eq_vec(&[false, true]));
+    /// ```
+    #[unstable(feature = "collections",
+               reason = "recently added as part of collections reform 2")]
+    pub fn split_off(&mut self, at: usize) -> Self {
+        let mut other = BitVec::new();
+
+        if at == 0 {
+            swap(self, &mut other);
+            return other;
+        } else if at > self.len() {
+            return other;
+        }
+
+        let w = at / u32::BITS;
+        let b = at % u32::BITS;
+        other.nbits = self.nbits - at;
+        self.nbits = at;
+        if b == 0 {
+            // Split at block boundary
+            other.storage = self.storage.split_off(w);
+            return other;
+        } else {
+            other.storage.reserve(self.storage.len() - w);
+
+            {
+                let mut cur = self.storage.iter().skip(w).peekable();
+                loop {
+                    match (cur.next(), cur.peek()) {
+                        (Some(x), Some(y)) => { other.storage.push((*x >> b) + (**y << (u32::BITS - b))) },
+                        (Some(x), None) => { other.storage.push((*x >> b)); break; },
+                        (None, Some(_)) => unreachable!(),
+                        (None, None) => unreachable!(),
+                    }
+                }
+            }
+
+            self.storage.truncate(w+1);
+        }
+
+        other
+    }
+
     /// Returns `true` if all bits are 0.
     ///
     /// # Examples

src/libcollectionstest/bit/vec.rs

@@ -630,6 +630,109 @@ fn test_bit_vec_extend() {
                                  0b01001001, 0b10010010, 0b10111101]));
 }
 
+#[test]
+fn test_bit_vec_append() {
+    // Append to BitVec, that holds a multiple of u32::BITS bits
+    let mut a = BitVec::from_bytes(&[0b10100000, 0b00010010, 0b10010010, 0b00110011]);
+    let mut b = BitVec::new();
+    b.push(false);
+    b.push(true);
+    b.push(true);
+
+    a.append(&mut b);
+
+    assert_eq!(a.len(), 35);
+    assert_eq!(b.len(), 0);
+    assert!(b.capacity() >= 3);
+
+    assert!(a.eq_vec(&[true, false, true, false, false, false, false, false,
+                       false, false, false, true, false, false, true, false,
+                       true, false, false, true, false, false, true, false,
+                       false, false, true, true, false, false, true, true,
+                       false, true, true]));
+
+    // Append to arbitrary BitVec
+    let mut a = BitVec::new();
+    a.push(true);
+    a.push(false);
+
+    let mut b = BitVec::from_bytes(&[0b10100000, 0b00010010, 0b10010010, 0b00110011, 0b10010101]);
+
+    a.append(&mut b);
+
+    assert_eq!(a.len(), 42);
+    assert_eq!(b.len(), 0);
+    assert!(b.capacity() >= 40);
+
+    assert!(a.eq_vec(&[true, false, true, false, true, false, false, false,
+                       false, false, false, false, false, true, false, false,
+                       true, false, true, false, false, true, false, false,
+                       true, false, false, false, true, true, false, false,
+                       true, true,  true, false, false, true, false, true,
+                       false, true]));
+}
+
+#[test]
+fn test_bit_vec_split_off() {
+    // Split at 0
+    let mut a = BitVec::new();
+    a.push(true);
+    a.push(false);
+    a.push(false);
+    a.push(true);
+
+    let b = a.split_off(0);
+
+    assert_eq!(a.len(), 0);
+    assert_eq!(b.len(), 4);
+
+    assert!(b.eq_vec(&[true, false, false, true]));
+
+    // Split behind last bit
+    a.truncate(0);
+    a.push(true);
+    a.push(false);
+    a.push(false);
+    a.push(true);
+
+    let b = a.split_off(5);
+
+    assert_eq!(a.len(), 4);
+    assert_eq!(b.len(), 0);
+
+    assert!(a.eq_vec(&[true, false, false, true]));
+
+    // Split at block boundary
+    let mut a = BitVec::from_bytes(&[0b10100000, 0b00010010, 0b10010010, 0b00110011, 0b11110011]);
+
+    let b = a.split_off(32);
+
+    assert_eq!(a.len(), 32);
+    assert_eq!(b.len(), 8);
+
+    assert!(a.eq_vec(&[true, false, true, false, false, false, false, false,
+                       false, false, false, true, false, false, true, false,
+                       true, false, false, true, false, false, true, false,
+                       false, false, true, true, false, false, true, true]));
+    assert!(b.eq_vec(&[true, true, true, true, false, false, true, true]));
+
+    // Don't split at block boundary
+    let mut a = BitVec::from_bytes(&[0b10100000, 0b00010010, 0b10010010, 0b00110011, 0b01101011, 0b10101101]);
+
+    let b = a.split_off(13);
+
+    assert_eq!(a.len(), 13);
+    assert_eq!(b.len(), 35);
+
+    assert!(a.eq_vec(&[true, false, true, false, false, false, false, false,
+                       false, false, false, true, false]));
+    assert!(b.eq_vec(&[false, true, false, true, false, false, true, false,
+                       false, true, false, false, false, true, true, false,
+                       false, true, true, false, true, true, false, true,
+                       false, true, true,  true, false, true, false, true,
+                       true, false, true]));
+}
+
 mod bench {
     use std::collections::BitVec;
     use std::u32;