Implementation of custom efficient serialization of structs in LMDB.

custom_serialization_structs.rs

use crate::{
    core::db::merkle_node::lmdb::LmdbError,
    model::{MerkleHash, MerkleTreeNodeType},
};

/// What is actually stored in the LMDB `merkle_links` table.
/// There is an [`LmdbLink`] for every Merkle tree node entry and Merkle tree node type.
///
/// Properites of stored Merkle links:
///   - Commit nodes do not have a parent (so their `parent_id` is None).
///   - File nodes do not have children.
#[derive(Debug, Clone)]
pub struct LmdbLink {
    pub(crate) parent_id: Option<MerkleHash>,
    /// The children of this node, if it has any.
    pub(crate) children: Vec<MerkleHash>,
}

/// What is actualy stored in the LMDB `merkle_tree_nodes` table.
/// Unlike the [`FileBackend`]'s `node` file, this is how **all** nodes are stored,
/// including file nodes.
#[derive(Debug, Clone)]
pub struct LmdbNode {
    /// The type of node that is encoded. Must align to the `data`.
    pub(crate) kind: MerkleTreeNodeType,
    /// The msgpack-serialized bytes for something that implements [`TMerkleTreeNode`].
    /// Will always be some variant of the concrete [`EMerkleTreeNode`] enum type.
    pub(crate) data: Vec<u8>,
}

//
//
// LmdbLink implementation
//
//

/// Grabs exactly size_of::<$type>() bytes and advances the buffer by that much.
/// If there isn't enough space in the buffer, then $error_variant is returned.
/// Converts all bytes to little-endian.
macro_rules! take_specialized {
    ($name:ident, $type:ty, $error_variant:expr, $final:ty, $convert:expr) => {
        /// Take exactly size_of::<$type>() bytes from the front of `advancing_buf`.
        /// Also move the slice pointer up that many bytes. Return $error_variant
        /// if there's not enough bytes to take. Finally, convert the bytes into
        /// litte-endian and convert into $final.
        #[inline(always)]
        pub(super) fn $name<'a>(advancing_buf: &mut &'a [u8]) -> Result<$final, LmdbError> {
            const SIZE: usize = size_of::<$type>();
            let Some(raw) = take::<SIZE>(advancing_buf) else {
                return Err($error_variant);
            };
            let value_le = <$type>::from_le_bytes(raw);
            Ok($convert(value_le))
        }
    };
}

take_specialized!(
    take_is_parent,
    u8,
    LmdbError::MissingParentIdFlag,
    u8,
    |x: u8| { x }
);
take_specialized!(
    take_num_children,
    u64,
    LmdbError::MissingNumChildren,
    u64,
    |x: u64| { x }
);
take_specialized!(
    take_parent_hash,
    u128,
    LmdbError::MissingParentHash,
    MerkleHash,
    |x: u128| -> MerkleHash { MerkleHash::new(x) }
);
take_specialized!(
    take_child_hash,
    u128,
    LmdbError::MissingChildHash,
    MerkleHash,
    |x: u128| -> MerkleHash { MerkleHash::new(x) }
);

/// ALL VALUES ARE STORED AS LITTLE-ENDIAN !!!
///
///
/// 0 - u8 - if 0, then no parent, if 1, then there's a parent merkle hash (u128)
/// 1 - u64 - number of children hashes. If 0, then there are no hashes written.
/// 2 -
/// 3 -
/// 4 -
/// 5 -
/// 6 -
/// 7 -
/// 8 -
/// 9 - if (0) byte was 1, then bytes 2-5 are the parent merkle hash (u128)
///     if (0) byte was 1, then we don't encode the parent hash: byte 2 is the start of children
/// 3
/// 4
/// 5
/// 6  - the start of children, if (1) was 0, then this is empty
impl LmdbLink {
    /// At a minimum, each LmdbLink is at least 9 bytes.
    /// This would be for [`LmdbLink { parent_id: none, children: vec![] }`].
    const MIN_SIZE: usize = size_of::<u8>() + size_of::<u64>();

    /// If there is a parent, then the minimum size is at least 17 bytes.
    /// This would be for [`LmdbLink { parent_id: Some(u128), children: vec![] }`].
    const MIN_SIZE_PARENT: usize = Self::MIN_SIZE + size_of::<u64>();

    /// Turn into bytes that are stored in LMDB.
    #[inline(always)]
    pub(super) fn serialize(self) -> Vec<u8> {
        // MerkleHash = u128 => 16 bytes
        // For N children there's N*16 bytes
        let size_of_children = size_of::<u128>() * self.children.len();

        // push the header (is_parent?, num children)
        // and push the parent merkle hash if it is present
        let num_children = self.children.len();
        let mut buffer = match self.parent_id {
            Some(parent) => {
                let mut init_buf_parent =
                    Vec::with_capacity(Self::MIN_SIZE_PARENT + size_of_children);
                // there is a parent, so set this flag to 1
                init_buf_parent.push(1_u8.to_le());
                // next bytes are the number of children
                Self::push_len(&mut init_buf_parent, num_children);
                // next bytes are the parent Merkle hash
                Self::push_merkle_hash(&mut init_buf_parent, &parent);
                init_buf_parent
            }
            None => {
                // there's always the first 9 bytes (parent present and number of children)
                // we know there's no parent, so the only extra space is for the children
                let mut init_buf_no_parent = Vec::with_capacity(Self::MIN_SIZE + size_of_children);
                // there's no parent, so set this flag to 0
                init_buf_no_parent.push(0);
                // next bytes are the number of children
                Self::push_len(&mut init_buf_no_parent, num_children);
                init_buf_no_parent
            }
        };

        // push the children
        for child_hash in self.children {
            Self::push_merkle_hash(&mut buffer, &child_hash);
        }

        buffer
    }

    /// Pushes 16 bytes of a u128 in little-endian format into the buffer.
    #[inline(always)]
    fn push_merkle_hash(buffer: &mut Vec<u8>, hash: &MerkleHash) {
        let payload: [u8; 16] = hash.to_le_bytes();
        buffer.extend(payload)
    }

    /// Always converts the usize into a u64, then converts it into little-endian format
    /// of 8 bytes and pushes it into the buffer.
    #[inline(always)]
    fn push_len(buffer: &mut Vec<u8>, length: usize) {
        let payload: [u8; 8] = (length as u64).to_le_bytes();
        buffer.extend(payload);
    }

    /// Convert from bytes that are stored in LMDB.
    #[inline(always)]
    pub(super) fn deserialize(data: &[u8]) -> Result<Self, LmdbError> {
        let mut advancing_buf: &[u8] = &data;

        let is_parent = take_is_parent(&mut advancing_buf)?;
        let num_children = take_num_children(&mut advancing_buf)?;

        let parent_id: Option<MerkleHash> = if is_parent == 0 {
            // no parent => start reading children immediately
            None
        } else if is_parent == 1 {
            // parent => read the parent value, then start reading children
            Some(take_parent_hash(&mut advancing_buf)?)
        } else {
            return Err(LmdbError::InvalidIsParent(is_parent));
        };

        let mut children = Vec::with_capacity(size_of::<MerkleHash>() * num_children as usize);
        for _ in 0..num_children {
            let child = take_child_hash(&mut advancing_buf)?;
            children.push(child);
        }

        Ok(Self {
            parent_id,
            children,
        })
    }

    /// Like [`deserialize`], but it only reads and obtains the parent ID in the link.
    #[inline(always)]
    pub(super) fn deserialize_parent_only(data: &[u8]) -> Result<Option<MerkleHash>, LmdbError> {
        let mut advancing_buf: &[u8] = &data;
        let is_parent = take_is_parent(&mut advancing_buf)?;
        if is_parent == 0 {
            Ok(None)
        } else if is_parent == 1 {
            // read the parent value
            Ok(Some(take_parent_hash(&mut advancing_buf)?))
        } else {
            Err(LmdbError::InvalidIsParent(is_parent))
        }
    }

    /// A reference to this node's optional parent hash.
    pub fn parent_id(&self) -> Option<&MerkleHash> {
        self.parent_id.as_ref()
    }

    /// A reference to this node's children.
    pub fn children(&self) -> &[MerkleHash] {
        &self.children
    }
}

/// Take `size_of::<T>()` bytes from the front of the `data` slice and advance it.
#[inline(always)]
fn take<'a, const N: usize>(data: &mut &'a [u8]) -> Option<[u8; N]> {
    let Some((head, tail)) = data.split_first_chunk::<N>() else {
        return None;
    };
    *data = tail;
    Some(*head)
}

//
//
// LmdbNode implementation
//
//

impl LmdbNode {
    /// Turn into bytes that are stored in LMDB.
    #[inline(always)]
    pub(super) fn serialize(mut self) -> Vec<u8> {
        let kind_as_byte = self.kind.to_u8().to_le();
        if self.data.capacity() > self.data.len() {
            // there's room for at least one byte
            self.data.insert(0, kind_as_byte);
            self.data
        } else {
            // not enough spare capacity in the vec - we have to grow and memcpy
            let mut s = Vec::with_capacity(self.data.len() + 1);
            s.push(kind_as_byte);
            s.extend_from_slice(&self.data);
            s
        }
    }

    /// Convert from bytes that are stored in LMDB.
    #[inline(always)]
    pub(super) fn deserialize(data: &[u8]) -> Result<Self, LmdbError> {
        let Some((head, tail)) = data.split_first_chunk::<1>() else {
            return Err(LmdbError::NoMerkleTreeNodeType);
        };
        let kind_as_byte = u8::from_le_bytes(*head);
        let kind = MerkleTreeNodeType::from_u8(kind_as_byte)?;
        Ok(Self {
            kind,
            data: tail.to_vec(),
        })
    }

    /// Only read enough of a serialized LmdbNode to extract its [`MerkleTreeNodeType`].
    #[inline(always)]
    pub(super) fn deserialize_kind(data: &[u8]) -> Result<MerkleTreeNodeType, LmdbError> {
        let Some((head, _)) = data.split_first_chunk::<1>() else {
            return Err(LmdbError::NoMerkleTreeNodeType);
        };
        let kind_as_byte = u8::from_le_bytes(*head);
        let kind = MerkleTreeNodeType::from_u8(kind_as_byte)?;
        Ok(kind)
    }

    /// The type of Merkle tree node that's serialized.
    pub fn kind(&self) -> MerkleTreeNodeType {
        self.kind
    }

    /// A reference to the msgpack-encoded bytes of the [`EMerkleTreeNode`].
    pub fn data(&self) -> &[u8] {
        &self.data
    }
}