sc-config proxy - add variant

multiversx · May 14, 2024 · bdefe36 · bdefe36
1 parent de13520
commit bdefe36
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 41 deletions.
diff --git a/framework/meta/src/cmd/contract/generate_proxy/proxy_generator.rs b/framework/meta/src/cmd/contract/generate_proxy/proxy_generator.rs
@@ -1,8 +1,8 @@
 use std::{fmt::Display, fs::File, io::Write};
 
 use multiversx_sc::abi::{
-    EndpointAbi, EnumVariantDescription, InputAbi, OutputAbi, StructFieldDescription, TypeContents,
-    TypeDescription,
+    ContractAbi, EndpointAbi, EnumVariantDescription, InputAbi, OutputAbi, StructFieldDescription,
+    TypeContents, TypeDescription,
 };
 
 use crate::cmd::contract::{meta_config::MetaConfig, sc_config::ProxyConfigSerde};
@@ -43,6 +43,7 @@ pub struct ProxyGenerator<'a> {
     pub meta_config: &'a MetaConfig,
     pub file: Option<&'a mut File>,
     pub proxy_config: &'a ProxyConfigSerde,
+    pub contract_abi: &'a ContractAbi,
 }
 
 impl<'a> ProxyGenerator<'a> {
@@ -55,6 +56,7 @@ impl<'a> ProxyGenerator<'a> {
             meta_config,
             file: Some(file),
             proxy_config,
+            contract_abi: &meta_config.original_contract_abi,
         }
     }
 
@@ -69,6 +71,7 @@ impl<'a> ProxyGenerator<'a> {
     }
 
     pub fn write_proxy_to_file(&mut self) {
+        self.extract_contract_abi();
         self.write_header();
         self.write_tx_proxy_type_def();
         self.write_impl_for_tx_proxy();
@@ -77,6 +80,16 @@ impl<'a> ProxyGenerator<'a> {
         self.write_types();
     }
 
+    fn extract_contract_abi(&mut self) {
+        if self.proxy_config.variant.is_some() {
+            for contract_variant in &self.meta_config.sc_config.contracts {
+                if self.proxy_config.variant.clone().unwrap() == contract_variant.abi.name {
+                    self.contract_abi = &contract_variant.abi;
+                }
+            }
+        }
+    }
+
     fn write_header(&mut self) {
         self.writeln(PRELUDE);
         match &self.proxy_config.override_import {
@@ -86,17 +99,16 @@ impl<'a> ProxyGenerator<'a> {
     }
 
     fn write_tx_proxy_type_def(&mut self) {
-        let proxy_type_name = proxy_type_name(&self.meta_config.original_contract_abi.name);
+        let proxy_type_name = proxy_type_name(&self.contract_abi.name);
         self.writeln(format!(
             r#"
 pub struct {proxy_type_name};"#
         ));
     }
 
     fn write_impl_for_tx_proxy(&mut self) {
-        let proxy_type_name = proxy_type_name(&self.meta_config.original_contract_abi.name);
-        let proxy_methods_type_name =
-            proxy_methods_type_name(&self.meta_config.original_contract_abi.name);
+        let proxy_type_name = proxy_type_name(&self.contract_abi.name);
+        let proxy_methods_type_name = proxy_methods_type_name(&self.contract_abi.name);
         self.writeln(format!(
             r#"
 impl<Env, From, To, Gas> TxProxyTrait<Env, From, To, Gas> for {proxy_type_name}
@@ -116,8 +128,7 @@ where
     }
 
     fn write_struct_tx_proxy_methods(&mut self) {
-        let proxy_methods_type_name =
-            proxy_methods_type_name(&self.meta_config.original_contract_abi.name);
+        let proxy_methods_type_name = proxy_methods_type_name(&self.contract_abi.name);
         self.writeln(format!(
             r#"
 pub struct {proxy_methods_type_name}<Env, From, To, Gas>
@@ -133,35 +144,22 @@ where
     }
 
     fn write_content(&mut self) {
-        if !self
-            .meta_config
-            .original_contract_abi
-            .constructors
-            .is_empty()
-        {
+        if !self.contract_abi.constructors.is_empty() {
             self.write_constructors();
         }
 
-        if !self
-            .meta_config
-            .original_contract_abi
-            .upgrade_constructors
-            .is_empty()
-        {
+        if !self.contract_abi.upgrade_constructors.is_empty() {
             self.write_upgrades();
         }
 
-        if !self.meta_config.original_contract_abi.endpoints.is_empty() {
+        if !self.contract_abi.endpoints.is_empty() {
             self.write_endpoints();
         }
     }
 
     fn write_types(&mut self) {
-        for (_, type_description) in &self.meta_config.original_contract_abi.type_descriptions.0 {
-            if self
-                .meta_config
-                .original_contract_abi
-                .get_crate_name_for_code()
+        for (_, type_description) in &self.contract_abi.type_descriptions.0 {
+            if self.contract_abi.get_crate_name_for_code()
                 != extract_struct_crate(type_description.names.rust.as_str())
             {
                 continue;
@@ -186,8 +184,7 @@ where
     }
 
     fn write_constructors(&mut self) {
-        let constructors: Vec<EndpointAbi> =
-            self.meta_config.original_contract_abi.constructors.clone();
+        let constructors: Vec<EndpointAbi> = self.contract_abi.constructors.clone();
 
         self.write_header_impl_constructor();
         for (i, constructor_abi) in constructors.into_iter().enumerate() {
@@ -205,8 +202,7 @@ where
     fn write_upgrades(&mut self) {
         self.write_header_impl_upgrade();
         for (i, upgrade) in self
-            .meta_config
-            .original_contract_abi
+            .contract_abi
             .upgrade_constructors
             .clone()
             .into_iter()
@@ -224,7 +220,7 @@ where
     }
 
     fn write_endpoints(&mut self) {
-        let endpoints: Vec<EndpointAbi> = self.meta_config.original_contract_abi.endpoints.clone();
+        let endpoints: Vec<EndpointAbi> = self.contract_abi.endpoints.clone();
 
         self.write_header_impl_endpoints();
         for (i, endpoint_abi) in endpoints.into_iter().enumerate() {
@@ -240,8 +236,7 @@ where
     }
 
     fn write_header_impl_constructor(&mut self) {
-        let proxy_methods_type_name =
-            proxy_methods_type_name(&self.meta_config.original_contract_abi.name);
+        let proxy_methods_type_name = proxy_methods_type_name(&self.contract_abi.name);
         self.writeln(format!(
             r#"
 #[rustfmt::skip]
@@ -256,8 +251,7 @@ where
     }
 
     fn write_header_impl_upgrade(&mut self) {
-        let proxy_methods_type_name =
-            proxy_methods_type_name(&self.meta_config.original_contract_abi.name);
+        let proxy_methods_type_name = proxy_methods_type_name(&self.contract_abi.name);
         self.writeln(format!(
             r#"
 #[rustfmt::skip]
@@ -273,8 +267,7 @@ where
     }
 
     fn write_header_impl_endpoints(&mut self) {
-        let proxy_methods_type_name =
-            proxy_methods_type_name(&self.meta_config.original_contract_abi.name);
+        let proxy_methods_type_name = proxy_methods_type_name(&self.contract_abi.name);
         self.writeln(format!(
             r#"
 #[rustfmt::skip]
@@ -606,10 +599,7 @@ where
 
     fn process_paths(&self, paths: &Vec<String>) -> Vec<String> {
         let mut processed_paths: Vec<String> = Vec::new();
-        let crate_name = self
-            .meta_config
-            .original_contract_abi
-            .get_crate_name_for_code();
+        let crate_name = self.contract_abi.get_crate_name_for_code();
 
         for path in paths {
             let type_rust_name = path.split("::").last().unwrap();
@@ -665,6 +655,7 @@ pub mod tests {
             meta_config: &meta_config,
             file: None,
             proxy_config: &ProxyConfigSerde::new(),
+            contract_abi: &meta_config.original_contract_abi,
         };
 
         let cleaned_path_unsanitized = proxy_generator.clean_paths(
@@ -696,6 +687,7 @@ pub mod tests {
             meta_config: &meta_config,
             file: None,
             proxy_config: &ProxyConfigSerde::new(),
+            contract_abi: &meta_config.original_contract_abi,
         };
 
         let cleaned_path_sanitized = proxy_generator.clean_paths(

diff --git a/framework/meta/src/cmd/contract/sc_config/sc_config_proxy.rs b/framework/meta/src/cmd/contract/sc_config/sc_config_proxy.rs
@@ -15,6 +15,9 @@ pub struct ProxyConfigSerde {
     #[serde(default)]
     #[serde(rename = "path-rename")]
     pub path_rename: Option<Vec<PathRename>>,
+
+    #[serde(default)]
+    pub variant: Option<String>,
 }
 
 impl ProxyConfigSerde {
@@ -23,6 +26,7 @@ impl ProxyConfigSerde {
             path: DEFAULT_PATH.to_string(),
             override_import: None,
             path_rename: None,
+            variant: None,
         }
     }
 }