fix(merkle): only store lower common levels into pruned branch

Author: Rexagon

src/cell/mod.rs

@@ -1337,6 +1337,13 @@ impl<'a> arbitrary::Arbitrary<'a> for LevelMask {
 #[derive(Clone)]
 pub struct LevelMaskIter(u8);
 
+impl LevelMaskIter {
+    /// Returns `true` if all levels were consumed.
+    pub fn is_empty(&self) -> bool {
+        self.0 == 0
+    }
+}
+
 impl Iterator for LevelMaskIter {
     type Item = u8;
 

src/merkle/pruned_branch.rs

@@ -8,10 +8,12 @@ pub fn make_pruned_branch(
     context: &dyn CellContext,
 ) -> Result<Cell, Error> {
     let descriptor = cell.descriptor();
-    let cell_level_mask = descriptor.level_mask();
+    let mut cell_level_mask = descriptor.level_mask();
 
     let mut builder = CellBuilder::new();
-    let level_mask = LevelMask::new(cell_level_mask.to_byte() | (1 << merkle_depth));
+
+    let new_level = 1 << merkle_depth;
+    let level_mask = LevelMask::new(cell_level_mask.to_byte() | new_level);
 
     builder.set_exotic(true);
 
@@ -20,6 +22,9 @@ pub fn make_pruned_branch(
         level_mask.to_byte(),
     ]));
 
+    // Only write levels lower than the new level.
+    cell_level_mask = LevelMask::new(cell_level_mask.to_byte() & (new_level - 1));
+
     for level in cell_level_mask {
         _ = builder.store_u256(cell.hash(level));
     }
@@ -33,6 +38,9 @@ pub fn make_pruned_branch(
 
 #[cfg(test)]
 mod test {
+    use crate::boc::Boc;
+    use crate::merkle::MerkleProof;
+
     use super::*;
 
     #[test]
@@ -56,4 +64,29 @@ mod test {
             make_pruned_branch(virtual_cell, 0, Cell::empty_context()).unwrap();
         assert_eq!(pruned_branch.as_ref(), virtual_pruned_branch.as_ref());
     }
+
+    #[test]
+    fn partial_pruned() -> anyhow::Result<()> {
+        let cell = CellBuilder::build_from((
+            CellBuilder::build_from((0xaa_u8, Cell::empty_cell()))?,
+            CellBuilder::build_from((0xbb_u8, Cell::empty_cell()))?,
+        ))?;
+        println!("Original: {}", Boc::encode_base64(&cell));
+
+        // Prune left cell.
+        let with_left_pruned = {
+            let usage_tree = UsageTree::new(UsageTreeMode::OnLoad);
+            let tracked = usage_tree.track(&cell);
+            tracked.reference(1).unwrap().touch_recursive();
+
+            MerkleProof::create(cell.as_ref(), usage_tree).build_raw_ext(Cell::empty_context())?
+        };
+        println!("Left pruned: {}", Boc::encode_base64(&with_left_pruned));
+
+        // Full pruned.
+        let pruned = make_pruned_branch(with_left_pruned.as_ref(), 0, Cell::empty_context())?;
+        println!("Full pruned: {}", Boc::encode_base64(&pruned));
+
+        Ok(())
+    }
 }