fix: concretize Self references in method signatures (#1001)

near · Jan 23, 2023 · 9aef9fd · 9aef9fd
1 parent 9ebe04e
commit 9aef9fd
Show file tree

Hide file tree

Showing 9 changed files with 132 additions and 17 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
@@ -156,7 +156,7 @@ impl ImplItemMethodInfo {
                         } else {
                             return syn::Error::new_spanned(
                                 &arg.ty,
-                                "Function parameters marked with  #[callback_vec] should have type Vec<T>",
+                                "Function parameters marked with #[callback_vec] should have type Vec<T>",
                             )
                             .into_compile_error();
                         };
@@ -257,7 +257,7 @@ fn generate_schema(ty: &Type, serializer_type: &SerializerType) -> TokenStream2
 }
 
 fn generate_abi_type(ty: &Type, serializer_type: &SerializerType) -> TokenStream2 {
-    let schema = generate_schema(ty, serializer_type);
+    let schema = generate_schema(&ty, serializer_type);
     match serializer_type {
         SerializerType::JSON => quote! {
             near_sdk::__private::AbiType::Json {

diff --git a/near-sdk-macros/src/core_impl/info_extractor/arg_info.rs b/near-sdk-macros/src/core_impl/info_extractor/arg_info.rs
@@ -1,5 +1,6 @@
-use crate::core_impl::info_extractor::serializer_attr::SerializerAttr;
-use crate::core_impl::info_extractor::SerializerType;
+use crate::core_impl::info_extractor::{SerializerAttr, SerializerType};
+use crate::core_impl::utils;
+use proc_macro2::TokenStream;
 use quote::ToTokens;
 use syn::{spanned::Spanned, Attribute, Error, Ident, Pat, PatType, Token, Type};
 
@@ -40,16 +41,15 @@ pub struct ArgInfo {
 
 impl ArgInfo {
     /// Extract near-sdk specific argument info.
-    pub fn new(original: &mut PatType) -> syn::Result<Self> {
+    pub fn new(original: &mut PatType, source_type: &TokenStream) -> syn::Result<Self> {
         let mut non_bindgen_attrs = vec![];
         let pat_reference;
         let pat_mutability;
-        let ident;
-        match original.pat.as_ref() {
+        let ident = match original.pat.as_ref() {
             Pat::Ident(pat_ident) => {
                 pat_reference = pat_ident.by_ref;
                 pat_mutability = pat_ident.mutability;
-                ident = pat_ident.ident.clone();
+                pat_ident.ident.clone()
             }
             _ => {
                 return Err(Error::new(
@@ -58,8 +58,9 @@ impl ArgInfo {
                 ));
             }
         };
+        *original.ty.as_mut() = utils::sanitize_self(&original.ty, source_type)?;
         let (reference, mutability, ty) = match original.ty.as_ref() {
-            x @ Type::Array(_) | x @ Type::Path(_) | x @ Type::Tuple(_) => {
+            x @ (Type::Array(_) | Type::Path(_) | Type::Tuple(_) | Type::Group(_)) => {
                 (None, None, (*x).clone())
             }
             Type::Reference(r) => (Some(r.and_token), r.mutability, (*r.elem.as_ref()).clone()),
@@ -69,7 +70,7 @@ impl ArgInfo {
         let mut bindgen_ty = BindgenArgType::Regular;
         // In the absence of serialization attributes this is a JSON serialization.
         let mut serializer_ty = SerializerType::JSON;
-        for attr in &mut original.attrs {
+        for attr in &original.attrs {
             let attr_str = attr.path.to_token_stream().to_string();
             match attr_str.as_str() {
                 "callback" | "callback_unwrap" => {

diff --git a/near-sdk-macros/src/core_impl/info_extractor/attr_sig_info.rs b/near-sdk-macros/src/core_impl/info_extractor/attr_sig_info.rs
@@ -1,5 +1,6 @@
 use super::{ArgInfo, BindgenArgType, InitAttr, MethodType, SerializerAttr, SerializerType};
-use proc_macro2::Span;
+use crate::core_impl::utils;
+use proc_macro2::{Span, TokenStream as TokenStream2};
 use quote::ToTokens;
 use syn::spanned::Spanned;
 use syn::{Attribute, Error, FnArg, Ident, Receiver, ReturnType, Signature};
@@ -37,6 +38,7 @@ impl AttrSigInfo {
     pub fn new(
         original_attrs: &mut Vec<Attribute>,
         original_sig: &mut Signature,
+        source_type: &TokenStream2,
     ) -> syn::Result<Self> {
         if original_sig.asyncness.is_some() {
             return Err(Error::new(
@@ -104,7 +106,7 @@ impl AttrSigInfo {
             match fn_arg {
                 FnArg::Receiver(r) => receiver = Some((*r).clone()),
                 FnArg::Typed(pat_typed) => {
-                    args.push(ArgInfo::new(pat_typed)?);
+                    args.push(ArgInfo::new(pat_typed, source_type)?);
                 }
             }
         }
@@ -132,7 +134,11 @@ impl AttrSigInfo {
         }
 
         *original_attrs = non_bindgen_attrs.clone();
-        let returns = original_sig.output.clone();
+        let mut returns = original_sig.output.clone();
+
+        if let ReturnType::Type(_, ref mut ty) = returns {
+            *ty.as_mut() = utils::sanitize_self(&*ty, source_type)?;
+        }
 
         let mut result = Self {
             ident,

diff --git a/near-sdk-macros/src/core_impl/info_extractor/impl_item_method_info.rs b/near-sdk-macros/src/core_impl/info_extractor/impl_item_method_info.rs
@@ -1,4 +1,5 @@
 use crate::core_impl::info_extractor::AttrSigInfo;
+use quote::ToTokens;
 use syn::{ImplItemMethod, Type, Visibility};
 
 /// Information extracted from `ImplItemMethod`.
@@ -15,7 +16,7 @@ impl ImplItemMethodInfo {
     /// Process the method and extract information important for near-sdk.
     pub fn new(original: &mut ImplItemMethod, struct_type: Type) -> syn::Result<Self> {
         let ImplItemMethod { attrs, sig, .. } = original;
-        let attr_signature_info = AttrSigInfo::new(attrs, sig)?;
+        let attr_signature_info = AttrSigInfo::new(attrs, sig, &struct_type.to_token_stream())?;
         let is_public = matches!(original.vis, Visibility::Public(_));
         Ok(Self { attr_signature_info, is_public, struct_type })
     }

diff --git a/near-sdk-macros/src/core_impl/info_extractor/item_trait_info.rs b/near-sdk-macros/src/core_impl/info_extractor/item_trait_info.rs
@@ -1,5 +1,6 @@
 use super::TraitItemMethodInfo;
 use inflector::Inflector;
+use quote::ToTokens;
 use syn::spanned::Spanned;
 use syn::{Error, Ident, ItemTrait, TraitItem};
 
@@ -30,7 +31,8 @@ impl ItemTraitInfo {
                     ))
                 }
                 TraitItem::Method(method) => {
-                    methods.push(TraitItemMethodInfo::new(method)?);
+                    methods
+                        .push(TraitItemMethodInfo::new(method, &original.ident.to_token_stream())?);
                     if method.default.is_some() {
                         return Err(Error::new(
                             method.span(),

diff --git a/near-sdk-macros/src/core_impl/info_extractor/trait_item_method_info.rs b/near-sdk-macros/src/core_impl/info_extractor/trait_item_method_info.rs
@@ -1,4 +1,5 @@
 use super::AttrSigInfo;
+use proc_macro2::TokenStream as TokenStream2;
 use syn::spanned::Spanned;
 use syn::{Error, LitStr, TraitItemMethod};
 
@@ -13,7 +14,7 @@ pub struct TraitItemMethodInfo {
 }
 
 impl TraitItemMethodInfo {
-    pub fn new(original: &mut TraitItemMethod) -> syn::Result<Self> {
+    pub fn new(original: &mut TraitItemMethod, trait_name: &TokenStream2) -> syn::Result<Self> {
         if original.default.is_some() {
             return Err(Error::new(
                 original.span(),
@@ -25,7 +26,7 @@ impl TraitItemMethodInfo {
 
         let TraitItemMethod { attrs, sig, .. } = original;
 
-        let attr_sig_info = AttrSigInfo::new(attrs, sig)?;
+        let attr_sig_info = AttrSigInfo::new(attrs, sig, trait_name)?;
 
         let ident_byte_str =
             LitStr::new(&attr_sig_info.ident.to_string(), attr_sig_info.ident.span());

diff --git a/near-sdk-macros/src/core_impl/utils/mod.rs b/near-sdk-macros/src/core_impl/utils/mod.rs
@@ -1,3 +1,5 @@
+use proc_macro2::{Group, TokenStream as TokenStream2, TokenTree};
+use quote::quote;
 use syn::{GenericArgument, Path, PathArguments, Type};
 
 /// Checks whether the given path is literally "Result".
@@ -73,3 +75,59 @@ pub(crate) fn extract_vec_type(ty: &Type) -> Option<&Type> {
         _ => None,
     }
 }
+
+fn _sanitize_self(typ: TokenStream2, replace_with: &TokenStream2) -> TokenStream2 {
+    let trees = typ.into_iter().map(|t| match t {
+        TokenTree::Ident(ident) if ident == "Self" => replace_with
+            .clone()
+            .into_iter()
+            .map(|mut t| {
+                t.set_span(ident.span());
+                t
+            })
+            .collect::<TokenStream2>(),
+        TokenTree::Group(group) => {
+            let stream = _sanitize_self(group.stream(), replace_with);
+            TokenTree::Group(Group::new(group.delimiter(), stream)).into()
+        }
+        rest => rest.into(),
+    });
+    trees.collect()
+}
+
+pub fn sanitize_self(typ: &Type, replace_with: &TokenStream2) -> syn::Result<Type> {
+    syn::parse2(_sanitize_self(quote! { #typ }, replace_with)).map_err(|original| {
+        syn::Error::new(original.span(), "Self sanitization failed. Please report this as a bug.")
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn sanitize_self_works() {
+        let typ: Type = syn::parse_str("Self").unwrap();
+        let replace_with: TokenStream2 = syn::parse_str("MyType").unwrap();
+        let sanitized = sanitize_self(&typ, &replace_with).unwrap();
+        assert_eq!(quote! { #sanitized }.to_string(), "MyType");
+
+        let typ: Type = syn::parse_str("Vec<Self>").unwrap();
+        let replace_with: TokenStream2 = syn::parse_str("MyType").unwrap();
+        let sanitized = sanitize_self(&typ, &replace_with).unwrap();
+        assert_eq!(quote! { #sanitized }.to_string(), "Vec < MyType >");
+
+        let typ: Type = syn::parse_str("Vec<Vec<Self>>").unwrap();
+        let replace_with: TokenStream2 = syn::parse_str("MyType").unwrap();
+        let sanitized = sanitize_self(&typ, &replace_with).unwrap();
+        assert_eq!(quote! { #sanitized }.to_string(), "Vec < Vec < MyType > >");
+
+        let typ: Type = syn::parse_str("Option<[(Self, Result<Self, ()>); 2]>").unwrap();
+        let replace_with: TokenStream2 = syn::parse_str("MyType").unwrap();
+        let sanitized = sanitize_self(&typ, &replace_with).unwrap();
+        assert_eq!(
+            quote! { #sanitized }.to_string(),
+            "Option < [(MyType , Result < MyType , () >) ; 2] >"
+        );
+    }
+}
diff --git a/near-sdk/compilation_tests/all.rs b/near-sdk/compilation_tests/all.rs
@@ -22,4 +22,5 @@ fn compilation_tests() {
     t.pass("compilation_tests/enum_near_bindgen.rs");
     t.pass("compilation_tests/schema_derive.rs");
     t.compile_fail("compilation_tests/schema_derive_invalids.rs");
+    t.pass("compilation_tests/self_support.rs");
 }
diff --git a/near-sdk/compilation_tests/self_support.rs b/near-sdk/compilation_tests/self_support.rs
@@ -0,0 +1,45 @@
+//! Method signature uses Self.
+
+use near_sdk::near_bindgen;
+use serde::{Deserialize, Serialize};
+
+#[near_bindgen]
+#[derive(Default, Serialize, Deserialize)]
+pub struct Ident {
+    value: u32,
+}
+
+#[near_bindgen]
+impl Ident {
+    pub fn plain_arg(_a: Self) {
+        unimplemented!()
+    }
+    pub fn plain_ret() -> Self {
+        unimplemented!()
+    }
+    pub fn plain_arg_ret(a: Self) -> Self {
+        a
+    }
+    pub fn nested_arg(_a: Vec<Self>) {
+        unimplemented!()
+    }
+    pub fn nested_ret() -> Vec<Self> {
+        unimplemented!()
+    }
+    pub fn nested_arg_ret(a: Vec<Self>) -> Vec<Self> {
+        a
+    }
+    pub fn deeply_nested_arg(_a: Option<[(Self, Result<Self, ()>); 2]>) {
+        unimplemented!()
+    }
+    pub fn deeply_nested_ret() -> Option<[(Self, Result<Self, ()>); 2]> {
+        unimplemented!()
+    }
+    pub fn deeply_nested_arg_ret(
+        a: Option<[(Self, Result<Self, ()>); 2]>,
+    ) -> Option<[(Self, Result<Self, ()>); 2]> {
+        a
+    }
+}
+
+fn main() {}