add set_raw and set_raw_copy

GCCFeli · Jan 25, 2022 · 1c209b6 · 1c209b6
1 parent a452211
commit 1c209b6
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 10 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -2,7 +2,7 @@
 name = "bitvec_simd"
 description = "bitvec implemented with wide"
 authors = ["GCCFeli"]
-version = "0.12.0"
+version = "0.13.0"
 edition = "2021"
 homepage = "https://github.com/gccfeli/bitvec_simd"
 license = "MIT"

diff --git a/benches/comparison.rs b/benches/comparison.rs
@@ -97,8 +97,8 @@ fn benchmark_bitvector_bitvec2(c: &mut Criterion) {
 }
 
 fn benchmark_bitvector_bitvec_n(c: &mut Criterion) {
-    let b1 = bitvec::bitvec![bitvec::order::Msb0, usize; 1; 100_000];
-    let b2 = bitvec::bitvec![bitvec::order::Msb0, usize; 0; 100_000];
+    let b1 = bitvec::bitvec![usize, bitvec::order::Msb0; 1; 100_000];
+    let b2 = bitvec::bitvec![usize, bitvec::order::Msb0; 0; 100_000];
     c.bench_function("bitvec 0.22", |b| {
         b.iter(|| {
             black_box(b1.clone() & b2.clone());
@@ -109,8 +109,8 @@ fn benchmark_bitvector_bitvec_n(c: &mut Criterion) {
 fn benchmark_bitvector_bitvec_n2(c: &mut Criterion) {
     c.bench_function("bitvec 0.22 with creation", |b| {
         b.iter(|| {
-            let b1 = bitvec::bitvec![bitvec::order::Msb0, usize; 1; 100_000];
-            let b2 = bitvec::bitvec![bitvec::order::Msb0, usize; 0; 100_000];
+            let b1 = bitvec::bitvec![usize, bitvec::order::Msb0; 1; 100_000];
+            let b2 = bitvec::bitvec![usize, bitvec::order::Msb0; 0; 100_000];
             black_box(b1 & b2);
         })
     });
@@ -119,7 +119,7 @@ fn benchmark_bitvector_bitvec_n2(c: &mut Criterion) {
 fn benchmark_bitvector_bitvec_n3(c: &mut Criterion) {
     c.bench_function("bitvec 0.22 resize false", |b| {
         b.iter(|| {
-            let mut b1 = bitvec::bitvec![bitvec::order::Msb0, usize; 1; 100_000];
+            let mut b1 = bitvec::bitvec![usize, bitvec::order::Msb0; 1; 100_000];
             black_box(b1.resize(200_000, false));
         })
     });
@@ -128,7 +128,7 @@ fn benchmark_bitvector_bitvec_n3(c: &mut Criterion) {
 fn benchmark_bitvector_bitvec_n4(c: &mut Criterion) {
     c.bench_function("bitvec 0.22 resize true", |b| {
         b.iter(|| {
-            let mut b1 = bitvec::bitvec![bitvec::order::Msb0, usize; 1; 100_000];
+            let mut b1 = bitvec::bitvec![usize, bitvec::order::Msb0; 1; 100_000];
             black_box(b1.resize(200_000, true));
         })
     });

diff --git a/src/lib.rs b/src/lib.rs
@@ -483,7 +483,6 @@ where
         }
     }
 
-
     /// Length of this bitvec.
     ///
     /// To get the number of elements, use `count_ones`
@@ -622,16 +621,39 @@ where
     pub fn set(&mut self, index: usize, flag: bool) {
         let (i, bytes, bits) = Self::bit_to_len(index);
         if self.nbits <= index {
-            let i = if bytes > 0 || bits > 0 { i + 1 } else { i };
+            let new_len = if bytes > 0 || bits > 0 { i + 1 } else { i };
             self.storage
-                .extend((0..i - self.storage.len()).map(move |_| <<A as Array>::Item as BitContainer<L>>::ZERO));
+                .extend((0..new_len - self.storage.len()).map(move |_| <<A as Array>::Item as BitContainer<L>>::ZERO));
             self.nbits = index + 1;
         }
         let mut arr = self.storage[i].to_array();
         arr[bytes] = Self::set_bit(flag, arr[bytes], bits as u32);
         self.storage[i] = <A as Array>::Item::from(arr);
     }
 
+    /// Copy content which ptr points to bitvec storage
+    /// Highly unsafe
+    pub unsafe fn set_raw_copy(&mut self, ptr: *mut A::Item, buffer_len: usize, nbits: usize) {
+        let new_len = (nbits + <A as Array>::Item::BIT_WIDTH - 1) / <A as Array>::Item::BIT_WIDTH;
+        assert!(new_len <= buffer_len);
+
+        if new_len > self.len() {
+            self.storage.extend((0..new_len - self.storage.len()).map(move |_| A::Item::ZERO));
+        }
+
+        for i in 0..(new_len as isize) {
+            self.storage[i as usize] = *ptr.offset(i);
+        }
+        self.nbits = nbits;
+    }
+
+    /// Directly set storage to ptr
+    /// Highly unsafe
+    pub unsafe fn set_raw(&mut self, ptr: *mut A::Item, buffer_len: usize, capacity: usize, nbits: usize) {
+        self.storage = SmallVec::from_raw_parts(ptr, buffer_len, capacity);
+        self.nbits = nbits;
+    }
+
     /// Set all items in bitvec to false
     pub fn set_all_false(&mut self) {
         self.storage.iter_mut().for_each(move |x| *x = <<A as Array>::Item as BitContainer<L>>::ZERO);
@@ -1440,6 +1462,43 @@ fn test_bitvec_creation() {
     assert_eq!(bitvec.get(64), None);
 }
 
+#[test]
+fn test_bitvec_set_raw_copy() {
+    let v = vec![7];
+    let buf = v.as_ptr();
+    let mut bitvec = unsafe { BitVec::from_raw_copy(buf, 1, 64) };
+    let ptr = bitvec.storage.as_mut_ptr();
+    let buffer_len = bitvec.storage.len();
+    let mut bitvec2 = BitVec::zeros(1);
+    unsafe { bitvec2.set_raw_copy(ptr, buffer_len, bitvec.nbits); }
+    assert_eq!(v.len(), 1); // ensure v lives long enough
+    assert_eq!(bitvec2.get(0), Some(true));
+    assert_eq!(bitvec2.get(1), Some(true));
+    assert_eq!(bitvec2.get(2), Some(true));
+    assert_eq!(bitvec2.get(3), Some(false));
+    assert_eq!(bitvec2.get(63), Some(false));
+    assert_eq!(bitvec2.get(64), None);
+}
+
+#[test]
+fn test_bitvec_set_raw() {
+    let mut bitvec = BitVec::ones(1025);
+    let ptr = bitvec.storage.as_mut_ptr();
+    let buffer_len = bitvec.storage.len();
+    let capacity = bitvec.storage.capacity();
+    let nbits = bitvec.nbits;
+    let spilled = bitvec.storage.spilled();
+    let mut bitvec2 = BitVec::zeros(1);
+    unsafe {
+        std::mem::forget(bitvec);
+        assert!(spilled);
+        bitvec2.set_raw(ptr, buffer_len, capacity, nbits);
+        assert_eq!(bitvec2.get(0), Some(true));
+        assert_eq!(bitvec2.get(1024), Some(true));
+        assert_eq!(bitvec2.get(1025), None);
+    }
+}
+
 #[test]
 fn test_bitvec_set_all() {
     let mut bitvec = BitVec::zeros(1000);