[MEX-592] Distribute only mex in fees collector

dex/pair/tests/pair_rs_test.rs

@@ -1600,6 +1600,8 @@ fn fees_collector_pair_test() {
                 sc.init(
                     managed_token_id!(LOCKED_TOKEN_ID),
                     managed_address!(&energy_factory_mock_addr),
+                    managed_address!(&energy_factory_mock_addr), // unused
+                    managed_token_id!(b"RANDTOK-123456"),        // unused
                     MultiValueEncoded::new(),
                 );
                 let _ = sc.known_contracts().insert(managed_address!(&pair_addr));

dex/router/Cargo.toml

@@ -29,3 +29,6 @@ path = "../../locked-asset/simple-lock"
 
 [dev-dependencies.multiversx-sc-scenario]
 version = "=0.53.2"
+
+[dev-dependencies.fees-collector]
+path = "../../energy-integration/fees-collector"

dex/router/src/factory.rs

@@ -132,24 +132,20 @@ pub trait FactoryModule: config::ConfigModule + read_pair_storage::ReadPairStora
         first_token_id: TokenIdentifier,
         second_token_id: TokenIdentifier,
     ) -> ManagedAddress {
-        let mut address = self
-            .pair_map()
-            .get(&PairTokens {
-                first_token_id: first_token_id.clone(),
-                second_token_id: second_token_id.clone(),
-            })
-            .unwrap_or_else(ManagedAddress::zero);
-
-        if address.is_zero() {
-            address = self
-                .pair_map()
-                .get(&PairTokens {
-                    first_token_id: second_token_id,
-                    second_token_id: first_token_id,
-                })
-                .unwrap_or_else(ManagedAddress::zero);
+        let mut opt_address = self.pair_map().get(&PairTokens {
+            first_token_id: first_token_id.clone(),
+            second_token_id: second_token_id.clone(),
+        });
+        if opt_address.is_some() {
+            return unsafe { opt_address.unwrap_unchecked() };
         }
-        address
+
+        opt_address = self.pair_map().get(&PairTokens {
+            first_token_id: second_token_id,
+            second_token_id: first_token_id,
+        });
+
+        opt_address.unwrap_or_else(ManagedAddress::zero)
     }
 
     fn get_pair_temporary_owner(&self, pair_address: &ManagedAddress) -> Option<ManagedAddress> {

dex/router/tests/router_setup/mod.rs

@@ -42,7 +42,7 @@ where
     RouterObjBuilder: 'static + Copy + Fn() -> router::ContractObj<DebugApi>,
     PairObjBuilder: 'static + Copy + Fn() -> pair::ContractObj<DebugApi>,
 {
-    pub blockchain_wrapper: BlockchainStateWrapper,
+    pub b_mock: BlockchainStateWrapper,
     pub owner_address: Address,
     pub user_address: Address,
     pub router_wrapper: ContractObjWrapper<router::ContractObj<DebugApi>, RouterObjBuilder>,
@@ -57,31 +57,23 @@ where
 {
     pub fn new(router_builder: RouterObjBuilder, pair_builder: PairObjBuilder) -> Self {
         let rust_zero = rust_biguint!(0u64);
-        let mut blockchain_wrapper = BlockchainStateWrapper::new();
-        let owner_addr = blockchain_wrapper.create_user_account(&rust_zero);
+        let mut b_mock = BlockchainStateWrapper::new();
+        let owner_addr = b_mock.create_user_account(&rust_zero);
 
-        let router_wrapper = blockchain_wrapper.create_sc_account(
+        let router_wrapper = b_mock.create_sc_account(
             &rust_zero,
             Some(&owner_addr),
             router_builder,
             ROUTER_WASM_PATH,
         );
 
-        let mex_pair_wrapper = blockchain_wrapper.create_sc_account(
-            &rust_zero,
-            Some(&owner_addr),
-            pair_builder,
-            PAIR_WASM_PATH,
-        );
+        let mex_pair_wrapper =
+            b_mock.create_sc_account(&rust_zero, Some(&owner_addr), pair_builder, PAIR_WASM_PATH);
 
-        let usdc_pair_wrapper = blockchain_wrapper.create_sc_account(
-            &rust_zero,
-            Some(&owner_addr),
-            pair_builder,
-            PAIR_WASM_PATH,
-        );
+        let usdc_pair_wrapper =
+            b_mock.create_sc_account(&rust_zero, Some(&owner_addr), pair_builder, PAIR_WASM_PATH);
 
-        blockchain_wrapper
+        b_mock
             .execute_tx(&owner_addr, &mex_pair_wrapper, &rust_zero, |sc| {
                 let first_token_id = managed_token_id!(WEGLD_TOKEN_ID);
                 let second_token_id = managed_token_id!(MEX_TOKEN_ID);
@@ -108,7 +100,7 @@ where
             })
             .assert_ok();
 
-        blockchain_wrapper
+        b_mock
             .execute_tx(&owner_addr, &usdc_pair_wrapper, &rust_zero, |sc| {
                 let first_token_id = managed_token_id!(WEGLD_TOKEN_ID);
                 let second_token_id = managed_token_id!(USDC_TOKEN_ID);
@@ -135,7 +127,7 @@ where
             })
             .assert_ok();
 
-        blockchain_wrapper
+        b_mock
             .execute_tx(&owner_addr, &router_wrapper, &rust_zero, |sc| {
                 sc.init(OptionalValue::None);
 
@@ -157,38 +149,38 @@ where
             .assert_ok();
 
         let lp_token_roles = [EsdtLocalRole::Mint, EsdtLocalRole::Burn];
-        blockchain_wrapper.set_esdt_local_roles(
+        b_mock.set_esdt_local_roles(
             mex_pair_wrapper.address_ref(),
             LPMEX_TOKEN_ID,
             &lp_token_roles[..],
         );
 
         let lp_token_roles = [EsdtLocalRole::Mint, EsdtLocalRole::Burn];
-        blockchain_wrapper.set_esdt_local_roles(
+        b_mock.set_esdt_local_roles(
             usdc_pair_wrapper.address_ref(),
             LPUSDC_TOKEN_ID,
             &lp_token_roles[..],
         );
 
-        let user_addr = blockchain_wrapper.create_user_account(&rust_biguint!(100_000_000));
-        blockchain_wrapper.set_esdt_balance(
+        let user_addr = b_mock.create_user_account(&rust_biguint!(100_000_000));
+        b_mock.set_esdt_balance(
             &user_addr,
             WEGLD_TOKEN_ID,
             &rust_biguint!(USER_TOTAL_WEGLD_TOKENS),
         );
-        blockchain_wrapper.set_esdt_balance(
+        b_mock.set_esdt_balance(
             &user_addr,
             MEX_TOKEN_ID,
             &rust_biguint!(USER_TOTAL_MEX_TOKENS),
         );
-        blockchain_wrapper.set_esdt_balance(
+        b_mock.set_esdt_balance(
             &user_addr,
             USDC_TOKEN_ID,
             &rust_biguint!(USER_TOTAL_USDC_TOKENS),
         );
 
         RouterSetup {
-            blockchain_wrapper,
+            b_mock,
             owner_address: owner_addr,
             user_address: user_addr,
             router_wrapper,
@@ -211,7 +203,7 @@ where
             },
         ];
 
-        self.blockchain_wrapper
+        self.b_mock
             .execute_esdt_multi_transfer(
                 &self.user_address,
                 &self.mex_pair_wrapper,
@@ -238,7 +230,7 @@ where
             },
         ];
 
-        self.blockchain_wrapper
+        self.b_mock
             .execute_esdt_multi_transfer(
                 &self.user_address,
                 &self.usdc_pair_wrapper,
@@ -261,7 +253,7 @@ where
     ) {
         let payment_amount_big = rust_biguint!(payment_amount);
 
-        self.blockchain_wrapper
+        self.b_mock
             .execute_esdt_transfer(
                 &self.user_address,
                 &self.router_wrapper,

dex/router/tests/router_test.rs

@@ -1,6 +1,7 @@
 #![allow(deprecated)]
 
 mod router_setup;
+use fees_collector::{fees_accumulation::FeesAccumulationModule, FeesCollector};
 use multiversx_sc::{
     codec::multi_types::OptionalValue,
     storage::mappers::StorageTokenWrapper,
@@ -131,17 +132,17 @@ fn test_multi_pair_swap() {
 
     router_setup.add_liquidity();
 
-    router_setup.blockchain_wrapper.check_esdt_balance(
+    router_setup.b_mock.check_esdt_balance(
         &router_setup.user_address,
         WEGLD_TOKEN_ID,
         &rust_biguint!(5_000_000_000),
     );
-    router_setup.blockchain_wrapper.check_esdt_balance(
+    router_setup.b_mock.check_esdt_balance(
         &router_setup.user_address,
         MEX_TOKEN_ID,
         &rust_biguint!(5_000_000_000),
     );
-    router_setup.blockchain_wrapper.check_esdt_balance(
+    router_setup.b_mock.check_esdt_balance(
         &router_setup.user_address,
         USDC_TOKEN_ID,
         &rust_biguint!(5_000_000_000),
@@ -164,17 +165,17 @@ fn test_multi_pair_swap() {
 
     router_setup.multi_pair_swap(MEX_TOKEN_ID, 100_000, &ops);
 
-    router_setup.blockchain_wrapper.check_esdt_balance(
+    router_setup.b_mock.check_esdt_balance(
         &router_setup.user_address,
         WEGLD_TOKEN_ID,
         &rust_biguint!(5_000_000_000), //unchanged
     );
-    router_setup.blockchain_wrapper.check_esdt_balance(
+    router_setup.b_mock.check_esdt_balance(
         &router_setup.user_address,
         MEX_TOKEN_ID,
         &rust_biguint!(4_999_900_000), //spent 100_000
     );
-    router_setup.blockchain_wrapper.check_esdt_balance(
+    router_setup.b_mock.check_esdt_balance(
         &router_setup.user_address,
         USDC_TOKEN_ID,
         &rust_biguint!(5_000_082_909), //gained 82_909
@@ -558,3 +559,70 @@ fn user_enable_pair_swaps_fail_test() {
         }),
     );
 }
+
+#[test]
+fn fees_collector_base_token_feature_test() {
+    let mut setup = RouterSetup::new(router::contract_obj, pair::contract_obj);
+
+    setup.add_liquidity();
+
+    let fc_wrapper = setup.b_mock.create_sc_account(
+        &rust_biguint!(0),
+        Some(&setup.owner_address),
+        fees_collector::contract_obj,
+        "fees collector path",
+    );
+
+    let router_address = setup.router_wrapper.address_ref().clone();
+    setup
+        .b_mock
+        .execute_tx(&setup.owner_address, &fc_wrapper, &rust_biguint!(0), |sc| {
+            sc.init(
+                managed_token_id!(b"LOCKED-123456"), // unused
+                managed_address!(&router_address),   // unused
+                managed_address!(&router_address),
+                managed_token_id!(WEGLD_TOKEN_ID),
+                MultiValueEncoded::new(),
+            );
+
+            let mut tokens = MultiValueEncoded::new();
+            tokens.push(managed_token_id!(WEGLD_TOKEN_ID));
+            tokens.push(managed_token_id!(USDC_TOKEN_ID));
+            tokens.push(managed_token_id!(CUSTOM_TOKEN_ID));
+
+            // must use qualified syntax, otherwise, you get complaints of multiple "config" modules
+            fees_collector::config::ConfigModule::add_known_tokens(&sc, tokens);
+
+            let _ = fees_collector::config::ConfigModule::known_contracts(&sc)
+                .insert(managed_address!(&setup.owner_address));
+        })
+        .assert_ok();
+
+    // try deposit USDC
+    setup
+        .b_mock
+        .set_esdt_balance(&setup.owner_address, USDC_TOKEN_ID, &rust_biguint!(1_000));
+
+    setup
+        .b_mock
+        .execute_esdt_transfer(
+            &setup.owner_address,
+            &fc_wrapper,
+            USDC_TOKEN_ID,
+            0,
+            &rust_biguint!(1_000),
+            |sc| {
+                sc.deposit_swap_fees();
+
+                // check fees were accumulate for WEGLD instead of USDC
+                assert!(sc
+                    .accumulated_fees(1, &managed_token_id!(USDC_TOKEN_ID))
+                    .is_empty());
+
+                assert!(!sc
+                    .accumulated_fees(1, &managed_token_id!(WEGLD_TOKEN_ID))
+                    .is_empty());
+            },
+        )
+        .assert_ok();
+}

energy-integration/fees-collector/src/additional_locked_tokens.rs

@@ -10,6 +10,9 @@ pub trait AdditionalLockedTokensModule:
     + crate::fees_accumulation::FeesAccumulationModule
     + crate::events::FeesCollectorEventsModule
     + week_timekeeping::WeekTimekeepingModule
+    + crate::external_sc_interactions::router::RouterInteractionsModule
+    + crate::external_sc_interactions::pair::PairInteractionsModule
+    + utils::UtilsModule
 {
     #[only_owner]
     #[endpoint(setLockedTokensPerBlock)]

energy-integration/fees-collector/src/claim.rs

@@ -24,6 +24,8 @@ pub trait ClaimModule:
     + sc_whitelist_module::SCWhitelistModule
     + multiversx_sc_modules::only_admin::OnlyAdminModule
     + crate::redistribute_rewards::RedistributeRewardsModule
+    + crate::external_sc_interactions::router::RouterInteractionsModule
+    + crate::external_sc_interactions::pair::PairInteractionsModule
 {
     #[endpoint(claimRewards)]
     fn claim_rewards_endpoint(

energy-integration/fees-collector/src/external_sc_interactions/mod.rs

@@ -0,0 +1,2 @@
+pub mod pair;
+pub mod router;

energy-integration/fees-collector/src/external_sc_interactions/pair.rs

@@ -0,0 +1,36 @@
+multiversx_sc::imports!();
+
+mod pair_proxy {
+    multiversx_sc::imports!();
+
+    #[multiversx_sc::proxy]
+    pub trait PairProxy {
+        #[payable("*")]
+        #[endpoint(swapTokensFixedInput)]
+        fn swap_tokens_fixed_input(
+            &self,
+            token_out: TokenIdentifier,
+            amount_out_min: BigUint,
+        ) -> EsdtTokenPayment;
+    }
+}
+
+const TOKEN_OUT_MIN: u32 = 1;
+
+#[multiversx_sc::module]
+pub trait PairInteractionsModule {
+    fn swap_to_common_token(
+        &self,
+        pair_address: ManagedAddress,
+        input_payment: EsdtTokenPayment,
+        token_out: TokenIdentifier,
+    ) -> EsdtTokenPayment {
+        self.pair_proxy_builder(pair_address)
+            .swap_tokens_fixed_input(token_out, BigUint::from(TOKEN_OUT_MIN))
+            .with_esdt_transfer(input_payment)
+            .execute_on_dest_context()
+    }
+
+    #[proxy]
+    fn pair_proxy_builder(&self, sc_address: ManagedAddress) -> pair_proxy::Proxy<Self::Api>;
+}