Export ToBytes & FromBytes for allowing persist custom struct

Signed-off-by: kexuan.yang <[email protected]>
yangkx-1024 · Oct 21, 2024 · 266bbe9 · 266bbe9
1 parent 7328534
commit 266bbe9
Show file tree

Hide file tree

Showing 8 changed files with 323 additions and 170 deletions.
diff --git a/src/core/buffer.rs b/src/core/buffer.rs
diff --git a/src/core/crc.rs b/src/core/crc.rs
@@ -55,7 +55,7 @@ mod tests {
 
     #[test]
     fn test_crc_buffer() {
-        let buffer = Buffer::encode("key", 1);
+        let buffer = Buffer::new("key", 1);
         let bytes = CrcEncoderDecoder.encode_to_bytes(&buffer, 0).unwrap();
         let decode_result = CrcEncoderDecoder.decode_bytes(bytes.as_slice(), 0).unwrap();
         assert_eq!(decode_result.len, bytes.len() as u32);

diff --git a/src/core/encrypt.rs b/src/core/encrypt.rs
@@ -178,12 +178,12 @@ mod tests {
     fn test_crypt_buffer() {
         let path = Path::new("./mmkv");
         let encryptor = Encryptor::init(path, TEST_KEY);
-        let buffer1 = Buffer::encode("key1", 1);
+        let buffer1 = Buffer::new("key1", 1);
         let bytes1 = encryptor.encode_to_bytes(&buffer1, 0).unwrap();
         let decode_result1 = encryptor.decode_bytes(bytes1.as_slice(), 0).unwrap();
         assert_eq!(decode_result1.len, bytes1.len() as u32);
         assert_eq!(decode_result1.buffer, Some(buffer1.clone()));
-        let buffer2 = Buffer::encode("key2", 2);
+        let buffer2 = Buffer::new("key2", 2);
         let bytes2 = encryptor.encode_to_bytes(&buffer2, 1).unwrap();
         let decode_result2 = encryptor.decode_bytes(bytes2.as_slice(), 1).unwrap();
         assert_eq!(decode_result2.len, bytes2.len() as u32);

diff --git a/src/core/iter.rs b/src/core/iter.rs
@@ -140,7 +140,7 @@ mod tests {
         let mut buffers: Vec<Buffer> = vec![];
         let test_encoder = &TestEncoderDecoder;
         for i in 0..10 {
-            let buffer = Buffer::encode(&i.to_string(), i);
+            let buffer = Buffer::new(&i.to_string(), i);
             mm.append(test_encoder.encode_to_bytes(&buffer, i as u32).unwrap());
             buffers.push(buffer);
         }

diff --git a/src/core/mmkv_impl.rs b/src/core/mmkv_impl.rs
@@ -162,35 +162,35 @@ mod tests {
         let config = &Config::new(Path::new(file_path), 100);
         let mm = MemoryMap::new(&config.file, 200);
         let mut mmkv = init(config);
-        mmkv.put("key1", Buffer::encode("key1", 1)).unwrap(); // + 17
+        mmkv.put("key1", Buffer::new("key1", 1)).unwrap(); // + 17
         drop(mmkv);
         assert_eq!(mm.write_offset(), 25);
 
         mmkv = init(config);
-        mmkv.put("key2", Buffer::encode("key2", 2)).unwrap(); // + 17
-        mmkv.put("key3", Buffer::encode("key3", 3)).unwrap(); // + 17
-        mmkv.put("key1", Buffer::encode("key1", 4)).unwrap(); // + 17
-        mmkv.put("key2", Buffer::encode("key2", 5)).unwrap(); // + 17
+        mmkv.put("key2", Buffer::new("key2", 2)).unwrap(); // + 17
+        mmkv.put("key3", Buffer::new("key3", 3)).unwrap(); // + 17
+        mmkv.put("key1", Buffer::new("key1", 4)).unwrap(); // + 17
+        mmkv.put("key2", Buffer::new("key2", 5)).unwrap(); // + 17
         drop(mmkv);
         assert_eq!(mm.write_offset(), 93);
 
         mmkv = init(config);
-        mmkv.put("key1", Buffer::encode("key1", 6)).unwrap(); // + 17, trim, 3 items remain
+        mmkv.put("key1", Buffer::new("key1", 6)).unwrap(); // + 17, trim, 3 items remain
         drop(mmkv);
         assert_eq!(mm.write_offset(), 59);
 
         mmkv = init(config);
-        assert_eq!(mmkv.get("key1").unwrap().decode::<i32>(), Ok(6));
-        assert_eq!(mmkv.get("key2").unwrap().decode::<i32>(), Ok(5));
-        mmkv.put("key4", Buffer::encode("key4", 4)).unwrap();
-        mmkv.put("key5", Buffer::encode("key5", 5)).unwrap(); // 93
-        mmkv.put("key6", Buffer::encode("key6", 6)).unwrap(); // expand, 110
+        assert_eq!(mmkv.get("key1").unwrap().parse::<i32>(), Ok(6));
+        assert_eq!(mmkv.get("key2").unwrap().parse::<i32>(), Ok(5));
+        mmkv.put("key4", Buffer::new("key4", 4)).unwrap();
+        mmkv.put("key5", Buffer::new("key5", 5)).unwrap(); // 93
+        mmkv.put("key6", Buffer::new("key6", 6)).unwrap(); // expand, 110
         drop(mmkv);
         assert_eq!(mm.write_offset(), 110);
         assert_eq!(config.file_size(), 200);
 
         mmkv = init(config);
-        mmkv.put("key7", Buffer::encode("key7", 7)).unwrap();
+        mmkv.put("key7", Buffer::new("key7", 7)).unwrap();
         drop(mmkv);
         assert_eq!(mm.write_offset(), 127);
 
@@ -208,32 +208,32 @@ mod tests {
         let config = &Config::new(Path::new(file), 100);
         let mm = MemoryMap::new(&config.file, 200);
         let mut mmkv = init(config);
-        mmkv.put("key1", Buffer::encode("key1", 1)).unwrap(); // + 24
+        mmkv.put("key1", Buffer::new("key1", 1)).unwrap(); // + 24
         drop(mmkv);
         assert_eq!(mm.write_offset(), 32);
 
         mmkv = init(config);
-        mmkv.put("key2", Buffer::encode("key2", 2)).unwrap(); // + 24
-        mmkv.put("key3", Buffer::encode("key3", 3)).unwrap(); // + 24
+        mmkv.put("key2", Buffer::new("key2", 2)).unwrap(); // + 24
+        mmkv.put("key3", Buffer::new("key3", 3)).unwrap(); // + 24
         drop(mmkv);
         assert_eq!(mm.write_offset(), 80);
 
         mmkv = init(config);
-        mmkv.put("key1", Buffer::encode("key1", 4)).unwrap(); // + 24 trim
-        mmkv.put("key2", Buffer::encode("key2", 5)).unwrap(); // + 24 trim
+        mmkv.put("key1", Buffer::new("key1", 4)).unwrap(); // + 24 trim
+        mmkv.put("key2", Buffer::new("key2", 5)).unwrap(); // + 24 trim
         drop(mmkv);
         assert_eq!(mm.write_offset(), 80);
 
         mmkv = init(config);
-        assert_eq!(mmkv.get("key1").unwrap().decode::<i32>(), Ok(4));
-        assert_eq!(mmkv.get("key2").unwrap().decode::<i32>(), Ok(5));
-        mmkv.put("key4", Buffer::encode("key4", 4)).unwrap(); // + 24
+        assert_eq!(mmkv.get("key1").unwrap().parse::<i32>(), Ok(4));
+        assert_eq!(mmkv.get("key2").unwrap().parse::<i32>(), Ok(5));
+        mmkv.put("key4", Buffer::new("key4", 4)).unwrap(); // + 24
         drop(mmkv);
         assert_eq!(mm.write_offset(), 104);
         assert_eq!(config.file_size(), 200);
 
         mmkv = init(config);
-        mmkv.put("key5", Buffer::encode("key5", 5)).unwrap(); // + 24
+        mmkv.put("key5", Buffer::new("key5", 5)).unwrap(); // + 24
         drop(mmkv);
         assert_eq!(mm.write_offset(), 128);
 
@@ -257,7 +257,7 @@ mod tests {
                     .unwrap()
                     .as_mut()
                     .unwrap()
-                    .put(key, Buffer::encode(key, i))
+                    .put(key, Buffer::new(key, i))
                     .unwrap();
             }
         };
@@ -268,7 +268,7 @@ mod tests {
                     let mut lock = mmkv.write().unwrap();
                     let mmkv = lock.as_mut().unwrap();
                     if i % 2 == 0 {
-                        mmkv.put(repeat_key, Buffer::encode(repeat_key, i)).unwrap();
+                        mmkv.put(repeat_key, Buffer::new(repeat_key, i)).unwrap();
                     } else {
                         mmkv.delete(repeat_key).unwrap();
                     }
@@ -283,7 +283,7 @@ mod tests {
         for i in 0..2 {
             for j in 0..loop_count {
                 let key = &format!("thread_{i}_key_{j}");
-                assert_eq!(mmkv.get(key).unwrap().decode::<i32>().unwrap(), j)
+                assert_eq!(mmkv.get(key).unwrap().parse::<i32>().unwrap(), j)
             }
         }
         assert_eq!(

diff --git a/src/lib.rs b/src/lib.rs
@@ -11,6 +11,7 @@
 //! mmkv.clear_data();
 //! ```
 //! For detailed API doc, see [MMKV]
+pub use crate::core::buffer::{FromBytes, ProvideTypeToken, ToBytes, TypeToken};
 pub use crate::log::LogLevel;
 pub use crate::log::Logger;
 pub use crate::mmkv::MMKV;

diff --git a/src/mmkv.rs b/src/mmkv.rs
@@ -5,7 +5,7 @@ use std::sync::{Arc, RwLock, Weak};
 
 use once_cell::sync::Lazy;
 
-use crate::core::buffer::{Buffer, FromBuffer, ToBuffer};
+use crate::core::buffer::{Buffer, FromBytes, ProvideTypeToken, ToBytes};
 use crate::core::config::Config;
 use crate::core::mmkv_impl::MmkvImpl;
 use crate::log::logger;
@@ -125,16 +125,70 @@ impl MMKV {
         path.join(DEFAULT_FILE_NAME)
     }
 
-    pub fn put<T: ToBuffer>(&self, key: &str, value: T) -> Result<()> {
+    /**
+    Types must implement [ProvideTypeToken] and [ToBytes] to be persisted in MMKV.
+
+    If you want to persist custom struct to MMKV,
+    your struct must implement [ToBytes] trait which serialize type to bytes,
+    and [FromBytes] trait which deserialize type from bytes.
+    For example:
+    ```
+    use mmkv::{FromBytes, MMKV, ProvideTypeToken, ToBytes, TypeToken};
+    #[derive(Clone, PartialEq, Eq, Debug, Hash)]
+    struct MyStruct {
+        int_value: i32,
+        str_value: String,
+    }
+
+    impl ProvideTypeToken for MyStruct {
+        fn type_token() -> TypeToken {
+            TypeToken::new(101)
+        }
+    }
+
+    impl ToBytes for &MyStruct {
+        fn to_bytes(&self) -> Vec<u8> {
+            let mut vec = vec![];
+            vec.extend(self.int_value.to_be_bytes());
+            vec.extend(self.str_value.as_bytes().to_vec());
+            vec
+        }
+    }
+
+    impl FromBytes for MyStruct {
+        fn from_bytes(bytes: &[u8]) -> mmkv::Result<Self> {
+            let int_len = size_of::<i32>();
+            let int_val = i32::from_be_bytes(bytes[0..int_len].try_into().unwrap());
+            let str_val = String::from_utf8(bytes[int_len..].try_into().unwrap()).unwrap();
+            Ok(MyStruct {
+                int_value: int_val,
+                str_value: str_val
+            })
+        }
+    }
+
+    let temp_dir = std::env::temp_dir();
+    let mmkv = MMKV::new(temp_dir.to_str().unwrap(), #[cfg(feature = "encryption")] "88C51C536176AD8A8EE4A06F62EE897E");
+    let my_struct = MyStruct {
+        int_value: 1,
+        str_value: "abc".to_string(),
+    };
+    mmkv.put("my_struct", &my_struct).unwrap();
+    let copy: MyStruct = mmkv.get("my_struct").unwrap();
+    assert_eq!(my_struct, copy)
+    ```
+    */
+    pub fn put<T: ProvideTypeToken + ToBytes>(&self, key: &str, value: T) -> Result<()> {
         match self.mmkv_impl.write() {
-            Ok(mut mmkv) => mmkv.put(key, Buffer::encode(key, value)),
+            Ok(mut mmkv) => mmkv.put(key, Buffer::new(key, value)),
             Err(e) => Err(LockError(e.to_string())),
         }
     }
 
-    pub fn get<T: FromBuffer>(&self, key: &str) -> Result<T> {
+    /// See [MMKV::put]
+    pub fn get<T: ProvideTypeToken + FromBytes>(&self, key: &str) -> Result<T> {
         match self.mmkv_impl.read() {
-            Ok(mmkv) => mmkv.get(key)?.decode(),
+            Ok(mmkv) => mmkv.get(key)?.parse(),
             Err(e) => Err(LockError(e.to_string())),
         }
     }

diff --git a/src/protos/kv.proto b/src/protos/kv.proto
@@ -1,22 +1,7 @@
 syntax = "proto3";
 
-enum Types {
-  I32 = 0;
-  STR = 1;
-  BYTE = 2;
-  I64 = 3;
-  F32 = 4;
-  F64 = 5;
-  BYTE_ARRAY = 6;
-  I32_ARRAY = 7;
-  I64_ARRAY = 8;
-  F32_ARRAY = 9;
-  F64_ARRAY = 10;
-  DELETED = 100;
-}
-
 message KV {
   string key = 1;
-  Types type = 2;
+  int32 type = 2;
   bytes value = 3;
 }