router pair map migration

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<bool>;

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<PairTokens<Self::Api>, ManagedAddress>;
 
+    #[storage_mapper("address_pair_map")]
+    fn address_pair_map(&self) -> MapMapper<ManagedAddress, PairTokens<Self::Api>>;
+
     #[view(getPairTemplateAddress)]
     #[storage_mapper("pair_template_address")]
     fn pair_template_address(&self) -> SingleValueMapper<ManagedAddress>;

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();
+    }
 }

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),
             ])));