rialo-types 0.2.0

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

use std::collections::{BTreeMap, HashMap};

use serde::{Deserialize, Serialize};

use crate::{
    websocket::BLS12381_PUBLIC_KEY_SIZE, OracleDutyConfig, OracleId, OracleUpdateResult,
    TimestampMs,
};

/// 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 oracle pipeline to schedule, assign, and finalize duties.
pub type CommitIndex = u64;

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

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

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

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

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

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

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

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

    use super::{AuthorityKeyBytes, OracleUpdateResult};

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

        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<HashMap<AuthorityKeyBytes, Option<OracleUpdateResult>>, D::Error>
    where
        D: Deserializer<'de>,
    {
        #[derive(Deserialize)]
        struct Entry {
            #[serde(with = "BigArray")]
            key: AuthorityKeyBytes,
            value: Option<OracleUpdateResult>,
        }

        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_oracle_duties_map_key_ordering() {
        let creator1 = Pubkey::new_unique();
        let creator2 = Pubkey::new_unique();
        let oracle_id1 = OracleId::new(creator1, "nonce1");
        let oracle_id2 = OracleId::new(creator2, "nonce2");
        // Ensure oracle_id1 < oracle_id2.
        let [oracle_id1, oracle_id2] = [min(oracle_id1, oracle_id2), max(oracle_id1, oracle_id2)];
        // Test ordering by target_timestamp
        let key1 = OracleDutiesMapKey {
            target_timestamp: 100,
            oracle_id: oracle_id1,
        };
        let key2 = OracleDutiesMapKey {
            target_timestamp: 200,
            oracle_id: oracle_id1,
        };
        assert!(key1 < key2);
        assert!(key2 > key1);

        // Test ordering by oracle_id when target_timestamp is equal
        let key3 = OracleDutiesMapKey {
            target_timestamp: 100,
            oracle_id: oracle_id1,
        };
        let key4 = OracleDutiesMapKey {
            target_timestamp: 100,
            oracle_id: oracle_id2,
        };
        assert!(key3 < key4);
        assert!(key4 > key3);

        let key5 = OracleDutiesMapKey {
            target_timestamp: 100,
            oracle_id: oracle_id2,
        };
        let key6 = OracleDutiesMapKey {
            target_timestamp: 200,
            oracle_id: oracle_id1,
        };
        assert!(key5 < key6);
        assert!(key6 > key5);

        // Test equality
        let key7 = OracleDutiesMapKey {
            target_timestamp: 100,
            oracle_id: oracle_id1,
        };
        let key8 = OracleDutiesMapKey {
            target_timestamp: 100,
            oracle_id: oracle_id1,
        };
        assert_eq!(key7, key8);
        assert!(key7 >= key8);
        assert!(key7 <= key8);

        // Test ordering in BTreeMap
        let mut map = BTreeMap::new();
        let key9 = OracleDutiesMapKey {
            target_timestamp: 200,
            oracle_id: oracle_id1,
        };
        let key10 = OracleDutiesMapKey {
            target_timestamp: 100,
            oracle_id: oracle_id2,
        };
        let key11 = OracleDutiesMapKey {
            target_timestamp: 100,
            oracle_id: oracle_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);
    }
}