hyli-model 0.14.0

use borsh::{BorshDeserialize, BorshSerialize};
use serde::{Deserialize, Serialize};
use sha3::{Digest, Sha3_256};
use std::{collections::HashMap, fmt::Display, sync::RwLock};

use crate::*;

#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
pub enum MempoolStatusEvent {
    WaitingDissemination {
        parent_data_proposal_hash: DataProposalHash,
        txs: Vec<Transaction>,
    },
    DataProposalCreated {
        parent_data_proposal_hash: DataProposalHash,
        data_proposal_hash: DataProposalHash,
        txs_metadatas: Vec<TransactionMetadata>,
    },
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub enum MempoolBlockEvent {
    BuiltSignedBlock(SignedBlock),
    StartedBuildingBlocks(BlockHeight),
}

#[derive(
    Debug,
    Clone,
    Serialize,
    Deserialize,
    BorshSerialize,
    BorshDeserialize,
    PartialEq,
    Eq,
    Hash,
    Ord,
    PartialOrd,
)]
#[cfg_attr(feature = "full", derive(utoipa::ToSchema))]
pub enum DataProposalParent {
    LaneRoot(LaneId),
    DP(DataProposalHash),
}

impl DataProposalParent {
    pub fn dp_hash(&self) -> Option<&DataProposalHash> {
        match self {
            DataProposalParent::DP(hash) => Some(hash),
            DataProposalParent::LaneRoot(_) => None,
        }
    }

    pub fn lane_id(&self) -> Option<&LaneId> {
        match self {
            DataProposalParent::LaneRoot(lane_id) => Some(lane_id),
            DataProposalParent::DP(_) => None,
        }
    }

    pub fn is_lane_root(&self) -> bool {
        matches!(self, DataProposalParent::LaneRoot(_))
    }

    pub fn as_tx_parent_hash(&self) -> DataProposalHash {
        match self {
            DataProposalParent::LaneRoot(lane_id) => DataProposalHash(lane_id.to_bytes()),
            DataProposalParent::DP(hash) => hash.clone(),
        }
    }
}

#[derive(Debug, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
#[readonly::make]
pub struct DataProposal {
    pub parent_data_proposal_hash: DataProposalParent,
    pub txs: Vec<Transaction>,
    /// Internal cache of the hash of the transaction
    #[borsh(skip)]
    hash_cache: RwLock<Option<DataProposalHash>>,
}

impl DataProposal {
    pub fn new(parent_data_proposal_hash: DataProposalHash, txs: Vec<Transaction>) -> Self {
        Self {
            parent_data_proposal_hash: DataProposalParent::DP(parent_data_proposal_hash),
            txs,
            hash_cache: RwLock::new(None),
        }
    }

    pub fn new_root(lane_id: LaneId, txs: Vec<Transaction>) -> Self {
        Self {
            parent_data_proposal_hash: DataProposalParent::LaneRoot(lane_id),
            txs,
            hash_cache: RwLock::new(None),
        }
    }

    pub fn remove_proofs(&mut self) {
        self.txs.iter_mut().for_each(|tx| {
            match &mut tx.transaction_data {
                TransactionData::VerifiedProof(proof_tx) => {
                    proof_tx.proof = None;
                }
                TransactionData::Proof(_) => {
                    // This can never happen.
                    // A DataProposal that has been processed has turned all TransactionData::Proof into TransactionData::VerifiedProof
                    unreachable!();
                }
                TransactionData::Blob(_) => {}
            }
        });
    }

    pub fn take_proofs(&mut self) -> HashMap<TxHash, ProofData> {
        self.txs
            .iter_mut()
            .filter_map(|tx| {
                match &mut tx.transaction_data {
                    TransactionData::VerifiedProof(proof_tx) => {
                        proof_tx.proof.take().map(|proof| (tx.hashed(), proof))
                    }
                    TransactionData::Proof(_) => {
                        // This can never happen.
                        // A DataProposal that has been processed has turned all TransactionData::Proof into TransactionData::VerifiedProof
                        unreachable!();
                    }
                    TransactionData::Blob(_) => None,
                }
            })
            .collect()
    }

    pub fn hydrate_proofs(&mut self, mut proofs: HashMap<TxHash, ProofData>) {
        self.txs.iter_mut().for_each(|tx| {
            let tx_hash = tx.hashed();
            if let TransactionData::VerifiedProof(ref mut vpt) = tx.transaction_data {
                if vpt.proof.is_none() {
                    if let Some(proof) = proofs.remove(&tx_hash) {
                        vpt.proof = Some(proof);
                    }
                }
            }
        });
    }

    /// This is used to set the hash of the DataProposal when we can trust we know it
    /// (specifically - deserializating from local storage)
    /// # Safety
    /// Marked unsafe so you think twice before using it, but this is safe rust-wise.
    pub unsafe fn unsafe_set_hash(&mut self, hash: &DataProposalHash) {
        self.hash_cache.write().unwrap().replace(hash.clone());
    }
}

impl Clone for DataProposal {
    fn clone(&self) -> Self {
        DataProposal {
            parent_data_proposal_hash: self.parent_data_proposal_hash.clone(),
            txs: self.txs.clone(),
            hash_cache: RwLock::new(self.hash_cache.read().unwrap().clone()),
        }
    }
}

impl PartialEq for DataProposal {
    fn eq(&self, other: &Self) -> bool {
        self.hashed() == other.hashed()
    }
}

impl Eq for DataProposal {}

impl DataSized for DataProposal {
    fn estimate_size(&self) -> usize {
        self.txs.iter().map(|tx| tx.estimate_size()).sum()
    }
}

#[derive(
    Default,
    Serialize,
    Deserialize,
    Debug,
    Clone,
    PartialEq,
    Eq,
    Hash,
    Ord,
    PartialOrd,
    BorshDeserialize,
    BorshSerialize,
)]
#[cfg_attr(feature = "full", derive(utoipa::ToSchema))]
pub struct TxId(pub DataProposalHash, pub TxHash);

#[derive(
    Clone,
    Default,
    Serialize,
    Deserialize,
    BorshSerialize,
    BorshDeserialize,
    PartialEq,
    Eq,
    Hash,
    Ord,
    PartialOrd,
)]
#[cfg_attr(feature = "full", derive(utoipa::ToSchema))]
pub struct DataProposalHash(#[serde(with = "crate::utils::hex_bytes")] pub Vec<u8>);

impl From<Vec<u8>> for DataProposalHash {
    fn from(v: Vec<u8>) -> Self {
        DataProposalHash(v)
    }
}
impl From<&[u8]> for DataProposalHash {
    fn from(v: &[u8]) -> Self {
        DataProposalHash(v.to_vec())
    }
}
impl<const N: usize> From<&[u8; N]> for DataProposalHash {
    fn from(v: &[u8; N]) -> Self {
        DataProposalHash(v.to_vec())
    }
}
impl DataProposalHash {
    pub fn from_hex(s: &str) -> Result<Self, hex::FromHexError> {
        crate::utils::decode_hex_string_checked(s).map(DataProposalHash)
    }
}

impl std::fmt::Debug for DataProposalHash {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "DataProposalHash({})", hex::encode(&self.0))
    }
}

impl Hashed<DataProposalHash> for DataProposal {
    fn hashed(&self) -> DataProposalHash {
        if let Some(hash) = self.hash_cache.read().unwrap().as_ref() {
            return hash.clone();
        }
        let mut hasher = Sha3_256::new();
        match &self.parent_data_proposal_hash {
            DataProposalParent::LaneRoot(lane_id) => {
                hasher.update(lane_id.to_string().as_bytes());
            }
            DataProposalParent::DP(parent_data_proposal_hash) => {
                hasher.update(&parent_data_proposal_hash.0);
            }
        }
        for tx in self.txs.iter() {
            hasher.update(&tx.hashed().0);
        }
        let hash = DataProposalHash(hasher.finalize().to_vec());
        *self.hash_cache.write().unwrap() = Some(hash.clone());
        hash
    }
}

// Warning: hashing DPs can be slow, so use with care
impl std::hash::Hash for DataProposal {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.hashed().hash(state);
    }
}

impl Display for DataProposalHash {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", hex::encode(&self.0))
    }
}
impl Display for DataProposal {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.hashed())
    }
}

pub type PoDA = AggregateSignature;
pub type Cut = Vec<(LaneId, DataProposalHash, LaneBytesSize, PoDA)>;

impl Display for TxId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}/{}", self.0, self.1)
    }
}

// Can't impl-display, but we can still make it a little nicer by default
pub struct CutDisplay<'a>(pub &'a Cut);
impl Display for CutDisplay<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut cut_str = String::new();
        for (lane_id, hash, size, _) in self.0.iter() {
            cut_str.push_str(&format!("{lane_id}:{hash}({size}), "));
        }
        write!(f, "{}", cut_str.trim_end())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn data_proposal_hash_from_hex_str_roundtrip() {
        let hex_str = "746573745f6470";
        let hash = DataProposalHash::from_hex(hex_str).expect("data proposal hash hex");
        assert_eq!(hash.0, b"test_dp".to_vec());
        assert_eq!(format!("{hash}"), hex_str);
        let json = serde_json::to_string(&hash).expect("serialize data proposal hash");
        assert_eq!(json, "\"746573745f6470\"");
        let decoded: DataProposalHash =
            serde_json::from_str(&json).expect("deserialize data proposal hash");
        assert_eq!(decoded, hash);
    }
}