cfx_storage/impls/snapshot_sync/offer/mpt_slicer.rs
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123
// Copyright 2019 Conflux Foundation. All rights reserved.
// Conflux is free software and distributed under GNU General Public License.
// See http://www.gnu.org/licenses/
pub struct MptSlicer<'a> {
    cursor: MptCursor<
        &'a mut dyn SnapshotMptTraitRead,
        BasicPathNode<&'a mut dyn SnapshotMptTraitRead>,
    >,
}
impl<'a> MptSlicer<'a> {
    pub fn new(mpt: &'a mut dyn SnapshotMptTraitRead) -> Result<Self> {
        let mut cursor = MptCursor::new(mpt);
        cursor.load_root()?;
        Ok(Self { cursor })
    }
    pub fn new_from_key(
        mpt: &'a mut dyn SnapshotMptTraitRead, key: &[u8],
    ) -> Result<Self> {
        let mut slicer = Self::new(mpt)?;
        slicer.cursor.open_path_for_key::<access_mode::Read>(key)?;
        Ok(slicer)
    }
    pub fn to_proof(&self) -> TrieProof { self.cursor.to_proof() }
    pub fn get_range_end_key(&self) -> Option<&[u8]> {
        // The cursor stops at the key which just exceed,the rlp_size_limit,
        // or at the root node.
        let key = self
            .cursor
            .get_path_nodes()
            .last()
            .unwrap()
            .get_path_to_node()
            .path_slice();
        if key.len() == 0 {
            None
        } else {
            Some(key)
        }
    }
    pub fn advance(&mut self, mut rlp_size_limit: u64) -> Result<()> {
        let current_node = self.cursor.current_node_mut();
        // First, check the value of this node, if we are the first time
        // visiting this node.
        if current_node.next_child_index == 0 {
            let maybe_value = current_node.value_as_slice();
            match maybe_value {
                MptValue::Some(value) => {
                    let key_value_size = rlp_key_value_len(
                        current_node.get_path_to_node().path_size(),
                        value.len(),
                    );
                    if rlp_size_limit <= key_value_size {
                        return Ok(());
                    } else {
                        rlp_size_limit -= key_value_size;
                    }
                }
                _ => {}
            }
        }
        for (
            this_child_index,
            &SubtreeMerkleWithSize {
                ref subtree_size, ..
            },
        ) in current_node
            .trie_node
            .get_children_table_ref()
            .iter()
            .set_start_index(current_node.next_child_index)
        {
            if *subtree_size <= rlp_size_limit {
                rlp_size_limit -= *subtree_size;
            } else {
                let child_node = unsafe {
                    // Mute Rust borrow checker because there is no way
                    // around. It's actually safe to open_child_index while
                    // we immutably borrows the trie_node, because
                    // open_child_index won't modify
                    // trie_node.
                    &mut *(current_node
                        as *const BasicPathNode<
                            &'a mut dyn SnapshotMptTraitRead,
                        >
                        as *mut BasicPathNode<&'a mut dyn SnapshotMptTraitRead>)
                }
                .open_child_index(this_child_index)?
                // Unwrap is fine because the child is guaranteed to exist.
                .unwrap();
                self.cursor.push_node(child_node);
                return self.advance(rlp_size_limit);
            }
        }
        // Pop-up because subtree isn't large enough.
        if rlp_size_limit > 0 && self.cursor.get_path_nodes().len() > 1 {
            self.cursor.pop_one_node()?;
            return self.advance(rlp_size_limit);
        }
        Ok(())
    }
}
use super::super::super::{
    super::{
        storage_db::snapshot_mpt::{
            SnapshotMptTraitRead, SubtreeMerkleWithSize,
        },
        utils::access_mode,
    },
    errors::*,
    merkle_patricia_trie::{mpt_cursor::*, *},
};
use primitives::MptValue;