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,105 @@ 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::from_bytes(&[0b10000000]);
+    /// let mut b = BitVec::from_bytes(&[0b01100001]);
+    ///
+    /// a.append(&mut b);
+    ///
+    /// assert_eq!(a.len(), 16);
+    /// assert_eq!(b.len(), 0);
+    /// assert!(a.eq_vec(&[true, false, false, false, false, false, false, false,
+    ///                    false, true, true, false, false, false, false, true]));
+    /// ```
+    #[unstable(feature = "bit_vec_append_split_off",
+               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 = "bit_vec_append_split_off",
+               reason = "recently added as part of collections reform 2")]
+    pub fn split_off(&mut self, at: usize) -> Self {
+        assert!(at <= self.len(), "`at` out of bounds");
+
+        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);
+        } else {
+            other.storage.reserve(self.storage.len() - w);
+
+            {
+                let mut iter = self.storage[w..].iter();
+                let mut last = *iter.next().unwrap();
+                for &cur in iter {
+                    other.storage.push((last >> b) | (cur << (u32::BITS - b)));
+                    last = cur;
+                }
+                other.storage.push(last >> b);
+            }
+
+            self.storage.truncate(w+1);
+            self.fix_last_block();
+        }
+
+        other
+    }
+
     /// Returns `true` if all bits are 0.
     ///
     /// # Examples

src/libcollectionstest/bit/vec.rs

@@ -630,6 +630,141 @@ 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]));
+
+    // Append to empty BitVec
+    let mut a = BitVec::new();
+    let mut b = BitVec::from_bytes(&[0b10100000, 0b00010010, 0b10010010, 0b00110011, 0b10010101]);
+
+    a.append(&mut b);
+
+    assert_eq!(a.len(), 40);
+    assert_eq!(b.len(), 0);
+    assert!(b.capacity() >= 40);
+
+    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,
+                       true, false, false, true, false, true, false, true]));
+
+    // Append empty BitVec
+    let mut a = BitVec::from_bytes(&[0b10100000, 0b00010010, 0b10010010, 0b00110011, 0b10010101]);
+    let mut b = BitVec::new();
+
+    a.append(&mut b);
+
+    assert_eq!(a.len(), 40);
+    assert_eq!(b.len(), 0);
+
+    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,
+                       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 at last bit
+    a.truncate(0);
+    a.push(true);
+    a.push(false);
+    a.push(false);
+    a.push(true);
+
+    let b = a.split_off(4);
+
+    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;