monotree: Correct path compression collapse in delete_key

src/sdk/src/monotree/bits.rs

@@ -25,6 +25,189 @@ use super::{
 };
 use crate::GenericResult;
 
+/// An owned representation of a bit path, compatible with Bits serialization.
+/// Internally stores bits normalized to start at position 0.
+/// Tracks the original offset for correct serialization.
+#[derive(Debug, Clone, PartialEq)]
+pub struct BitsOwned {
+    /// Bits stored normalized (starting at bit 0, MSB-first)
+    path: Vec<u8>,
+    /// Number of valid bits
+    len: BitsLen,
+    /// Original starting bit position (for serialization)
+    offset: BitsLen,
+}
+
+impl BitsOwned {
+    /// Create empty `BitsOwned` with given length and offset
+    #[inline]
+    pub fn new(len: BitsLen, offset: BitsLen) -> Self {
+        let num_bytes = len.div_ceil(8) as usize;
+        Self { path: vec![0u8; num_bytes], len, offset }
+    }
+
+    /// Create from raw normalized bytes (bits start at position 0)
+    #[inline]
+    fn from_raw(path: Vec<u8>, len: BitsLen, offset: BitsLen) -> Self {
+        Self { path, len, offset }
+    }
+
+    #[inline]
+    pub fn len(&self) -> BitsLen {
+        self.len
+    }
+
+    #[inline]
+    pub fn is_empty(&self) -> bool {
+        self.len == 0
+    }
+
+    #[inline]
+    pub fn first(&self) -> bool {
+        debug_assert!(self.is_empty(), "cannot get first bit of empty BitsOwned");
+        self.path[0] & 0x80 != 0
+    }
+
+    /// Serialize to bytes in monotree format: `[start:2][end:2][path_bytes]`
+    /// Shifts bits to match the original offset position.
+    pub fn to_bytes(&self) -> GenericResult<Vec<u8>> {
+        let start = self.offset;
+        let end = self.offset + self.len;
+
+        let byte_start = (start / 8) as usize;
+        let byte_end = end.div_ceil(8) as usize;
+        let num_path_bytes = byte_end - byte_start;
+
+        let mut result = Vec::with_capacity(4 + num_path_bytes);
+        result.extend_from_slice(&start.to_be_bytes());
+        result.extend_from_slice(&end.to_be_bytes());
+
+        if num_path_bytes == 0 {
+            return Ok(result)
+        }
+
+        let bit_shift = (start % 8) as usize;
+
+        if bit_shift == 0 {
+            // Fast path: no shifting needed
+            result.extend_from_slice(&self.path[..num_path_bytes]);
+        } else {
+            // Shift bits right by bit_shift to align with offset
+            for i in 0..num_path_bytes {
+                let hi = if i == 0 { 0 } else { self.path.get(i - 1).copied().unwrap_or(0) };
+                let lo = self.path.get(i).copied().unwrap_or(0);
+                let byte = (hi << (8 - bit_shift)) | (lo >> bit_shift);
+                result.push(byte);
+            }
+        }
+
+        // Mask leading bits (before start)
+        let lead_mask = 0xFFu8 >> (start % 8) as u8;
+        result[4] &= lead_mask;
+
+        // Mask trailing bits (after end)
+        let trail = (end % 8) as u8;
+        if trail != 0 {
+            let trail_mask = 0xFFu8 << (8 - trail);
+            let last_idx = result.len() - 1;
+            result[last_idx] &= trail_mask;
+        }
+
+        Ok(result)
+    }
+}
+
+/// Extract bits from a `Bits` into a normalized `Vec<u8>` (starting at bit 0)
+#[inline]
+fn extract_normalized(bits: &Bits) -> Vec<u8> {
+    let len = bits.len();
+    if len == 0 {
+        return Vec::new();
+    }
+
+    let num_bytes = len.div_ceil(8) as usize;
+    let bit_offset = (bits.range.start % 8) as usize;
+    let byte_start = (bits.range.start / 8) as usize;
+    let byte_end = bits.range.end.div_ceil(8) as usize;
+    let src = &bits.path[byte_start..byte_end];
+
+    let mut result = vec![0u8; num_bytes];
+
+    if bit_offset == 0 {
+        // Fast path: already byte-aligned
+        result[..num_bytes.min(src.len())].copy_from_slice(&src[..num_bytes.min(src.len())]);
+    } else {
+        // Shift left by bit_offset to normalize
+        for (i, bit) in result.iter_mut().enumerate().take(num_bytes) {
+            let hi = src.get(i).copied().unwrap_or(0);
+            let lo = src.get(i + 1).copied().unwrap_or(0);
+            *bit = (hi << bit_offset) | (lo >> (8 - bit_offset));
+        }
+    }
+
+    // Mask trailing bits
+    let trail = (len % 8) as usize;
+    if trail != 0 && !result.is_empty() {
+        let last = result.len() - 1;
+        result[last] &= 0xFF << (8 - trail);
+    }
+
+    result
+}
+
+/// Append normalized bits from src to dst at bit position dst_len
+#[inline]
+fn append_normalized(dst: &mut Vec<u8>, dst_len: BitsLen, src: &[u8], src_len: BitsLen) {
+    if src_len == 0 {
+        return
+    }
+
+    let total_len = dst_len + src_len;
+    let new_bytes = total_len.div_ceil(8) as usize;
+    dst.resize(new_bytes, 0);
+
+    let bit_offset = (dst_len % 8) as usize;
+    let byte_offset = (dst_len / 8) as usize;
+
+    if bit_offset == 0 {
+        // Fast path: byte-aligned append
+        let copy_len = src.len().min(new_bytes - byte_offset);
+        dst[byte_offset..byte_offset + copy_len].copy_from_slice(&src[..copy_len]);
+    } else {
+        // Shift src right by bit_offset and OR into dst
+        for i in 0..src.len() {
+            let shifted_hi = src[i] >> bit_offset;
+            let shifted_lo = src[i] << (8 - bit_offset);
+
+            dst[byte_offset + i] |= shifted_hi;
+            if byte_offset + i + 1 < dst.len() {
+                dst[byte_offset + i + 1] |= shifted_lo;
+            }
+        }
+    }
+}
+
+/// Merge a `BitsOwned` with a `Bits`, appending b's bits after a's.
+#[inline]
+pub fn merge_owned_and_bits(a: &BitsOwned, b: &Bits) -> BitsOwned {
+    let a_len = a.len();
+    let b_len = b.len();
+    let total_len = a_len + b_len;
+
+    if total_len == 0 {
+        return BitsOwned::from_raw(Vec::new(), 0, a.offset);
+    }
+
+    // Start with a copy of a's path
+    let mut path = a.path.clone();
+
+    // Extract and append b's normalized bits
+    let b_norm = extract_normalized(b);
+    append_normalized(&mut path, a_len, &b_norm, b_len);
+
+    BitsOwned::from_raw(path, total_len, a.offset)
+}
+
 #[derive(Debug, Clone, PartialEq)]
 /// `BitVec` implementation based on bytes slice.
 pub struct Bits<'a> {
@@ -97,4 +280,32 @@ impl<'a> Bits<'a> {
     pub fn len_common_bits(a: &Self, b: &Self) -> BitsLen {
         len_lcp(a.path, &a.range, b.path, &b.range)
     }
+
+    /// Convert to `BitsOwned`, preserving the range offset.
+    #[inline]
+    pub fn to_bits_owned(&self) -> BitsOwned {
+        let path = extract_normalized(self);
+        BitsOwned::from_raw(path, self.len(), self.range.start)
+    }
+
+    /// Merge two `Bits` into a new `BitsOwned`, preserving the first's offset.
+    #[inline]
+    pub fn merge(a: &Self, b: &Self) -> BitsOwned {
+        let a_len = a.len();
+        let b_len = b.len();
+        let total_len = a_len + b_len;
+
+        if total_len == 0 {
+            return BitsOwned::from_raw(Vec::new(), 0, a.range.start);
+        }
+
+        // Extract normalized bits from both
+        let mut path = extract_normalized(a);
+        let b_norm = extract_normalized(b);
+
+        // Append b's bits
+        append_normalized(&mut path, a_len, &b_norm, b_len);
+
+        BitsOwned::from_raw(path, total_len, a.range.start)
+    }
 }

src/sdk/src/monotree/tests.rs

@@ -212,3 +212,25 @@ fn monotree_test_deterministic_ordering() {
     assert_eq!(root1, None);
     assert_eq!(root2, None);
 }
+
+#[test]
+fn monotree_test_insert_remove_identity() {
+    let keys = random_hashes(3);
+    let values = random_hashes(3);
+
+    println!("{:?}", keys);
+    println!("{:?}", values);
+
+    let db = MemoryDb::new();
+    let mut tree = Monotree::new(db);
+    let mut root = tree.inserts(None, &keys, &values).unwrap();
+    let original_root = root;
+
+    // Insert then remove a new key
+    let temp_key = random_hashes(1)[0];
+    let temp_value = random_hashes(1)[0];
+    root = tree.insert(root.as_ref(), &temp_key, &temp_value).unwrap();
+    root = tree.remove(root.as_ref(), &temp_key).unwrap();
+
+    assert_eq!(root, original_root);
+}

src/sdk/src/monotree/tree.rs

@@ -21,7 +21,7 @@ use hashbrown::{HashMap, HashSet};
 use sled_overlay::{sled::Tree, SledDbOverlay};
 
 use super::{
-    bits::Bits,
+    bits::{merge_owned_and_bits, Bits, BitsOwned},
     node::{Node, Unit},
     utils::{get_sorted_indices, slice_to_hash},
     Hash, Proof, HASH_LEN, ROOT_KEY,
@@ -372,6 +372,64 @@ impl<D: MonotreeStorageAdapter> Monotree<D> {
         Ok(Some(hash))
     }
 
+    /// Create and store a soft node using owned bits.
+    fn put_soft_node_owned(
+        &mut self,
+        target_hash: &[u8],
+        bits: &BitsOwned,
+    ) -> GenericResult<Option<Hash>> {
+        let bits_bytes = bits.to_bytes()?;
+        let node_bytes = [target_hash, &bits_bytes[..], &[0x00u8]].concat();
+        let node_hash = Self::hash_digest(&node_bytes);
+        self.db.put(&node_hash, node_bytes)?;
+        Ok(Some(node_hash))
+    }
+
+    /// Create hard node with owned left bits and preserved right bits (unchanged sibling).
+    fn put_hard_node_mixed(
+        &mut self,
+        left_hash: &[u8],
+        left_bits: &BitsOwned,
+        right: &Unit,
+    ) -> GenericResult<Option<Hash>> {
+        let lb_bytes = left_bits.to_bytes()?;
+        let rb_bytes = right.bits.to_bytes()?;
+
+        let (lh, lb, rh, rb) = if right.bits.first() {
+            (left_hash, &lb_bytes[..], right.hash, &rb_bytes[..])
+        } else {
+            (right.hash, &rb_bytes[..], left_hash, &lb_bytes[..])
+        };
+
+        let node_bytes = [lh, lb, rb, rh, &[0x01u8]].concat();
+        let node_hash = Self::hash_digest(&node_bytes);
+        self.db.put(&node_hash, node_bytes)?;
+        Ok(Some(node_hash))
+    }
+
+    /// Collapse a path through soft nodes, accumulating bit prefixes.
+    /// Returns `(target_hash, accumulated_bits)`.
+    fn collapse_to_target(
+        &mut self,
+        hash: &[u8],
+        prefix: BitsOwned,
+    ) -> GenericResult<(Hash, BitsOwned)> {
+        let Some(bytes) = self.db.get(hash)? else {
+            // Leaf value - return it with the accumulated prefix
+            return Ok((slice_to_hash(hash), prefix))
+        };
+
+        let node = Node::from_bytes(&bytes)?;
+        match node {
+            Node::Soft(Some(child)) => {
+                let merged = merge_owned_and_bits(&prefix, &child.bits);
+                self.collapse_to_target(child.hash, merged)
+            }
+            Node::Hard(_, _) => Ok((slice_to_hash(hash), prefix)),
+            _ => unreachable!("unexpected node type in collapse_to_target"),
+        }
+    }
+
     /// Recursively insert a bytes (in forms of Bits) and a leaf into the tree.
     ///
     /// Optimisation in `monotree` is mainly to compress the path as much as possible
@@ -462,19 +520,64 @@ impl<D: MonotreeStorageAdapter> Monotree<D> {
         let n = Bits::len_common_bits(&unit.bits, &bits);
 
         match n {
-            n if n == bits.len() => match right {
-                Some(_) => self.put_node(Node::new(None, right)),
-                None => Ok(None),
-            },
+            // Found the exact key to delete
+            n if n == bits.len() => {
+                match right {
+                    Some(ref sibling) => {
+                        // Collapse sibling path through any soft nodes
+                        let prefix = sibling.bits.to_bits_owned();
+                        let (target, merged_bits) =
+                            self.collapse_to_target(sibling.hash, prefix)?;
+                        self.put_soft_node_owned(&target, &merged_bits)
+                    }
+                    None => Ok(None),
+                }
+            }
+            // Recurse into subtree
             n if n == unit.bits.len() => {
                 let hash = self.delete_key(unit.hash, bits.drop(n))?;
                 match (hash, &right) {
                     (None, None) => Ok(None),
-                    (None, Some(_)) => self.put_node(Node::new(None, right)),
-                    (Some(ref hash), _) => {
-                        let unit = unit.to_owned();
-                        let left = Some(Unit { hash, ..unit });
-                        self.put_node(Node::new(left, right))
+
+                    (None, Some(sibling)) => {
+                        // Child deleted, collapse sibling
+                        let prefix = sibling.bits.to_bits_owned();
+                        let (target, merged_bits) =
+                            self.collapse_to_target(sibling.hash, prefix)?;
+                        self.put_soft_node_owned(&target, &merged_bits)
+                    }
+
+                    (Some(ref new_child), None) => {
+                        // Child modified, no sibling - collapse through
+                        let prefix = unit.bits.to_bits_owned();
+                        let (target, merged_bits) = self.collapse_to_target(new_child, prefix)?;
+                        self.put_soft_node_owned(&target, &merged_bits)
+                    }
+
+                    (Some(ref new_child), Some(sibling)) => {
+                        // Child modified, sibling exists - check if we need to inline soft node
+                        match self.db.get(new_child)? {
+                            Some(child_bytes) => {
+                                match Node::from_bytes(&child_bytes)? {
+                                    Node::Soft(Some(inner)) => {
+                                        // Inline the soft node: merge parent bits + soft node bits
+                                        let merged = Bits::merge(&unit.bits, &inner.bits);
+                                        self.put_hard_node_mixed(inner.hash, &merged, sibling)
+                                    }
+                                    Node::Hard(_, _) => {
+                                        // Can't inline hard node - keep reference
+                                        let parent_bits = unit.bits.to_bits_owned();
+                                        self.put_hard_node_mixed(new_child, &parent_bits, sibling)
+                                    }
+                                    _ => unreachable!(),
+                                }
+                            }
+                            None => {
+                                // new_child is a leaf valuie
+                                let parent_bits = unit.bits.to_bits_owned();
+                                self.put_hard_node_mixed(new_child, &parent_bits, sibling)
+                            }
+                        }
                     }
                 }
             }