From 1f4cc00bb5ce08c2107a6774444b9bfe75ad60bf Mon Sep 17 00:00:00 2001
From: Ivan Kalinin <i.kalinin@dexpa.io>
Date: Tue, 15 Aug 2023 16:26:27 +0200
Subject: [PATCH] Eliminate intermediate serialization for dict insert

---
 src/cell/builder.rs | 79 +++++++++++++++++++++++++++++++++++++++++++++
 src/dict/mod.rs     | 10 +++---
 src/dict/typed.rs   |  9 +++---
 src/util.rs         | 11 +++++++
 4 files changed, 99 insertions(+), 10 deletions(-)

diff --git a/src/cell/builder.rs b/src/cell/builder.rs
index e8b2195a..5b33b7e3 100644
--- a/src/cell/builder.rs
+++ b/src/cell/builder.rs
@@ -300,6 +300,14 @@ impl CellBuilder {
         unsafe { CellSlice::new_unchecked(IntermediateDataCell::wrap(self)) }
     }
 
+    /// Returns a slice which contains builder data and references.
+    ///
+    /// NOTE: intermediate cell hash is undefined.
+    pub fn as_full_slice(&self) -> CellSlice<'_> {
+        // SAFETY: we interpret cell builder data as ordinary cell
+        unsafe { CellSlice::new_unchecked(IntermediateFullCell::wrap(self)) }
+    }
+
     /// Returns an underlying cell data.
     #[inline]
     pub fn raw_data(&self) -> &[u8; 128] {
@@ -1059,6 +1067,77 @@ impl CellImpl for IntermediateDataCell {
     }
 }
 
+#[repr(transparent)]
+struct IntermediateFullCell(CellBuilder);
+
+impl IntermediateFullCell {
+    #[inline(always)]
+    const fn wrap(value: &CellBuilder) -> &Self {
+        // SAFETY: IntermediateFullCell is #[repr(transparent)]
+        unsafe { &*(value as *const CellBuilder as *const Self) }
+    }
+}
+
+impl CellImpl for IntermediateFullCell {
+    fn descriptor(&self) -> CellDescriptor {
+        CellDescriptor {
+            d1: CellDescriptor::compute_d1(LevelMask::EMPTY, false, self.0.references.len() as u8),
+            d2: CellDescriptor::compute_d2(self.0.bit_len),
+        }
+    }
+
+    fn data(&self) -> &[u8] {
+        self.0.raw_data()
+    }
+
+    fn bit_len(&self) -> u16 {
+        self.0.bit_len
+    }
+
+    fn reference(&self, index: u8) -> Option<&DynCell> {
+        match self.0.references.get(index) {
+            Some(cell) => Some(cell.as_ref()),
+            None => None,
+        }
+    }
+
+    fn reference_cloned(&self, index: u8) -> Option<Cell> {
+        self.0.references.get(index).cloned()
+    }
+
+    fn virtualize(&self) -> &DynCell {
+        self
+    }
+
+    fn hash(&self, _: u8) -> &HashBytes {
+        panic!("Hash for an intermediate data cell is not defined");
+    }
+
+    fn depth(&self, _: u8) -> u16 {
+        0
+    }
+
+    fn take_first_child(&mut self) -> Option<Cell> {
+        None
+    }
+
+    fn replace_first_child(&mut self, parent: Cell) -> Result<Cell, Cell> {
+        Err(parent)
+    }
+
+    fn take_next_child(&mut self) -> Option<Cell> {
+        None
+    }
+
+    #[cfg(feature = "stats")]
+    fn stats(&self) -> CellTreeStats {
+        CellTreeStats {
+            bit_count: self.0.bit_len as u64,
+            cell_count: 1 + self.0.references.len() as u64,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/dict/mod.rs b/src/dict/mod.rs
index f9bdf37d..5788c415 100644
--- a/src/dict/mod.rs
+++ b/src/dict/mod.rs
@@ -196,7 +196,7 @@ pub fn dict_insert(
     root: &Option<Cell>,
     key: &mut CellSlice,
     key_bit_len: u16,
-    value: &CellSlice,
+    value: &dyn Store,
     mode: SetMode,
     finalizer: &mut dyn Finalizer,
 ) -> Result<(Option<Cell>, bool), Error> {
@@ -293,7 +293,7 @@ pub fn dict_insert_owned(
     root: &Option<Cell>,
     key: &mut CellSlice,
     key_bit_len: u16,
-    value: &CellSlice,
+    value: &dyn Store,
     mode: SetMode,
     finalizer: &mut dyn Finalizer,
 ) -> Result<(Option<Cell>, bool, Option<CellSliceParts>), Error> {
@@ -906,12 +906,12 @@ pub fn dict_remove_bound_owned(
 fn make_leaf(
     key: &CellSlice,
     key_bit_len: u16,
-    value: &CellSlice,
+    value: &dyn Store,
     finalizer: &mut dyn Finalizer,
 ) -> Result<Cell, Error> {
     let mut builder = CellBuilder::new();
     ok!(write_label(key, key_bit_len, &mut builder));
-    ok!(builder.store_slice(value));
+    ok!(value.store_into(&mut builder, finalizer));
     builder.build_ext(finalizer)
 }
 
@@ -921,7 +921,7 @@ fn split_edge(
     prefix: &mut CellSlice,
     lcp: &CellSlice,
     key: &mut CellSlice,
-    value: &CellSlice,
+    value: &dyn Store,
     finalizer: &mut dyn Finalizer,
 ) -> Result<Cell, Error> {
     // Advance the key
diff --git a/src/dict/typed.rs b/src/dict/typed.rs
index 11dff644..60a04861 100644
--- a/src/dict/typed.rs
+++ b/src/dict/typed.rs
@@ -677,14 +677,13 @@ where
         V: Store,
     {
         let (new_root, changed) = {
-            let mut builder = CellBuilder::new();
-            ok!(key.store_into(&mut builder, &mut Cell::default_finalizer()));
-            let value = ok!(CellBuilder::build_from_ext(value, finalizer));
+            let mut key_builder = CellBuilder::new();
+            ok!(key.store_into(&mut key_builder, &mut Cell::default_finalizer()));
             ok!(dict_insert(
                 &self.root,
-                &mut builder.as_data_slice(),
+                &mut key_builder.as_data_slice(),
                 K::BITS,
-                &ok!(value.as_ref().as_slice()),
+                value,
                 mode,
                 finalizer
             ))
diff --git a/src/util.rs b/src/util.rs
index fd4539e0..dd128547 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -87,6 +87,17 @@ impl<T, const N: usize> ArrayVec<T, N> {
         self.len += 1;
     }
 
+    /// Returns a reference to an element.
+    pub fn get(&self, n: u8) -> Option<&T> {
+        if n < self.len {
+            let references = self.inner.as_ptr() as *const T;
+            // SAFETY: {len} elements were initialized, n < len
+            Some(unsafe { &*references.add(n as usize) })
+        } else {
+            None
+        }
+    }
+
     /// Returns the inner data without dropping its elements.
     ///
     /// # Safety