draft: working aliasing system

near · Feb 14, 2024 · 760a900 · 760a900
1 parent 2705096
commit 760a900
Show file tree

Hide file tree

Showing 10 changed files with 156 additions and 86 deletions.
diff --git a/near-sdk-macros/src/core_impl/abi/abi_generator.rs b/near-sdk-macros/src/core_impl/abi/abi_generator.rs
@@ -88,7 +88,13 @@ impl ImplItemMethodInfo {
     pub fn abi_struct(&self) -> TokenStream2 {
         let attr_signature_info = &self.attr_signature_info;
 
-        let function_name_str = attr_signature_info.ident.to_string();
+        let ident_ = attr_signature_info.ident.to_string();
+        let function_name_str = match &attr_signature_info.method_kind {
+            MethodKind::View(m) => m.alias.clone().unwrap_or(ident_.to_string()),
+            MethodKind::Call(m) => m.alias.clone().unwrap_or(ident_.to_string()),
+            MethodKind::Init(m) => m.alias.clone().unwrap_or(ident_.to_string()),
+        };
+
         let function_doc = match parse_rustdoc(&attr_signature_info.non_bindgen_attrs) {
             Some(doc) => quote! { ::std::option::Option::Some(::std::string::String::from(#doc)) },
             None => quote! { ::std::option::Option::None },

diff --git a/near-sdk-macros/src/core_impl/code_generator/ext.rs b/near-sdk-macros/src/core_impl/code_generator/ext.rs
@@ -1,4 +1,7 @@
-use crate::core_impl::{serializer, AttrSigInfo};
+use crate::{
+    core_impl::{serializer, AttrSigInfo},
+    MethodKind,
+};
 use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
 use quote::{format_ident, quote, ToTokens};
 use syn::{parse_quote, Attribute, Generics, Path, Signature};
@@ -116,8 +119,21 @@ fn generate_ext_function(attr_signature_info: &AttrSigInfo) -> TokenStream2 {
     let serialize =
         serializer::generate_serializer(attr_signature_info, &attr_signature_info.input_serializer);
 
-    let AttrSigInfo { non_bindgen_attrs, ident, original_sig, .. } = attr_signature_info;
-    let ident_str = ident.to_string();
+    let AttrSigInfo { non_bindgen_attrs, ident, original_sig, method_kind, .. } =
+        attr_signature_info;
+
+    let ident_ = ident.clone().to_string();
+    let ident_str = match method_kind {
+        MethodKind::Call(m) => &m.alias,
+        MethodKind::Init(m) => &m.alias,
+        MethodKind::View(m) => &m.alias,
+    }
+    .as_ref()
+    .unwrap_or(&ident_)
+    .to_string();
+
+    let ident = &Ident::new(&ident_str, ident.span());
+
     let mut new_non_bindgen_attrs = TokenStream2::new();
     for attribute in non_bindgen_attrs.iter() {
         if is_fn_attribute_to_forward(attribute) {
@@ -154,15 +170,15 @@ mod tests {
     fn ext_gen() {
         let st: ItemStruct = parse_quote! { struct Test { a: u8 } };
         let actual = generate_ext_structs(&st.ident, Some(&st.generics));
-       
+
         local_insta_assert_snapshot!(pretty_print_syn_str(&actual).unwrap());
     }
 
     #[test]
     fn module_ext_gen() {
         let ident: Ident = parse_quote! { Test };
         let actual = generate_ext_structs(&ident, None);
-    
+
         local_insta_assert_snapshot!(pretty_print_syn_str(&actual).unwrap());
     }
 
@@ -177,7 +193,7 @@ mod tests {
             #[warn(unused)]
             pub fn method(&self) { }
         };
-        let method_info = ImplItemMethodInfo::new(&mut method, false, impl_type).unwrap().unwrap();
+        let method_info = ImplItemMethodInfo::new(&mut method, None, impl_type).unwrap().unwrap();
         let actual = generate_ext_function(&method_info.attr_signature_info);
 
         local_insta_assert_snapshot!(pretty_print_syn_str(&actual).unwrap());
@@ -189,9 +205,9 @@ mod tests {
         let mut method: ImplItemFn = parse_quote! {
             pub fn method(&self, k: &String) { }
         };
-        let method_info = ImplItemMethodInfo::new(&mut method, false, impl_type).unwrap().unwrap();
+        let method_info = ImplItemMethodInfo::new(&mut method, None, impl_type).unwrap().unwrap();
         let actual = generate_ext_function(&method_info.attr_signature_info);
-     
+
         local_insta_assert_snapshot!(pretty_print_syn_str(&actual).unwrap());
     }
 
@@ -201,9 +217,9 @@ mod tests {
         let mut method: syn::ImplItemFn = parse_quote! {
           pub fn borsh_test(&mut self, #[serializer(borsh)] a: String) {}
         };
-        let method_info = ImplItemMethodInfo::new(&mut method, false, impl_type).unwrap().unwrap();
+        let method_info = ImplItemMethodInfo::new(&mut method, None, impl_type).unwrap().unwrap();
         let actual = generate_ext_function(&method_info.attr_signature_info);
-       
+
         local_insta_assert_snapshot!(pretty_print_syn_str(&actual).unwrap());
     }
 }
diff --git a/near-sdk-macros/src/core_impl/code_generator/impl_item_method_info.rs b/near-sdk-macros/src/core_impl/code_generator/impl_item_method_info.rs
@@ -9,18 +9,10 @@ impl ImplItemMethodInfo {
     pub fn method_wrapper(&self) -> TokenStream2 {
         let non_bindgen_attrs = self.non_bindgen_attrs_tokens();
 
-        // If the method has an alias, use it, otherwise use the method name.
-        let abi_alias = match &self.attr_signature_info.method_kind {
-            MethodKind::Call(call_method) => &call_method.alias,
-            MethodKind::Init(init_method) => &init_method.alias,
-            MethodKind::View(view_method) => &view_method.alias,
-        }
-        .as_ref()
-        .map(|alias| syn::Ident::new(&alias, self.attr_signature_info.ident.span()))
-        .unwrap_or(self.attr_signature_info.ident.clone());
-
         let panic_hook = self.panic_hook_tokens();
 
+        let ident = self.abi_alias();
+
         let arg_struct = self.arg_struct_tokens();
         let arg_parsing = self.arg_parsing_tokens();
 
@@ -43,7 +35,7 @@ impl ImplItemMethodInfo {
             #non_bindgen_attrs
             #[cfg(target_arch = "wasm32")]
             #[no_mangle]
-            pub extern "C" fn #abi_alias() {
+            pub extern "C" fn #ident() {
                 #panic_hook
                 #is_private_check
                 #deposit_check
@@ -282,11 +274,18 @@ impl ImplItemMethodInfo {
 
     fn method_invocation_tokens(&self) -> TokenStream2 {
         use MethodKind::*;
-
         let ident = &self.attr_signature_info.ident;
         let arg_list = self.attr_signature_info.arg_list();
 
         let method_invocation = || {
+            if let Some((_, path, _)) = &self.attr_signature_info.trait_ {
+                if let Some(trait_ident) = path.get_ident() {
+                    return quote! {
+                        #trait_ident::#ident(&contract, #arg_list)
+                    };
+                };
+            }
+
             quote! {
                 contract.#ident(#arg_list)
             }
@@ -410,4 +409,15 @@ impl ImplItemMethodInfo {
             }
         })
     }
+
+    fn abi_alias(&self) -> syn::Ident {
+        match &self.attr_signature_info.method_kind {
+            MethodKind::Call(call_method) => &call_method.alias,
+            MethodKind::Init(init_method) => &init_method.alias,
+            MethodKind::View(view_method) => &view_method.alias,
+        }
+        .as_ref()
+        .map(|alias| syn::Ident::new(&alias, self.attr_signature_info.ident.span()))
+        .unwrap_or(self.attr_signature_info.ident.clone())
+    }
 }