feat: allow imported FieldElement type in macros (#285)

xJonathanLEI · Jan 21, 2023 · e3f3e4f · e3f3e4f
1 parent 54d2890
commit e3f3e4f
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 10 deletions.
diff --git a/starknet-macros/Cargo.toml b/starknet-macros/Cargo.toml
@@ -18,3 +18,7 @@ proc-macro = true
 [dependencies]
 starknet-core = { version = "0.2.0", path = "../starknet-core" }
 syn = "1.0.96"
+
+[features]
+default = []
+use_imported_type = []
diff --git a/starknet-macros/src/lib.rs b/starknet-macros/src/lib.rs
@@ -15,8 +15,12 @@ pub fn selector(input: TokenStream) -> TokenStream {
     let selector_raw = selector_value.into_mont();
 
     format!(
-        "::starknet::core::types::FieldElement::from_mont([{}, {}, {}, {}])",
-        selector_raw[0], selector_raw[1], selector_raw[2], selector_raw[3],
+        "{}::from_mont([{}, {}, {}, {}])",
+        field_element_path(),
+        selector_raw[0],
+        selector_raw[1],
+        selector_raw[2],
+        selector_raw[3],
     )
     .parse()
     .unwrap()
@@ -32,8 +36,12 @@ pub fn short_string(input: TokenStream) -> TokenStream {
     let felt_raw = felt_value.into_mont();
 
     format!(
-        "::starknet::core::types::FieldElement::from_mont([{}, {}, {}, {}])",
-        felt_raw[0], felt_raw[1], felt_raw[2], felt_raw[3],
+        "{}::from_mont([{}, {}, {}, {}])",
+        field_element_path(),
+        felt_raw[0],
+        felt_raw[1],
+        felt_raw[2],
+        felt_raw[3],
     )
     .parse()
     .unwrap()
@@ -54,8 +62,12 @@ pub fn felt(input: TokenStream) -> TokenStream {
     let felt_raw = felt_value.into_mont();
 
     format!(
-        "::starknet::core::types::FieldElement::from_mont([{}, {}, {}, {}])",
-        felt_raw[0], felt_raw[1], felt_raw[2], felt_raw[3],
+        "{}::from_mont([{}, {}, {}, {}])",
+        field_element_path(),
+        felt_raw[0],
+        felt_raw[1],
+        felt_raw[2],
+        felt_raw[3],
     )
     .parse()
     .unwrap()
@@ -71,8 +83,12 @@ pub fn felt_dec(input: TokenStream) -> TokenStream {
     let felt_raw = felt_value.into_mont();
 
     format!(
-        "::starknet::core::types::FieldElement::from_mont([{}, {}, {}, {}])",
-        felt_raw[0], felt_raw[1], felt_raw[2], felt_raw[3],
+        "{}::from_mont([{}, {}, {}, {}])",
+        field_element_path(),
+        felt_raw[0],
+        felt_raw[1],
+        felt_raw[2],
+        felt_raw[3],
     )
     .parse()
     .unwrap()
@@ -88,9 +104,23 @@ pub fn felt_hex(input: TokenStream) -> TokenStream {
     let felt_raw = felt_value.into_mont();
 
     format!(
-        "::starknet::core::types::FieldElement::from_mont([{}, {}, {}, {}])",
-        felt_raw[0], felt_raw[1], felt_raw[2], felt_raw[3],
+        "{}::from_mont([{}, {}, {}, {}])",
+        field_element_path(),
+        felt_raw[0],
+        felt_raw[1],
+        felt_raw[2],
+        felt_raw[3],
     )
     .parse()
     .unwrap()
 }
+
+#[cfg(feature = "use_imported_type")]
+fn field_element_path() -> &'static str {
+    "FieldElement"
+}
+
+#[cfg(not(feature = "use_imported_type"))]
+fn field_element_path() -> &'static str {
+    "::starknet::core::types::FieldElement"
+}