Custom serialization refactor

storage/src/lib.rs

@@ -2,7 +2,6 @@ mod memory_backend;
 mod storage;
 
 mod rc_key;
-mod serde_rc;
 mod sled_backend;
 mod storage_backend;
 mod storage_entity;

storage/src/rc_key.rs

@@ -9,6 +9,7 @@ use std::rc::Rc;
  * using RcKey, we still effectively use the pointer as the key, but we also hold a reference to the
  * rc, so it cannot be dropped and reused, which would create a false association.
  */
+#[derive(Clone)]
 pub struct RcKey(Rc<dyn Any>);
 
 impl RcKey {

storage/src/serde_rc.rs

@@ -1,17 +0,0 @@
-use std::rc::Rc;
-
-use serde::{Deserialize, Deserializer, Serialize, Serializer};
-
-pub fn serialize_rc<T: Serialize, S: Serializer>(
-  data: &Rc<T>,
-  serializer: S,
-) -> Result<S::Ok, S::Error> {
-  data.serialize(serializer)
-}
-
-pub fn deserialize_rc<'de, T: Deserialize<'de>, D: Deserializer<'de>>(
-  deserializer: D,
-) -> Result<Rc<T>, D::Error> {
-  let data = T::deserialize(deserializer)?;
-  Ok(Rc::new(data))
-}

storage/src/storage_entry.rs

@@ -8,3 +8,50 @@ pub struct StorageEntry {
   pub(crate) refs: Vec<StorageEntryPtr>,
   pub(crate) data: Vec<u8>,
 }
+
+pub struct StorageEntryReader<'a> {
+  pub(crate) entry: &'a StorageEntry,
+  pub(crate) refs_i: usize,
+  pub(crate) data_i: usize,
+}
+
+impl<'a> StorageEntryReader<'a> {
+  pub fn new(entry: &'a StorageEntry) -> Self {
+    Self {
+      entry,
+      refs_i: 0,
+      data_i: 0,
+    }
+  }
+
+  pub fn read_ref(&mut self) -> Option<StorageEntryPtr> {
+    if self.refs_i >= self.entry.refs.len() {
+      return None;
+    }
+
+    let ptr = self.entry.refs[self.refs_i];
+    self.refs_i += 1;
+    Some(ptr)
+  }
+
+  pub fn read_u8(&mut self) -> Option<u8> {
+    if self.data_i >= self.entry.data.len() {
+      return None;
+    }
+
+    let byte = self.entry.data[self.data_i];
+    self.data_i += 1;
+    Some(byte)
+  }
+
+  pub fn read_u64(&mut self) -> Option<u64> {
+    if self.data_i + 8 > self.entry.data.len() {
+      return None;
+    }
+
+    let bytes = self.entry.data.get(self.data_i..self.data_i + 8)?;
+    self.data_i += 8;
+
+    Some(u64::from_le_bytes(bytes.try_into().unwrap()))
+  }
+}

storage/src/storage_ops.rs

@@ -18,7 +18,6 @@ pub trait StorageOps<E> {
   fn get_head(&mut self, ptr: StorageHeadPtr) -> Result<Option<StorageVal>, E>;
   fn set_head(&mut self, ptr: StorageHeadPtr, value: Option<&StorageVal>) -> Result<(), E>;
 
-  fn store_with_replacements(&mut self, value: &StorageVal) -> Result<StorageEntryPtr, E>;
   fn store(&mut self, value: &StorageVal) -> Result<StorageEntryPtr, E>;
 
   fn ref_delta<T>(&mut self, ptr: StoragePtr<T>, delta: i64) -> Result<(), E>;
@@ -61,7 +60,7 @@ where
 
   fn set_head(&mut self, head_ptr: StorageHeadPtr, value: Option<&StorageVal>) -> Result<(), E> {
     if let Some(value) = value {
-      let entry_ptr = self.store_with_replacements(value)?;
+      let entry_ptr = self.store(value)?;
       self.ref_delta(entry_ptr, 1)?;
 
       if let Some(old_entry_ptr) = self.read(head_ptr)? {
@@ -78,14 +77,6 @@ where
     }
   }
 
-  fn store_with_replacements(&mut self, value: &StorageVal) -> Result<StorageEntryPtr, E> {
-    if let Some(ptr) = value.maybe_replace_store(self)? {
-      return Ok(ptr);
-    }
-
-    self.store(value)
-  }
-
   fn store(&mut self, value: &StorageVal) -> Result<StorageEntryPtr, E> {
     let ptr = StoragePtr::random(&mut thread_rng());
     let entry = value.to_storage_entry(self)?;

storage/src/storage_val.rs

@@ -1,43 +1,31 @@
 use std::rc::Rc;
 
-use serde::{Deserialize, Serialize};
-
 use crate::rc_key::RcKey;
-use crate::serde_rc::{deserialize_rc, serialize_rc};
 use crate::storage_entity::StorageEntity;
-use crate::storage_entry::StorageEntry;
+use crate::storage_entry::{StorageEntry, StorageEntryReader};
 use crate::storage_ops::StorageOps;
 use crate::storage_ptr::StorageEntryPtr;
 
 #[cfg(test)]
 use crate::{Storage, StorageBackend};
 
-#[derive(Default, Debug, Serialize, Deserialize, Clone)]
+#[derive(Default, Debug, Clone)]
 pub enum StorageVal {
   #[default]
   Void,
   Number(u64),
   Ptr(StorageEntryPtr),
-  Ref(u64),
-  Compound(
-    #[serde(serialize_with = "serialize_rc", deserialize_with = "deserialize_rc")]
-    Rc<StorageCompoundVal>,
-  ),
+  Compound(Rc<StorageCompoundVal>),
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug)]
 pub enum StorageCompoundVal {
   Array(StorageArray),
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug)]
 pub struct StorageArray {
   pub items: Vec<StorageVal>,
-
-  // Skipping serialization because they're stored in the entry. When converting from an entry, we
-  // copy (todo: move?) the refs from there.
-  #[serde(skip)]
-  pub refs: Vec<StorageEntryPtr>,
 }
 
 impl StorageVal {
@@ -49,73 +37,71 @@ impl StorageVal {
     storage.sb.transaction(|sb| self.numbers_impl(sb))
   }
 
-  pub fn maybe_replace_store<E, SO: StorageOps<E>>(
+  fn write_to_entry<E, SO: StorageOps<E>>(
     &self,
     tx: &mut SO,
-  ) -> Result<Option<StorageEntryPtr>, E> {
-    if let Some(id) = self.cache_key() {
-      if let Some(ptr) = tx.cache_get(id) {
-        return Ok(Some(ptr));
+    entry: &mut StorageEntry,
+  ) -> Result<(), E> {
+    match self {
+      StorageVal::Void => {
+        entry.data.push(0);
       }
-    }
+      StorageVal::Number(n) => {
+        entry.data.push(1);
+        entry.data.extend(n.to_le_bytes());
+      }
+      StorageVal::Ptr(ptr) => {
+        entry.data.push(2);
+        entry.refs.push(*ptr);
+      }
+      StorageVal::Compound(compound) => 'b: {
+        let key = RcKey::from(compound.clone());
 
-    Ok(match &self {
-      StorageVal::Void | StorageVal::Number(_) | StorageVal::Ptr(_) | StorageVal::Ref(_) => None,
-      StorageVal::Compound(compound) => match &**compound {
-        StorageCompoundVal::Array(arr) => 'b: {
-          let mut replacements: Vec<(usize, StorageEntryPtr)> = Vec::new();
-
-          for i in 0..arr.items.len() {
-            if let Some(ptr) = arr.items[i].maybe_replace_store(tx)? {
-              replacements.push((i, ptr));
-            }
-          }
-
-          let key = RcKey::from(compound.clone());
-
-          if replacements.is_empty() {
-            break 'b Some(tx.store_and_cache(self, key)?);
-          }
-
-          let mut new_arr = Vec::<StorageVal>::new();
-          let mut new_refs = arr.refs.clone();
-
-          let mut replacements_iter = replacements.iter();
-          let mut next_replacement = replacements_iter.next();
-
-          for (i, item) in arr.items.iter().enumerate() {
-            if let Some((j, entry_ptr)) = next_replacement {
-              if *j == i {
-                new_arr.push(StorageVal::Ref(new_refs.len() as u64));
-                new_refs.push(*entry_ptr);
-                next_replacement = replacements_iter.next();
-                continue;
-              }
-            }
-
-            new_arr.push(item.clone());
-          }
-
-          Some(tx.store_and_cache(
-            &StorageVal::Compound(Rc::new(StorageCompoundVal::Array(StorageArray {
-              items: new_arr,
-              refs: new_refs,
-            }))),
-            key,
-          )?)
+        if let Some(ptr) = tx.cache_get(key.clone()) {
+          entry.data.push(2);
+          entry.refs.push(ptr);
+          break 'b;
         }
-      },
-    })
+
+        let ptr = tx.store_and_cache(self, key)?;
+        entry.data.push(2);
+        entry.refs.push(ptr);
+      }
+    };
+
+    Ok(())
   }
 
-  fn cache_key(&self) -> Option<RcKey> {
-    match self {
-      StorageVal::Void => None,
-      StorageVal::Number(_) => None,
-      StorageVal::Ptr(_) => None,
-      StorageVal::Ref(_) => None,
-      StorageVal::Compound(compound) => Some(RcKey::from(compound.clone())),
-    }
+  fn read_from_entry<E, SO: StorageOps<E>>(
+    _tx: &mut SO,
+    reader: &mut StorageEntryReader,
+  ) -> Result<StorageVal, E> {
+    let tag = reader.read_u8().unwrap();
+
+    Ok(match tag {
+      0 => StorageVal::Void,
+      1 => {
+        let n = reader.read_u64().unwrap();
+        StorageVal::Number(n)
+      }
+      2 => {
+        let ptr = reader.read_ref().unwrap();
+        StorageVal::Ptr(ptr)
+      }
+      3 => {
+        let len = reader.read_u64().unwrap();
+        let mut items = Vec::new();
+
+        for _ in 0..len {
+          items.push(StorageVal::read_from_entry(_tx, reader)?);
+        }
+
+        let compound = Rc::new(StorageCompoundVal::Array(StorageArray { items }));
+
+        StorageVal::Compound(compound)
+      }
+      _ => panic!("Invalid tag"),
+    })
   }
 
   #[cfg(test)]
@@ -127,20 +113,12 @@ impl StorageVal {
         let entry = tx.read(*ptr)?.unwrap();
         Self::from_storage_entry(tx, entry)?.numbers_impl(tx)
       }
-      StorageVal::Ref(_) => {
-        panic!("Can't lookup ref (shouldn't hit this case)")
-      }
       StorageVal::Compound(compound) => match &**compound {
         StorageCompoundVal::Array(arr) => {
           let mut numbers = Vec::new();
 
           for item in &arr.items {
-            if let StorageVal::Ref(i) = item {
-              let item = &StorageVal::Ptr(arr.refs[*i as usize]);
-              numbers.extend(item.numbers_impl(tx)?);
-            } else {
-              numbers.extend(item.numbers_impl(tx)?);
-            }
+            numbers.extend(item.numbers_impl(tx)?);
           }
 
           Ok(numbers)
@@ -151,41 +129,38 @@ impl StorageVal {
 }
 
 impl StorageEntity for StorageVal {
-  fn to_storage_entry<E, Tx: StorageOps<E>>(&self, _tx: &mut Tx) -> Result<StorageEntry, E> {
-    Ok(StorageEntry {
+  fn to_storage_entry<E, Tx: StorageOps<E>>(&self, tx: &mut Tx) -> Result<StorageEntry, E> {
+    let mut entry = StorageEntry {
       ref_count: 1,
-      refs: match self {
-        StorageVal::Void | StorageVal::Number(_) | StorageVal::Ptr(_) | StorageVal::Ref(_) => {
-          vec![]
-        }
-        StorageVal::Compound(compound) => match &**compound {
-          StorageCompoundVal::Array(arr) => arr.refs.clone(),
-        },
-      },
-      data: bincode::serialize(self).unwrap(),
-    })
-  }
+      refs: Vec::new(),
+      data: Vec::new(),
+    };
 
-  fn from_storage_entry<E, Tx: StorageOps<E>>(
-    _tx: &mut Tx,
-    entry: StorageEntry,
-  ) -> Result<Self, E> {
-    let StorageEntry {
-      ref_count: _,
-      refs,
-      data,
-    } = entry;
-
-    let mut val = bincode::deserialize::<StorageVal>(&data).unwrap();
-
-    if let StorageVal::Compound(compound) = &mut val {
-      match Rc::get_mut(compound).expect("Should be single ref") {
+    match self {
+      StorageVal::Compound(compound) => match &**compound {
         StorageCompoundVal::Array(arr) => {
-          arr.refs = refs;
+          entry.data.push(3);
+
+          entry
+            .data
+            .extend_from_slice(&(arr.items.len() as u64).to_le_bytes());
+
+          for item in &arr.items {
+            item.write_to_entry(tx, &mut entry)?;
+          }
         }
+      },
+      _ => {
+        self.write_to_entry(tx, &mut entry)?;
       }
     };
 
-    Ok(val)
+    Ok(entry)
+  }
+
+  fn from_storage_entry<E, Tx: StorageOps<E>>(tx: &mut Tx, entry: StorageEntry) -> Result<Self, E> {
+    let mut reader = StorageEntryReader::new(&entry);
+
+    Self::read_from_entry(tx, &mut reader)
   }
 }

storage/src/tests.rs

@@ -52,8 +52,7 @@ mod tests_ {
           storage_head_ptr(b"test"),
           Some(&StorageVal::Compound(Rc::new(StorageCompoundVal::Array(
             StorageArray {
-              items: vec![StorageVal::Ref(0), StorageVal::Ref(1)],
-              refs: vec![key0, key1],
+              items: vec![StorageVal::Ptr(key0), StorageVal::Ptr(key1)],
             },
           )))),
         )
@@ -89,14 +88,11 @@ mod tests_ {
               items: vec![
                 StorageVal::Compound(Rc::new(StorageCompoundVal::Array(StorageArray {
                   items: vec![StorageVal::Number(1), StorageVal::Number(2)],
-                  refs: vec![],
                 }))),
                 StorageVal::Compound(Rc::new(StorageCompoundVal::Array(StorageArray {
                   items: vec![StorageVal::Number(3), StorageVal::Number(4)],
-                  refs: vec![],
                 }))),
               ],
-              refs: vec![],
             },
           )))),
         )
@@ -120,7 +116,6 @@ mod tests_ {
     fn impl_<SB: StorageBackend>(storage: &mut Storage<SB>) {
       let mut arr = StorageVal::Compound(Rc::new(StorageCompoundVal::Array(StorageArray {
         items: vec![StorageVal::Number(123)],
-        refs: vec![],
       })));
 
       let depth = 3;
@@ -128,7 +123,6 @@ mod tests_ {
       for _ in 0..depth {
         arr = StorageVal::Compound(Rc::new(StorageCompoundVal::Array(StorageArray {
           items: vec![arr.clone(), arr],
-          refs: vec![],
         })));
       }
 
@@ -160,7 +154,6 @@ mod tests_ {
     fn impl_<SB: StorageBackend>(storage: &mut Storage<SB>) {
       let mut arr = StorageVal::Compound(Rc::new(StorageCompoundVal::Array(StorageArray {
         items: vec![StorageVal::Number(123)],
-        refs: vec![],
       })));
 
       // 2^100 = 1267650600228229401496703205376
@@ -170,7 +163,6 @@ mod tests_ {
       for _ in 0..100 {
         arr = StorageVal::Compound(Rc::new(StorageCompoundVal::Array(StorageArray {
           items: vec![arr.clone(), arr],
-          refs: vec![],
         })));
       }