From 624b09ecd9719175a2c5285347b04d7c5c3009b5 Mon Sep 17 00:00:00 2001
From: librelois <c@elo.tf>
Date: Sat, 5 Jun 2021 18:26:33 +0200
Subject: [PATCH] fix(kv_typed): deser empty BTreeSet

---
 tools/kv_typed/src/as_bytes.rs         | 14 ++++++++
 tools/kv_typed/src/from_bytes.rs       | 46 +++++++++++++++++++++-----
 tools/kv_typed/tests/test_db_schema.rs |  5 +++
 3 files changed, 57 insertions(+), 8 deletions(-)

diff --git a/tools/kv_typed/src/as_bytes.rs b/tools/kv_typed/src/as_bytes.rs
index 46e0e3f..4e79a4f 100644
--- a/tools/kv_typed/src/as_bytes.rs
+++ b/tools/kv_typed/src/as_bytes.rs
@@ -61,6 +61,20 @@ where
     }
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_btreeset_as_bytes() {
+        BTreeSet::<u64>::new().as_bytes(|bytes| assert_eq!(bytes, &[]));
+    }
+    #[test]
+    fn test_hashset_as_bytes() {
+        HashSet::<u64>::new().as_bytes(|bytes| assert_eq!(bytes, &[]));
+    }
+}
+
 macro_rules! impl_as_bytes_for_le_numbers {
     ($($T:ty),*) => {$(
         impl AsBytes for $T {
diff --git a/tools/kv_typed/src/from_bytes.rs b/tools/kv_typed/src/from_bytes.rs
index 138f9ff..5ae7245 100644
--- a/tools/kv_typed/src/from_bytes.rs
+++ b/tools/kv_typed/src/from_bytes.rs
@@ -97,10 +97,14 @@ where
     type Err = LayoutVerifiedErr;
 
     fn from_bytes(bytes: &[u8]) -> Result<Self, Self::Err> {
-        let layout_verified = zerocopy::LayoutVerified::<_, [T]>::new_slice(bytes)
-            .ok_or(LayoutVerifiedErr(stringify!(BTreeSet<T>)))?;
-        let slice = layout_verified.into_slice();
-        Ok(slice.iter().copied().collect())
+        if bytes.is_empty() {
+            Ok(Self::new())
+        } else {
+            let layout_verified = zerocopy::LayoutVerified::<_, [T]>::new_slice(bytes)
+                .ok_or(LayoutVerifiedErr(stringify!(BTreeSet<T>)))?;
+            let slice = layout_verified.into_slice();
+            Ok(slice.iter().copied().collect())
+        }
     }
 }
 
@@ -111,10 +115,14 @@ where
     type Err = LayoutVerifiedErr;
 
     fn from_bytes(bytes: &[u8]) -> Result<Self, Self::Err> {
-        let layout_verified = zerocopy::LayoutVerified::<_, [T]>::new_slice(bytes)
-            .ok_or(LayoutVerifiedErr(stringify!(HashSet<T>)))?;
-        let slice = layout_verified.into_slice();
-        Ok(slice.iter().copied().collect())
+        if bytes.is_empty() {
+            Ok(Self::new())
+        } else {
+            let layout_verified = zerocopy::LayoutVerified::<_, [T]>::new_slice(bytes)
+                .ok_or(LayoutVerifiedErr(stringify!(HashSet<T>)))?;
+            let slice = layout_verified.into_slice();
+            Ok(slice.iter().copied().collect())
+        }
     }
 }
 
@@ -125,3 +133,25 @@ impl FromBytes for IVec {
         Ok(Self::from(bytes))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_btreeset_from_bytes() -> Result<(), LayoutVerifiedErr> {
+        assert_eq!(
+            <BTreeSet<u64> as FromBytes>::from_bytes(&[])?,
+            BTreeSet::<u64>::new()
+        );
+        Ok(())
+    }
+    #[test]
+    fn test_hashset_from_bytes() -> Result<(), LayoutVerifiedErr> {
+        assert_eq!(
+            <HashSet<u64> as FromBytes>::from_bytes(&[])?,
+            HashSet::<u64>::new()
+        );
+        Ok(())
+    }
+}
diff --git a/tools/kv_typed/tests/test_db_schema.rs b/tools/kv_typed/tests/test_db_schema.rs
index c312ce2..dead76e 100644
--- a/tools/kv_typed/tests/test_db_schema.rs
+++ b/tools/kv_typed/tests/test_db_schema.rs
@@ -116,6 +116,11 @@ fn test_db<B: Backend>(db: &TestV1Db<B>) -> KvResult<()> {
     })?;
 
     // Test get_ref_slice
+    db.col4_write().upsert(U64BE(3), BTreeSet::new())?;
+    db.col4().get_ref_slice(&U64BE(3), |numbers| {
+        assert_eq!(numbers, &[]);
+        Ok(())
+    })?;
     db.col4_write()
         .upsert(U64BE(4), (&[3, 2, 4, 1]).iter().copied().collect())?;
     db.col4().get_ref_slice(&U64BE(4), |numbers| {
-- 
GitLab