//! Types related to sparse trie nodes and masks.

use crate::Nibbles;
use alloy_primitives::map::HashMap;
use alloy_trie::{nodes::TrieNode, TrieMask};

/// Branch node masks containing `hash_mask` and `tree_mask`.
///
/// Consolidates `hash_mask` and `tree_mask` into a single struct, reducing `HashMap` overhead
/// when storing masks by path. Instead of two separate `HashMap<Nibbles, TrieMask>`,
/// we use a single `HashMap<Nibbles, BranchNodeMasks>`.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct BranchNodeMasks {
    /// Hash mask indicating which children are stored as hashes.
    pub hash_mask: TrieMask,
    /// Tree mask indicating which children are complete subtrees.
    pub tree_mask: TrieMask,
}

impl BranchNodeMasks {
    /// Creates masks from optional hash and tree masks.
    ///
    /// Returns `None` if both masks are `None`, otherwise defaults missing masks to empty.
    pub fn from_optional(hash_mask: Option<TrieMask>, tree_mask: Option<TrieMask>) -> Option<Self> {
        match (hash_mask, tree_mask) {
            (Some(h), Some(t)) => Some(Self { hash_mask: h, tree_mask: t }),
            (Some(h), None) => Some(Self { hash_mask: h, tree_mask: TrieMask::default() }),
            (None, Some(t)) => Some(Self { hash_mask: TrieMask::default(), tree_mask: t }),
            (None, None) => None,
        }
    }

    /// Returns true if both interior masks are empty.
    pub const fn is_empty(&self) -> bool {
        self.hash_mask.is_empty() && self.tree_mask.is_empty()
    }

    /// Updates this set of masks for `nibble` based on that child's mask bits.
    pub const fn set_child_bits(&mut self, nibble: u8, hash_mask_bit: bool, tree_mask_bit: bool) {
        if hash_mask_bit {
            self.hash_mask.set_bit(nibble);
        } else {
            self.hash_mask.unset_bit(nibble);
        }
        if tree_mask_bit {
            self.tree_mask.set_bit(nibble);
        } else {
            self.tree_mask.unset_bit(nibble);
        }
    }
}

/// Map from nibble path to branch node masks.
pub type BranchNodeMasksMap = HashMap<Nibbles, BranchNodeMasks>;

/// Carries all information needed by a sparse trie to reveal a particular node.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProofTrieNode {
    /// Path of the node.
    pub path: Nibbles,
    /// The node itself.
    pub node: TrieNode,
    /// Tree and hash masks for the node, if known.
    /// Both masks are always set together (from database branch nodes).
    pub masks: Option<BranchNodeMasks>,
}
