From e78c9e8879ebc579ea32f0c4582a4edf2b7a91aa Mon Sep 17 00:00:00 2001
From: walter <w41ter.l@gmail.com>
Date: Tue, 11 Oct 2022 16:52:25 +0800
Subject: [PATCH] Add WriteBatch::data method (#690)

---
 src/write_batch.rs        | 23 +++++++++++++++++++++
 tests/test_write_batch.rs | 43 ++++++++++++++++++++++++++++++++++++++-
 2 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/src/write_batch.rs b/src/write_batch.rs
index 683e457..5032274 100644
--- a/src/write_batch.rs
+++ b/src/write_batch.rs
@@ -90,6 +90,20 @@ unsafe extern "C" fn writebatch_delete_callback(state: *mut c_void, k: *const c_
 }
 
 impl<const TRANSACTION: bool> WriteBatchWithTransaction<TRANSACTION> {
+    /// Construct with a reference to a byte array serialized by [`WriteBatch`].
+    pub fn from_data(data: &[u8]) -> Self {
+        unsafe {
+            let ptr = data.as_ptr();
+            let len = data.len();
+            Self {
+                inner: ffi::rocksdb_writebatch_create_from(
+                    ptr as *const libc::c_char,
+                    len as size_t,
+                ),
+            }
+        }
+    }
+
     pub fn len(&self) -> usize {
         unsafe { ffi::rocksdb_writebatch_count(self.inner) as usize }
     }
@@ -103,6 +117,15 @@ impl<const TRANSACTION: bool> WriteBatchWithTransaction<TRANSACTION> {
         }
     }
 
+    /// Return a reference to a byte array which represents a serialized version of the batch.
+    pub fn data(&self) -> &[u8] {
+        unsafe {
+            let mut batch_size: size_t = 0;
+            let batch_data = ffi::rocksdb_writebatch_data(self.inner, &mut batch_size);
+            std::slice::from_raw_parts(batch_data as _, batch_size as usize)
+        }
+    }
+
     pub fn is_empty(&self) -> bool {
         self.len() == 0
     }
diff --git a/tests/test_write_batch.rs b/tests/test_write_batch.rs
index c1627aa..28ea265 100644
--- a/tests/test_write_batch.rs
+++ b/tests/test_write_batch.rs
@@ -12,9 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
+
 use pretty_assertions::assert_eq;
 
-use rocksdb::WriteBatch;
+use rocksdb::{WriteBatch, WriteBatchIterator};
 
 #[test]
 fn test_write_batch_clear() {
@@ -25,3 +27,42 @@ fn test_write_batch_clear() {
     assert_eq!(batch.len(), 0);
     assert!(batch.is_empty());
 }
+
+#[test]
+fn test_write_batch_with_serialized_data() {
+    struct Iterator {
+        data: HashMap<Vec<u8>, Vec<u8>>,
+    }
+
+    impl WriteBatchIterator for Iterator {
+        fn put(&mut self, key: Box<[u8]>, value: Box<[u8]>) {
+            match self.data.remove(key.as_ref()) {
+                Some(expect) => {
+                    assert_eq!(value.as_ref(), expect.as_slice());
+                }
+                None => {
+                    panic!("key not exists");
+                }
+            }
+        }
+
+        fn delete(&mut self, _: Box<[u8]>) {
+            panic!("invalid delete operation");
+        }
+    }
+
+    let mut kvs: HashMap<Vec<u8>, Vec<u8>> = HashMap::default();
+    kvs.insert(vec![1], vec![2]);
+    kvs.insert(vec![2], vec![3]);
+    kvs.insert(vec![1, 2, 3, 4, 5], vec![4]);
+
+    let mut b1 = WriteBatch::default();
+    for (k, v) in &kvs {
+        b1.put(k, v);
+    }
+    let data = b1.data();
+
+    let b2 = WriteBatch::from_data(data);
+    let mut it = Iterator { data: kvs };
+    b2.iterate(&mut it);
+}