rialo-feature-management-program-interface 0.3.0-alpha.0

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

//! Feature state management

extern crate alloc;

use alloc::{
    collections::BTreeMap,
    string::{String, ToString},
    vec::Vec,
};

use borsh::{BorshDeserialize, BorshSerialize};
use rialo_s_pubkey::Pubkey;

/// The program's global state
///
/// This is the top-level structure stored in the storage account.
#[derive(Clone, Debug, PartialEq, BorshSerialize, BorshDeserialize)]
pub struct FeaturesState {
    /// Authority pubkey
    authority: Pubkey,
    /// The features map
    pub features_map: BTreeMap<String, (u64, u64)>, // feature name -> (start_time, end_time)
}

#[cfg(test)]
pub const DETERMINISTIC_TEST_KEYPAIR: &str =
    "57Vqb7tHij5NhQnTgrgYXA19pC8ZVHQoCHpapSiQ8LJaeUvTcSBzKoB6CazhR6VtxmyVAbWnoeDSzD1Vm672NaKp";

impl FeaturesState {
    /// Create a new state with the given authority
    pub fn new(authority: Pubkey) -> Self {
        Self {
            authority,
            features_map: BTreeMap::new(),
        }
    }

    #[cfg(test)]
    pub fn new_for_test() -> Self {
        use rialo_s_keypair::Keypair;
        use rialo_s_signer::Signer;

        let deterministic_test_pubkey: Pubkey =
            Keypair::from_base58_string(DETERMINISTIC_TEST_KEYPAIR)
                .try_pubkey()
                .expect("Failed to get pubkey from deterministic test keypair");
        Self {
            authority: deterministic_test_pubkey,
            features_map: BTreeMap::new(),
        }
    }

    pub fn get_authority(&self) -> &rialo_s_pubkey::Pubkey {
        &self.authority
    }

    pub fn set_authority(&mut self, new_authority: rialo_s_pubkey::Pubkey) {
        self.authority = new_authority;
    }

    /// Serialize the state
    pub fn serialize(&self) -> Result<Vec<u8>, borsh::io::Error> {
        borsh::to_vec(self)
    }

    /// Deserialize the state
    pub fn deserialize(data: &[u8]) -> Result<Self, borsh::io::Error> {
        borsh::from_slice(data)
    }

    /// Check if the feature is currently active
    pub fn is_active(&self, feature_name: &str, current_time: u64) -> bool {
        if let Some((start_time, end_time)) = self.features_map.get(feature_name) {
            current_time >= *start_time && current_time < *end_time
        } else {
            false
        }
    }

    /// Add or update a feature
    pub fn upsert(&mut self, name: String, start_time: u64, end_time: u64) -> Result<(), String> {
        if start_time >= end_time {
            return Err("Invalid time range: start_time must be less than end_time".to_string());
        }

        // Check if adding a new feature would exceed the maximum count
        if !self.features_map.contains_key(&name)
            && self.features_map.len() >= crate::MAX_FEATURE_COUNT
        {
            return Err(alloc::format!(
                "Maximum feature count ({}) exceeded",
                crate::MAX_FEATURE_COUNT
            ));
        }

        self.features_map.insert(name, (start_time, end_time));
        Ok(())
    }

    /// Get a feature's time range in (start, end) format
    pub fn get(&self, name: &str) -> Option<&(u64, u64)> {
        self.features_map.get(name)
    }
}

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

    #[test]
    fn test_new_state() {
        let authority = Pubkey::new_unique();
        let state = FeaturesState::new(authority);

        assert_eq!(state.get_authority(), &authority);
        assert!(state.features_map.is_empty());
    }

    #[test]
    fn test_upsert_feature() {
        let mut state = FeaturesState::new_for_test();

        let result = state.upsert("feature1".to_string(), 100, 200);
        assert!(result.is_ok());
        assert_eq!(state.get("feature1"), Some(&(100, 200)));
    }

    #[test]
    fn test_upsert_invalid_time_range() {
        let mut state = FeaturesState::new_for_test();

        // start_time >= end_time should fail
        let result = state.upsert("feature1".to_string(), 200, 100);
        assert!(result.is_err());

        let result = state.upsert("feature2".to_string(), 100, 100);
        assert!(result.is_err());
    }

    #[test]
    fn test_is_active() {
        let mut state = FeaturesState::new_for_test();
        state.upsert("feature1".to_string(), 100, 200).unwrap();

        // Before start time
        assert!(!state.is_active("feature1", 50));

        // At start time
        assert!(state.is_active("feature1", 100));

        // During active period
        assert!(state.is_active("feature1", 150));

        // At end time (should be inactive)
        assert!(!state.is_active("feature1", 200));

        // After end time
        assert!(!state.is_active("feature1", 250));

        // Non-existent feature
        assert!(!state.is_active("nonexistent", 150));
    }

    #[test]
    fn test_serialize_deserialize() {
        let mut state = FeaturesState::new_for_test();
        state.upsert("feature1".to_string(), 100, 200).unwrap();
        state.upsert("feature2".to_string(), 300, 400).unwrap();

        let serialized = state.serialize().expect("Serialization failed");
        let deserialized = FeaturesState::deserialize(&serialized).expect("Deserialization failed");

        assert_eq!(state, deserialized);
    }

    #[test]
    fn test_set_authority() {
        let mut state = FeaturesState::new_for_test();
        let new_authority = Pubkey::new_unique();

        state.set_authority(new_authority);
        assert_eq!(state.get_authority(), &new_authority);
    }
}