diff --git a/dex/router/src/contract.rs b/dex/router/src/contract.rs index 8e9f685e5..56e9cba45 100644 --- a/dex/router/src/contract.rs +++ b/dex/router/src/contract.rs @@ -322,6 +322,8 @@ pub trait Router: .unwrap_or_else(ManagedAddress::zero); } + self.address_pair_map().remove(&pair_address); + pair_address } @@ -394,6 +396,20 @@ pub trait Router: self.pair_creation_enabled().set(enabled); } + #[only_owner] + #[endpoint(migratePairMap)] + fn migrate_pair_map(&self) { + let pair_map = self.pair_map(); + let mut address_pair_map = self.address_pair_map(); + require!( + address_pair_map.is_empty(), + "The destination mapper must be empty" + ); + for (pair_tokens, address) in pair_map.iter() { + address_pair_map.insert(address, pair_tokens); + } + } + #[view(getPairCreationEnabled)] #[storage_mapper("pair_creation_enabled")] fn pair_creation_enabled(&self) -> SingleValueMapper; diff --git a/dex/router/src/factory.rs b/dex/router/src/factory.rs index bd4b803ca..c0badf49c 100644 --- a/dex/router/src/factory.rs +++ b/dex/router/src/factory.rs @@ -69,6 +69,13 @@ pub trait FactoryModule { }, new_address.clone(), ); + self.address_pair_map().insert( + new_address.clone(), + PairTokens { + first_token_id: first_token_id.clone(), + second_token_id: second_token_id.clone(), + }, + ); self.pair_temporary_owner().insert( new_address.clone(), ( @@ -167,9 +174,7 @@ pub trait FactoryModule { fn check_is_pair_sc(&self, pair_address: &ManagedAddress) { require!( - self.pair_map() - .values() - .any(|address| &address == pair_address), + self.address_pair_map().contains_key(pair_address), "Not a pair SC" ); } @@ -215,6 +220,9 @@ pub trait FactoryModule { #[storage_mapper("pair_map")] fn pair_map(&self) -> MapMapper, ManagedAddress>; + #[storage_mapper("address_pair_map")] + fn address_pair_map(&self) -> MapMapper>; + #[view(getPairTemplateAddress)] #[storage_mapper("pair_template_address")] fn pair_template_address(&self) -> SingleValueMapper; diff --git a/dex/router/tests/router_setup/mod.rs b/dex/router/tests/router_setup/mod.rs index 98313d35a..54f6330a3 100644 --- a/dex/router/tests/router_setup/mod.rs +++ b/dex/router/tests/router_setup/mod.rs @@ -282,4 +282,17 @@ where ) .assert_ok(); } + + pub fn migrate_pair_map(&mut self) { + self.blockchain_wrapper + .execute_tx( + &self.owner_address, + &self.router_wrapper, + &rust_biguint!(0u64), + |sc| { + sc.migrate_pair_map(); + }, + ) + .assert_ok(); + } } diff --git a/dex/router/tests/router_test.rs b/dex/router/tests/router_test.rs index 3d395316d..17e74be21 100644 --- a/dex/router/tests/router_test.rs +++ b/dex/router/tests/router_test.rs @@ -133,6 +133,7 @@ fn test_router_upgrade_pair() { #[test] fn test_multi_pair_swap() { let mut router_setup = RouterSetup::new(router::contract_obj, pair::contract_obj); + router_setup.migrate_pair_map(); router_setup.add_liquidity(); @@ -253,6 +254,8 @@ fn user_enable_pair_swaps_through_router_test() { managed_address!(pair_wrapper.address_ref()), ); + sc.migrate_pair_map(); + sc.add_common_tokens_for_user_pairs(MultiValueEncoded::from(ManagedVec::from(vec![ managed_token_id!(USDC_TOKEN_ID), ]))); @@ -438,6 +441,8 @@ fn user_enable_pair_swaps_fail_test() { managed_address!(pair_wrapper.address_ref()), ); + sc.migrate_pair_map(); + sc.add_common_tokens_for_user_pairs(MultiValueEncoded::from(ManagedVec::from(vec![ managed_token_id!(USDC_TOKEN_ID), ])));