feat: implement treasury contract (#994)

Cargo.toml

@@ -1,5 +1,5 @@
 [workspace]
-members = ["creator-keys", "creator-keys-factory", "referral-contract", "lp-reward-contract"]
+members = ["creator-keys", "creator-keys-factory", "referral-contract", "treasury-contract"]
 resolver = "2"
 
 [workspace.package]

treasury-contract/Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "treasury-contract"
+version = "0.1.0"
+edition = "2021"
+publish = false
+
+[lib]
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+soroban-sdk = { workspace = true }
+
+[dev-dependencies]
+soroban-sdk = { workspace = true, features = ["testutils"] }

treasury-contract/src/lib.rs

@@ -0,0 +1,105 @@
+#![no_std]
+
+use soroban_sdk::{contract, contractimpl, contracttype, Address, Env, String, Symbol, Vec};
+
+#[contracttype]
+#[derive(Clone)]
+pub enum DataKey {
+    Admin,
+    Balance,
+    Epoch,
+    EpochHistory(u32),
+}
+
+#[contracttype]
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct EpochDistribution {
+    pub epoch: u32,
+    pub recipients: Vec<Address>,
+    pub amounts: Vec<i128>,
+}
+
+#[contract]
+pub struct TreasuryContract;
+
+#[contractimpl]
+impl TreasuryContract {
+    pub fn initialize(env: Env, admin: Address) {
+        if env.storage().instance().has(&DataKey::Admin) {
+            panic!("Already initialized");
+        }
+        env.storage().instance().set(&DataKey::Admin, &admin);
+        env.storage().instance().set(&DataKey::Balance, &0i128);
+        env.storage().instance().set(&DataKey::Epoch, &0u32);
+    }
+
+    pub fn collect_fee(env: Env, key_id: String, amount: i128) {
+        if amount <= 0 {
+            panic!("Amount must be positive");
+        }
+
+        let mut balance: i128 = env.storage().instance().get(&DataKey::Balance).unwrap_or(0);
+        balance += amount;
+        env.storage().instance().set(&DataKey::Balance, &balance);
+
+        env.events()
+            .publish((Symbol::new(&env, "FeeCollected"), key_id), amount);
+    }
+
+    pub fn get_treasury_balance(env: Env) -> i128 {
+        env.storage().instance().get(&DataKey::Balance).unwrap_or(0)
+    }
+
+    pub fn distribute(env: Env, recipients: Vec<Address>, amounts: Vec<i128>) {
+        let admin: Address = env
+            .storage()
+            .instance()
+            .get(&DataKey::Admin)
+            .expect("Not initialized");
+        admin.require_auth();
+
+        if recipients.len() != amounts.len() {
+            panic!("Recipients and amounts length mismatch");
+        }
+
+        let mut total_amount: i128 = 0;
+        for amount in amounts.iter() {
+            if amount <= 0 {
+                panic!("Amount must be positive");
+            }
+            total_amount += amount;
+        }
+
+        let mut balance: i128 = env.storage().instance().get(&DataKey::Balance).unwrap_or(0);
+        if total_amount > balance {
+            panic!("Insufficient balance");
+        }
+
+        balance -= total_amount;
+        env.storage().instance().set(&DataKey::Balance, &balance);
+
+        let epoch: u32 = env.storage().instance().get(&DataKey::Epoch).unwrap_or(0);
+        let distribution = EpochDistribution {
+            epoch,
+            recipients: recipients.clone(),
+            amounts: amounts.clone(),
+        };
+
+        env.storage()
+            .persistent()
+            .set(&DataKey::EpochHistory(epoch), &distribution);
+        env.storage().instance().set(&DataKey::Epoch, &(epoch + 1));
+
+        env.events()
+            .publish((Symbol::new(&env, "FeeDistributed"), epoch), distribution);
+    }
+
+    pub fn get_distribution_history(env: Env, epoch: u32) -> Option<EpochDistribution> {
+        env.storage()
+            .persistent()
+            .get(&DataKey::EpochHistory(epoch))
+    }
+}
+
+#[cfg(test)]
+mod test;

treasury-contract/src/test.rs

@@ -0,0 +1,63 @@
+#![cfg(test)]
+
+use super::*;
+use soroban_sdk::{testutils::Address as _, vec, Address, Env, String};
+
+#[test]
+fn test_treasury() {
+    let env = Env::default();
+    let contract_id = env.register(TreasuryContract, ());
+    let client = TreasuryContractClient::new(&env, &contract_id);
+
+    let admin = Address::generate(&env);
+    client.initialize(&admin);
+
+    let key_id = String::from_str(&env, "key123");
+
+    // Collect fee
+    client.collect_fee(&key_id, &1000);
+    assert_eq!(client.get_treasury_balance(), 1000);
+
+    client.collect_fee(&key_id, &500);
+    assert_eq!(client.get_treasury_balance(), 1500);
+
+    // Distribute
+    let recipient1 = Address::generate(&env);
+    let recipient2 = Address::generate(&env);
+    let recipients = vec![&env, recipient1.clone(), recipient2.clone()];
+    let amounts = vec![&env, 700, 300];
+
+    // Mock admin auth
+    env.mock_all_auths();
+
+    client.distribute(&recipients, &amounts);
+
+    assert_eq!(client.get_treasury_balance(), 500);
+
+    let history = client.get_distribution_history(&0).unwrap();
+    assert_eq!(history.epoch, 0);
+    assert_eq!(history.recipients, recipients);
+    assert_eq!(history.amounts, amounts);
+}
+
+#[test]
+#[should_panic(expected = "Insufficient balance")]
+fn test_distribute_insufficient_balance() {
+    let env = Env::default();
+    let contract_id = env.register(TreasuryContract, ());
+    let client = TreasuryContractClient::new(&env, &contract_id);
+
+    let admin = Address::generate(&env);
+    client.initialize(&admin);
+
+    let key_id = String::from_str(&env, "key123");
+    client.collect_fee(&key_id, &100);
+
+    let recipient = Address::generate(&env);
+    let recipients = vec![&env, recipient];
+    let amounts = vec![&env, 200];
+
+    env.mock_all_auths();
+
+    client.distribute(&recipients, &amounts);
+}