rialo-types 0.12.2

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

use std::collections::BTreeMap;

use serde::{Deserialize, Serialize};

use crate::{
    rex_info::{RexId, RexUpdateResult, TimestampMs},
    websocket::BLS12381_PUBLIC_KEY_SIZE,
    RexDutyConfig,
};

/// Raw 96-byte authority public key (BLS12-381).
/// Used to identify which authority submitted a particular update.
/// The authority key represents the identity of an authority in the committee.
pub type AuthorityKeyBytes = [u8; BLS12381_PUBLIC_KEY_SIZE];

/// Monotonically increasing consensus round/commit index.
///
/// Used throughout the REX pipeline to schedule, assign, and finalize duties.
pub type CommitIndex = u64;

/// Map structure for tracking REX duties and their execution status.
///
/// The map is keyed by `RexDutiesMapKey`.
///
/// The value is `RexDutiesMapValue` which tracks validator assignments and their update results.
///
/// This structure enables efficient tracking of REX duties ordered by round priority,
/// with the most recent rounds processed first.
pub type RexDutiesMap = BTreeMap<RexDutiesMapKey, RexDutiesMapValue>;

/// Key type for the REX duties map, consisting of a target commit and REX identifier.
#[derive(Clone, Copy, Debug, Serialize, Deserialize, Eq, PartialEq, Hash)]
pub struct RexDutiesMapKey {
    pub target_timestamp: TimestampMs,
    pub rex_id: RexId,
}

// Explicit implementation of `Ord` to make sure `target_timestamp` remains most important in the key comparison.
impl Ord for RexDutiesMapKey {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.target_timestamp
            .cmp(&other.target_timestamp)
            .then_with(|| self.rex_id.cmp(&other.rex_id))
    }
}

impl PartialOrd for RexDutiesMapKey {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

/// Value type for the REX duties map, tracking validator assignments and their update results.
#[derive(Clone, Debug, Serialize, Deserialize, Default, PartialEq, Eq)]
pub struct RexDutiesMapValue {
    /// A map of authority key bytes to their REX update results.
    /// - Key: The validator's 96-byte authority public key (BLS12-381)
    /// - Value: The result of the validator's REX update,
    ///   where `None` indicates the duty is still pending and `Some(result)` indicates completion
    #[serde(with = "serde_utils")]
    pub updates: BTreeMap<AuthorityKeyBytes, Option<RexUpdateResult>>,

    /// The configuration parameters that define how this REX duty is scheduled and finalized,
    /// including request delay, collection window size, and commitment buffer.
    pub config: RexDutyConfig,
}

// Helper module for serializing BTreeMap with array keys
mod serde_utils {
    use std::collections::BTreeMap;

    use serde::{Deserialize, Deserializer, Serialize, Serializer};
    use serde_big_array::BigArray;

    use crate::{rex_info::RexUpdateResult, AuthorityKeyBytes};

    pub fn serialize<S>(
        map: &BTreeMap<AuthorityKeyBytes, Option<RexUpdateResult>>,
        serializer: S,
    ) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        #[derive(Serialize)]
        struct Entry {
            #[serde(with = "BigArray")]
            key: AuthorityKeyBytes,
            value: Option<RexUpdateResult>,
        }

        let entries: Vec<_> = map
            .iter()
            .map(|(key, value)| Entry {
                key: *key,
                value: value.clone(),
            })
            .collect();
        entries.serialize(serializer)
    }

    pub fn deserialize<'de, D>(
        deserializer: D,
    ) -> Result<BTreeMap<AuthorityKeyBytes, Option<RexUpdateResult>>, D::Error>
    where
        D: Deserializer<'de>,
    {
        #[derive(Deserialize)]
        struct Entry {
            #[serde(with = "BigArray")]
            key: AuthorityKeyBytes,
            value: Option<RexUpdateResult>,
        }

        let entries = Vec::<Entry>::deserialize(deserializer)?;
        Ok(entries.into_iter().map(|e| (e.key, e.value)).collect())
    }
}

#[cfg(test)]
mod tests {
    use std::cmp::{max, min};

    use rialo_s_pubkey::Pubkey;

    use super::*;

    #[test]
    fn test_rex_duties_map_key_ordering() {
        let creator1 = Pubkey::new_unique();
        let creator2 = Pubkey::new_unique();
        let rex_id1 = RexId::new(creator1, "nonce1");
        let rex_id2 = RexId::new(creator2, "nonce2");
        // Ensure rex_id1 < rex_id2.
        let [rex_id1, rex_id2] = [min(rex_id1, rex_id2), max(rex_id1, rex_id2)];
        // Test ordering by target_timestamp
        let key1 = RexDutiesMapKey {
            target_timestamp: 100,
            rex_id: rex_id1,
        };
        let key2 = RexDutiesMapKey {
            target_timestamp: 200,
            rex_id: rex_id1,
        };
        assert!(key1 < key2);
        assert!(key2 > key1);

        // Test ordering by rex_id when target_timestamp is equal
        let key3 = RexDutiesMapKey {
            target_timestamp: 100,
            rex_id: rex_id1,
        };
        let key4 = RexDutiesMapKey {
            target_timestamp: 100,
            rex_id: rex_id2,
        };
        assert!(key3 < key4);
        assert!(key4 > key3);

        let key5 = RexDutiesMapKey {
            target_timestamp: 100,
            rex_id: rex_id2,
        };
        let key6 = RexDutiesMapKey {
            target_timestamp: 200,
            rex_id: rex_id1,
        };
        assert!(key5 < key6);
        assert!(key6 > key5);

        // Test equality
        let key7 = RexDutiesMapKey {
            target_timestamp: 100,
            rex_id: rex_id1,
        };
        let key8 = RexDutiesMapKey {
            target_timestamp: 100,
            rex_id: rex_id1,
        };
        assert_eq!(key7, key8);
        assert!(key7 >= key8);
        assert!(key7 <= key8);

        // Test ordering in BTreeMap
        let mut map = BTreeMap::new();
        let key9 = RexDutiesMapKey {
            target_timestamp: 200,
            rex_id: rex_id1,
        };
        let key10 = RexDutiesMapKey {
            target_timestamp: 100,
            rex_id: rex_id2,
        };
        let key11 = RexDutiesMapKey {
            target_timestamp: 100,
            rex_id: rex_id1,
        };

        map.insert(key9, "first");
        map.insert(key10, "second");
        map.insert(key11, "third");

        let keys: Vec<_> = map.keys().copied().collect();
        assert_eq!(keys[0], key11);
        assert_eq!(keys[1], key10);
        assert_eq!(keys[2], key9);
    }
}