Alligned with master

resources/invariant_schema.json

@@ -123,6 +123,64 @@
         "List": "Any"
       }
     },
+    {
+      "name": "get_protocol_fee",
+      "is_mutable": false,
+      "args": [],
+      "return_ty": "Any"
+    },
+    {
+      "name": "withdraw_protocol_fee",
+      "is_mutable": true,
+      "args": [
+        {
+          "name": "pool_key",
+          "ty": "Any"
+        }
+      ],
+      "return_ty": {
+        "Result": {
+          "ok": "Unit",
+          "err": "U32"
+        }
+      }
+    },
+    {
+      "name": "change_protocol_fee",
+      "is_mutable": true,
+      "args": [
+        {
+          "name": "protocol_fee",
+          "ty": "Any"
+        }
+      ],
+      "return_ty": {
+        "Result": {
+          "ok": "Unit",
+          "err": "U32"
+        }
+      }
+    },
+    {
+      "name": "change_fee_receiver",
+      "is_mutable": true,
+      "args": [
+        {
+          "name": "pool_key",
+          "ty": "Any"
+        },
+        {
+          "name": "fee_receiver",
+          "ty": "Key"
+        }
+      ],
+      "return_ty": {
+        "Result": {
+          "ok": "Unit",
+          "err": "U32"
+        }
+      }
+    },
     {
       "name": "is_tick_initialized",
       "is_mutable": false,

src/contracts/entrypoints.rs

@@ -1,6 +1,9 @@
 use super::{FeeTier, Pool, PoolKey, Position, Tick};
 use crate::{
-    math::{liquidity::Liquidity, sqrt_price::SqrtPrice, token_amount::TokenAmount},
+    math::{
+        liquidity::Liquidity, percentage::Percentage, sqrt_price::SqrtPrice,
+        token_amount::TokenAmount,
+    },
     InvariantError,
 };
 use odra::{prelude::vec::Vec, types::Address};
@@ -26,8 +29,16 @@ pub trait Entrypoints {
     ) -> Result<Pool, InvariantError>;
     fn get_pools(&self) -> Vec<PoolKey>;
 
-    fn is_tick_initialized(&self, key: PoolKey, index: i32) -> bool;
+    fn get_protocol_fee(&self) -> Percentage;
+    fn withdraw_protocol_fee(&mut self, pool_key: PoolKey) -> Result<(), InvariantError>;
+    fn change_protocol_fee(&mut self, protocol_fee: Percentage) -> Result<(), InvariantError>;
+    fn change_fee_receiver(
+        &mut self,
+        pool_key: PoolKey,
+        fee_receiver: Address,
+    ) -> Result<(), InvariantError>;
 
+    fn is_tick_initialized(&self, key: PoolKey, index: i32) -> bool;
     fn get_tick(&self, key: PoolKey, index: i32) -> Result<Tick, InvariantError>;
 
     fn claim_fee(&mut self, index: u32) -> Result<(TokenAmount, TokenAmount), InvariantError>;

src/e2e/change_fee_receiver.rs

@@ -0,0 +1,67 @@
+use crate::math::percentage::Percentage;
+use crate::InvariantDeployer;
+use crate::InvariantError;
+use crate::{FeeTier, PoolKey};
+use decimal::*;
+use odra::test_env;
+use odra::types::casper_types::ContractPackageHash;
+use odra::types::Address;
+use odra::types::U128;
+
+#[test]
+fn test_change_fee_reciever() {
+    let token_0 = Address::Contract(ContractPackageHash::from([0x01; 32]));
+    let token_1 = Address::Contract(ContractPackageHash::from([0x02; 32]));
+
+    let deployer = test_env::get_account(0);
+    test_env::set_caller(deployer);
+    let mut invariant = InvariantDeployer::init(Percentage::new(U128::from(0)));
+
+    let fee_tier = FeeTier::new(Percentage::from_scale(5, 1), 1).unwrap();
+    invariant.add_fee_tier(fee_tier).unwrap();
+
+    let exist = invariant.fee_tier_exist(fee_tier);
+    assert!(exist);
+
+    let init_tick = 0;
+    invariant
+        .create_pool(token_0, token_1, fee_tier, init_tick)
+        .unwrap();
+
+    let new_receiver = test_env::get_account(1);
+    let pool_key = PoolKey::new(token_0, token_1, fee_tier).unwrap();
+
+    invariant
+        .change_fee_receiver(pool_key, new_receiver)
+        .unwrap();
+
+    let pool = invariant.get_pool(token_0, token_1, fee_tier).unwrap();
+    assert_eq!(pool.fee_receiver, new_receiver);
+}
+
+#[test]
+fn test_not_admin_change_fee_reciever() {
+    let token_0 = Address::Contract(ContractPackageHash::from([0x01; 32]));
+    let token_1 = Address::Contract(ContractPackageHash::from([0x02; 32]));
+
+    let deployer = test_env::get_account(0);
+    test_env::set_caller(deployer);
+    let mut invariant = InvariantDeployer::init(Percentage::new(U128::from(0)));
+
+    let fee_tier = FeeTier::new(Percentage::from_scale(5, 1), 1).unwrap();
+    invariant.add_fee_tier(fee_tier).unwrap();
+
+    let exist = invariant.fee_tier_exist(fee_tier);
+    assert!(exist);
+
+    let init_tick = 0;
+    invariant
+        .create_pool(token_0, token_1, fee_tier, init_tick)
+        .unwrap();
+
+    let new_receiver = test_env::get_account(1);
+    let pool_key = PoolKey::new(token_0, token_1, fee_tier).unwrap();
+    test_env::set_caller(new_receiver);
+    let result = invariant.change_fee_receiver(pool_key, new_receiver);
+    assert_eq!(result, Err(InvariantError::NotAdmin));
+}

src/e2e/change_protocol_fee.rs

@@ -0,0 +1,38 @@
+use crate::math::percentage::Percentage;
+use crate::InvariantDeployer;
+use crate::InvariantError;
+use decimal::*;
+use odra::test_env;
+use odra::types::U128;
+
+#[test]
+fn test_change_protocol_fee() {
+    let deployer = test_env::get_account(0);
+    test_env::set_caller(deployer);
+    let mut invariant = InvariantDeployer::init(Percentage::new(U128::from(0)));
+
+    let protocol_fee = invariant.get_protocol_fee();
+    assert_eq!(protocol_fee, Percentage::new(U128::from(0)));
+
+    let new_fee = Percentage::new(U128::from(1));
+    invariant.change_protocol_fee(new_fee).unwrap();
+
+    let protocol_fee = invariant.get_protocol_fee();
+    assert_eq!(protocol_fee, new_fee);
+}
+
+#[test]
+fn test_change_protocol_fee_not_admin() {
+    let deployer = test_env::get_account(0);
+    test_env::set_caller(deployer);
+    let mut invariant = InvariantDeployer::init(Percentage::new(U128::from(0)));
+
+    let protocol_fee = invariant.get_protocol_fee();
+    assert_eq!(protocol_fee, Percentage::new(U128::from(0)));
+
+    let new_fee = Percentage::new(U128::from(1));
+    test_env::set_caller(test_env::get_account(1));
+    let result = invariant.change_protocol_fee(new_fee);
+
+    assert_eq!(result, Err(InvariantError::NotAdmin));
+}

src/e2e/mod.rs

@@ -1,4 +1,6 @@
 pub mod add_fee_tier;
+pub mod change_fee_receiver;
+pub mod change_protocol_fee;
 pub mod constructor;
 pub mod create_pool;
 pub mod remove_fee_tier;

src/lib.rs

@@ -5,6 +5,8 @@ extern crate alloc;
 pub mod contracts;
 pub mod math;
 
+use odra_modules::erc20::Erc20Ref;
+
 #[cfg(test)]
 pub mod e2e;
 
@@ -21,7 +23,6 @@ use odra::types::event::OdraEvent;
 use odra::types::{Address, U256};
 use odra::{contract_env, Event};
 use odra::{OdraType, UnwrapOrRevert, Variable};
-use odra_modules::erc20::Erc20Ref;
 
 #[derive(OdraType, Debug, PartialEq)]
 pub enum InvariantError {
@@ -267,6 +268,63 @@ impl Entrypoints for Invariant {
         self.pool_keys.get().unwrap_or_revert().get_all()
     }
 
+    pub fn get_protocol_fee(&self) -> Percentage {
+        let state = self.state.get().unwrap_or_revert();
+        state.protocol_fee
+    }
+
+    pub fn withdraw_protocol_fee(&mut self, pool_key: PoolKey) -> Result<(), InvariantError> {
+        let caller = contract_env::caller();
+        let mut pool = self.pools.get(pool_key)?;
+
+        if caller != pool.fee_receiver {
+            return Err(InvariantError::NotFeeReceiver);
+        }
+
+        let (fee_protocol_token_x, fee_protocol_token_y) = pool.withdraw_protocol_fee(pool_key);
+
+        Erc20Ref::at(&pool_key.token_x).transfer(&pool.fee_receiver, &fee_protocol_token_x.get());
+        Erc20Ref::at(&pool_key.token_y).transfer(&pool.fee_receiver, &fee_protocol_token_y.get());
+
+        self.pools.update(pool_key, &pool)?;
+
+        Ok(())
+    }
+
+    pub fn change_protocol_fee(&mut self, protocol_fee: Percentage) -> Result<(), InvariantError> {
+        let caller = contract_env::caller();
+        let mut state = self.state.get().unwrap_or_revert();
+
+        if caller != state.admin {
+            return Err(InvariantError::NotAdmin);
+        }
+
+        state.protocol_fee = protocol_fee;
+
+        self.state.set(state);
+
+        Ok(())
+    }
+
+    pub fn change_fee_receiver(
+        &mut self,
+        pool_key: PoolKey,
+        fee_receiver: Address,
+    ) -> Result<(), InvariantError> {
+        let caller = contract_env::caller();
+        let state = self.state.get().unwrap_or_revert();
+        let mut pool = self.pools.get(pool_key)?;
+
+        if caller != state.admin {
+            return Err(InvariantError::NotAdmin);
+        }
+
+        pool.fee_receiver = fee_receiver;
+        self.pools.update(pool_key, &pool)?;
+
+        Ok(())
+    }
+
     pub fn is_tick_initialized(&self, key: PoolKey, index: i32) -> bool {
         self.tickmap.get(index, key.fee_tier.tick_spacing, key)
     }