Allow &self/self/&mut self on contract fns

soroban-sdk-macros/src/derive_fn.rs

@@ -6,9 +6,18 @@ use syn::{
     punctuated::Punctuated,
     spanned::Spanned,
     token::{Colon, Comma},
-    Attribute, Error, FnArg, Ident, Pat, PatIdent, PatType, Path, Type, TypePath, TypeReference,
+    Attribute, Error, FnArg, Ident, Pat, PatIdent, PatType, Path, Receiver, Type, TypePath,
+    TypeReference,
 };
 
+enum SpecialFirstInput {
+    BorrowEnv,
+    OwnedEnv,
+    BorrowMutSelf,
+    BorrowSelf,
+    OwnedSelf,
+}
+
 #[allow(clippy::too_many_arguments)]
 pub fn derive_fn(
     crate_path: &Path,
@@ -23,7 +32,7 @@ pub fn derive_fn(
     let mut errors = Vec::<Error>::new();
 
     // Prepare the env input.
-    let env_input = inputs.first().and_then(|a| match a {
+    let special_first_input = inputs.first().and_then(|a| match a {
         FnArg::Typed(pat_type) => {
             let mut is_ref = false;
             let mut ty = &*pat_type.ty;
@@ -37,21 +46,44 @@ pub fn derive_fn(
             }) = ty
             {
                 if segments.last().map_or(false, |s| s.ident == "Env") {
-                    Some(is_ref)
+                    if is_ref {
+                        Some(SpecialFirstInput::BorrowEnv)
+                    } else {
+                        Some(SpecialFirstInput::OwnedEnv)
+                    }
                 } else {
                     None
                 }
             } else {
                 None
             }
         }
-        FnArg::Receiver(_) => None,
+        FnArg::Receiver(Receiver {
+            reference: Some(_),
+            mutability: Some(_),
+            ..
+        }) => Some(SpecialFirstInput::BorrowMutSelf),
+        FnArg::Receiver(Receiver {
+            reference: Some(_),
+            mutability: None,
+            ..
+        }) => Some(SpecialFirstInput::BorrowSelf),
+        FnArg::Receiver(Receiver {
+            reference: None,
+            mutability: None,
+            ..
+        }) => Some(SpecialFirstInput::OwnedSelf),
+        FnArg::Receiver(Receiver {
+            reference: None,
+            mutability: Some(_),
+            ..
+        }) => None,
     });
 
     // Prepare the argument inputs.
     let (wrap_args, wrap_calls): (Vec<_>, Vec<_>) = inputs
         .iter()
-        .skip(if env_input.is_some() { 1 } else { 0 })
+        .skip(if special_first_input.is_some() { 1 } else { 0 })
         .enumerate()
         .map(|(i, a)| match a {
             FnArg::Typed(_) => {
@@ -92,14 +124,13 @@ pub fn derive_fn(
         "use `{}::new(&env, &contract_id).{}` instead",
         client_ident, &ident
     );
-    let env_call = if let Some(is_ref) = env_input {
-        if is_ref {
-            quote! { &env, }
-        } else {
-            quote! { env.clone(), }
-        }
-    } else {
-        quote! {}
+    let env_call = match special_first_input {
+        None => quote! {},
+        Some(SpecialFirstInput::BorrowEnv) => quote! { &env, },
+        Some(SpecialFirstInput::OwnedEnv) => quote! { env.clone(), },
+        Some(SpecialFirstInput::BorrowSelf) => quote! { &env, },
+        Some(SpecialFirstInput::BorrowMutSelf) => quote! { &env, },
+        Some(SpecialFirstInput::OwnedSelf) => quote! { env.clone(), },
     };
     let slice_args: Vec<TokenStream2> = (0..wrap_args.len()).map(|n| quote! { args[#n] }).collect();
     let use_trait = if let Some(t) = trait_ident {

soroban-sdk-macros/src/lib.rs

@@ -125,6 +125,7 @@ pub fn contractspecfn(metadata: TokenStream, input: TokenStream) -> TokenStream
 struct ContractArgs {
     #[darling(default = "default_crate_path")]
     crate_path: Path,
+    v2: bool,
 }
 
 #[proc_macro_attribute]
@@ -140,19 +141,57 @@ pub fn contract(metadata: TokenStream, input: TokenStream) -> TokenStream {
         Err(e) => return e.write_errors().into(),
     };
 
-    let input2: TokenStream2 = input.clone().into();
+    let input_orig: TokenStream2 = input.clone().into();
 
     let item = parse_macro_input!(input as ItemStruct);
+    match &item.fields {
+        Fields::Named(_) | Fields::Unnamed(_) => {
+            return Error::new(
+                item.span(),
+                "contract type must be a unit struct with no fields",
+            )
+            .into_compile_error()
+            .into()
+        }
+        Fields::Unit => (),
+    }
+
+    let crate_path = &args.crate_path;
 
     let ty = &item.ident;
     let ty_str = quote!(#ty).to_string();
 
     let client_ident = format!("{ty_str}Client");
     let fn_set_registry_ident = format_ident!("__{ty_str}_fn_set_registry");
-    let crate_path = &args.crate_path;
     let client = derive_client_type(&args.crate_path, &ty_str, &client_ident);
+
+    let code = if args.v2 {
+        quote! {
+            pub struct #ty {
+                env: #crate_path::Env;
+            }
+
+            #[cfg(any(test, feature = "testutils"))]
+            impl #crate_path::testutils::ContractStruct for #ty {
+                fn new(env: #crate_path::Env) -> Self {
+                    Self { env }
+                }
+            }
+
+            impl #ty {
+                pub fn env(&self) -> &#crate_path::Env {
+                    &self.env
+                }
+            }
+        }
+    } else {
+        quote! {
+            #input_orig
+        }
+    };
     quote! {
-        #input2
+        #code
+
         #client
 
         #[cfg(any(test, feature = "testutils"))]

soroban-sdk/src/env.rs

@@ -445,7 +445,8 @@ use crate::{
     auth,
     testutils::{
         budget::Budget, Address as _, AuthSnapshot, AuthorizedInvocation, ContractFunctionSet,
-        EventsSnapshot, Generators, Ledger as _, MockAuth, MockAuthContract, Snapshot,
+        ContractStruct, EventsSnapshot, Generators, Ledger as _, MockAuth, MockAuthContract,
+        Snapshot,
     },
     Bytes, BytesN,
 };
@@ -596,10 +597,9 @@ impl Env {
     ///     let contract_id = env.register_contract(None, HelloContract);
     /// }
     /// ```
-    pub fn register_contract<'a, T: ContractFunctionSet + 'static>(
+    pub fn register_contract<'a, T: ContractStruct + ContractFunctionSet + 'static>(
         &self,
         contract_id: impl Into<Option<&'a Address>>,
-        contract: T,
     ) -> Address {
         struct InternalContractFunctionSet<T: ContractFunctionSet>(pub(crate) T);
         impl<T: ContractFunctionSet> internal::ContractFunctionSet for InternalContractFunctionSet<T> {
@@ -631,6 +631,7 @@ impl Env {
         } else {
             Address::generate(self)
         };
+        let contract = T::new(self.clone());
         self.env_impl
             .register_test_contract(
                 contract_id.to_object(),
@@ -859,7 +860,7 @@ impl Env {
     /// ```
     pub fn mock_auths(&self, auths: &[MockAuth]) {
         for a in auths {
-            self.register_contract(a.address, MockAuthContract);
+            self.register_contract::<MockAuthContract>(a.address);
         }
         let auths = auths
             .iter()

soroban-sdk/src/testutils.rs

@@ -197,6 +197,11 @@ impl Generators {
     }
 }
 
+#[doc(hidden)]
+pub trait ContractStruct {
+    fn new(env: Env) -> Self;
+}
+
 #[doc(hidden)]
 pub trait ContractFunctionSet {
     fn call(&self, func: &str, env: Env, args: &[Val]) -> Option<Val>;