rialo-types 0.18.1

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! This module contains types used to describe the REX network.

use std::{fmt, mem::size_of};

use borsh::{BorshDeserialize, BorshSerialize};
use serde::{Deserialize, Serialize};

/// The representation of an epoch ID.
#[derive(
    BorshSerialize,
    BorshDeserialize,
    Clone,
    Copy,
    Debug,
    Default,
    Deserialize,
    Eq,
    Hash,
    Ord,
    PartialEq,
    PartialOrd,
    Serialize,
)]
#[serde(transparent)]
#[repr(transparent)]
pub struct RexEpochId(u64);

impl RexEpochId {
    /// The maximum size used when serializing an [`RexEpochId`] using [`borsh`].
    pub const BORSH_SERIALIZED_SIZE: usize = size_of::<u64>();

    /// Wraps a raw epoch number.
    pub const fn new(epoch: u64) -> Self {
        Self(epoch)
    }

    /// Returns the epoch as a `u64`.
    pub const fn as_u64(self) -> u64 {
        self.0
    }
}

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

/// Index of a committee in the REX configuration.
#[derive(
    BorshSerialize,
    BorshDeserialize,
    Clone,
    Copy,
    Debug,
    Default,
    Deserialize,
    Eq,
    Hash,
    Ord,
    PartialEq,
    PartialOrd,
    Serialize,
)]
pub struct CommitteeIndex(pub u32);

impl CommitteeIndex {
    /// The maximum size used when serializing a [`CommitteeIndex`] using [`borsh`].
    pub const BORSH_SERIALIZED_SIZE: usize = size_of::<u32>();

    /// Committee 0: the Standard REX committee.
    ///
    /// The TEE Registry orders the committees of a REX configuration by kind,
    /// in this fixed order; use these named indices instead of hardcoding
    /// `0`/`1`.
    pub const STANDARD_REX: Self = Self(0);

    /// Committee 1: the DKG / crypto-service committee.
    pub const CRYPTO_SERVICE: Self = Self(1);

    /// Returns the index as a `usize`, e.g. for indexing the committee `Vec`
    /// of a REX configuration.
    pub const fn as_usize(self) -> usize {
        self.0 as usize
    }
}

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

/// The REX committees established at genesis, by kind.
///
/// This is the closed, semantic view for code that targets a specific
/// committee: matching on it is exhaustive, and an invalid committee cannot
/// be constructed. [`CommitteeIndex`] remains the open wire representation
/// (a REX configuration's committee `Vec` is not capped by the type system),
/// so conversion from an index is fallible.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum GenesisCommittee {
    /// Committee 0: the Standard REX committee.
    StandardRex,
    /// Committee 1: the DKG / crypto-service committee.
    CryptoService,
}

impl GenesisCommittee {
    /// All committees established at genesis, in committee-index order.
    pub const ALL: [Self; 2] = [Self::StandardRex, Self::CryptoService];

    /// Number of REX committees established at genesis.
    pub const COUNT: usize = Self::ALL.len();

    /// The positional committee index this committee occupies in a REX
    /// configuration.
    pub const fn index(self) -> CommitteeIndex {
        match self {
            Self::StandardRex => CommitteeIndex::STANDARD_REX,
            Self::CryptoService => CommitteeIndex::CRYPTO_SERVICE,
        }
    }
}

impl From<GenesisCommittee> for CommitteeIndex {
    fn from(committee: GenesisCommittee) -> Self {
        committee.index()
    }
}

/// Error returned when a [`CommitteeIndex`] does not name a committee
/// established at genesis.
#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)]
#[error("committee index {0} does not name a genesis committee")]
pub struct UnknownGenesisCommittee(pub CommitteeIndex);

impl TryFrom<CommitteeIndex> for GenesisCommittee {
    type Error = UnknownGenesisCommittee;

    fn try_from(index: CommitteeIndex) -> Result<Self, Self::Error> {
        Self::ALL
            .into_iter()
            .find(|committee| committee.index() == index)
            .ok_or(UnknownGenesisCommittee(index))
    }
}

/// Index of a node within a REX committee.
#[derive(
    BorshSerialize,
    BorshDeserialize,
    Clone,
    Copy,
    Debug,
    Default,
    Deserialize,
    Eq,
    Ord,
    PartialEq,
    PartialOrd,
    Serialize,
)]
pub struct NodeIndex(pub u32);

impl NodeIndex {
    /// The maximum size used when serializing a [`NodeIndex`] using [`borsh`].
    pub const BORSH_SERIALIZED_SIZE: usize = size_of::<u32>();
}

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

/// The identifier to locate a REX node in a REX configuration.
#[derive(
    BorshSerialize,
    BorshDeserialize,
    Clone,
    Copy,
    Debug,
    Default,
    Deserialize,
    Eq,
    Ord,
    PartialEq,
    PartialOrd,
    Serialize,
)]
pub struct NodeId {
    pub epoch: RexEpochId,
    pub committee: CommitteeIndex,
    pub index: NodeIndex,
}

impl NodeId {
    /// The maximum size used when serializing an [`NodeId`] using [`borsh`].
    pub const BORSH_SERIALIZED_SIZE: usize = RexEpochId::BORSH_SERIALIZED_SIZE
        + CommitteeIndex::BORSH_SERIALIZED_SIZE
        + NodeIndex::BORSH_SERIALIZED_SIZE;
}

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

    #[test]
    fn test_genesis_committee_indices_round_trip() {
        for committee in GenesisCommittee::ALL {
            assert_eq!(GenesisCommittee::try_from(committee.index()), Ok(committee));
        }
    }

    #[test]
    fn test_genesis_committee_all_is_in_committee_index_order() {
        for (position, committee) in GenesisCommittee::ALL.into_iter().enumerate() {
            assert_eq!(committee.index().as_usize(), position);
        }
    }

    #[test]
    fn test_genesis_committee_rejects_index_beyond_the_genesis_committees() {
        let beyond = CommitteeIndex(GenesisCommittee::COUNT as u32);
        assert_eq!(
            GenesisCommittee::try_from(beyond),
            Err(UnknownGenesisCommittee(beyond))
        );
    }

    // `RexEpochId` is borsh-serialized on-chain (REX configuration, crypto-service
    // state) and serde-serialized on the wire. Pin that it is byte-identical to
    // the underlying `u64` in every encoding, so `#[serde(transparent)]` /
    // `#[repr(transparent)]` leave existing account and wire bytes unchanged.
    #[test]
    fn test_epoch_id_is_byte_identical_to_u64() {
        let raw: u64 = 0x00ad_beef_dead_beef;
        let e = RexEpochId::new(raw);
        assert_eq!(borsh::to_vec(&e).unwrap(), raw.to_le_bytes());
        assert_eq!(RexEpochId::BORSH_SERIALIZED_SIZE, raw.to_le_bytes().len());
        assert_eq!(bcs::to_bytes(&e).unwrap(), bcs::to_bytes(&raw).unwrap());
        assert_eq!(
            bincode::serialize(&e).unwrap(),
            bincode::serialize(&raw).unwrap()
        );
        assert_eq!(serde_json::to_string(&RexEpochId::new(7)).unwrap(), "7");
        assert_eq!(
            serde_json::from_str::<RexEpochId>("7").unwrap(),
            RexEpochId::new(7)
        );
    }

    // `committee_index` is part of the `DBMap<ValidatorRexScoreKey, _>` key
    // (typed-store encodes keys as bincode big-endian fixint) and of the
    // BCS-encoded `WriteSet::validator_rex_scores` value. Pin that the newtype
    // is byte-identical to the underlying `u32` in every encoding, so wrapping
    // the field leaves existing on-disk keys and serialized values unchanged.
    #[test]
    fn test_committee_index_is_byte_identical_to_u32() {
        let raw: u32 = 0xdead_beef;
        let c = CommitteeIndex(raw);
        assert_eq!(borsh::to_vec(&c).unwrap(), raw.to_le_bytes());
        assert_eq!(
            CommitteeIndex::BORSH_SERIALIZED_SIZE,
            raw.to_le_bytes().len()
        );
        assert_eq!(bcs::to_bytes(&c).unwrap(), bcs::to_bytes(&raw).unwrap());
        assert_eq!(
            bincode::serialize(&c).unwrap(),
            bincode::serialize(&raw).unwrap()
        );
        assert_eq!(serde_json::to_string(&CommitteeIndex(7)).unwrap(), "7");
        assert_eq!(
            serde_json::from_str::<CommitteeIndex>("7").unwrap(),
            CommitteeIndex(7)
        );
    }

    // `CommitteeIndex::Ord` orders the in-memory `BTreeSet`
    // `modified_validator_rex_scores`, which fixes the entry order of the
    // serialized `WriteSet::validator_rex_scores` Vec — this must stay
    // deterministic across nodes. (Not the RocksDB key order, which is
    // byte-lexicographic and covered by the byte-identity test above.)
    #[test]
    fn test_committee_index_ord_matches_u32() {
        let raws: [u32; 5] = [0, 1, 2, 100, u32::MAX];
        for &a in &raws {
            for &b in &raws {
                assert_eq!(
                    CommitteeIndex(a).cmp(&CommitteeIndex(b)),
                    a.cmp(&b),
                    "Ord mismatch for {a} vs {b}"
                );
            }
        }
    }
}