rialo-types 0.1.8

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

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

use serde::{Deserialize, Serialize};

use crate::{OracleDutyConfig, OracleId, OracleUpdateResult};

/// Raw 32-byte public key for an authority (a set of validator nodes).
/// Used to identify which authority submitted a particular update.
pub type ProtocolKeyBytes = [u8; 32];

/// 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 proposal round and oracle identifier and retries made.
#[derive(Clone, Copy, Debug, Serialize, Deserialize, Eq, PartialEq, Hash)]
pub struct OracleDutiesMapKey {
    pub target_commit: u64,
    pub oracle_id: OracleId,
}

// Explicit implementation of `Ord` to make sure `target_commit` remains most important in the key comparison.
impl Ord for OracleDutiesMapKey {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.target_commit
            .cmp(&other.target_commit)
            .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 {
    /// This is a `HashMap` where:
    /// - Key (`ProtocolKeyBytes`): The validator's 32-byte protocol public key
    /// - Value (`Option<OracleUpdateResult>`): The result of the validator's oracle update,
    ///   where `None` indicates the duty is still pending and `Some(result)` indicates completion
    pub updates: HashMap<ProtocolKeyBytes, 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,

    /// Which retry attempt this duty is on.
    /// The value is zero for the first attempt and gets increasing with each attempt.
    pub retries_made: u32,
}

#[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_commit
        let key1 = OracleDutiesMapKey {
            target_commit: 100,
            oracle_id: oracle_id1,
        };
        let key2 = OracleDutiesMapKey {
            target_commit: 200,
            oracle_id: oracle_id1,
        };
        assert!(key1 < key2);
        assert!(key2 > key1);

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

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

        // Test equality
        let key7 = OracleDutiesMapKey {
            target_commit: 100,
            oracle_id: oracle_id1,
        };
        let key8 = OracleDutiesMapKey {
            target_commit: 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_commit: 200,
            oracle_id: oracle_id1,
        };
        let key10 = OracleDutiesMapKey {
            target_commit: 100,
            oracle_id: oracle_id2,
        };
        let key11 = OracleDutiesMapKey {
            target_commit: 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);
    }
}