diff --git a/soroban-sdk-macros/src/lib.rs b/soroban-sdk-macros/src/lib.rs index 32b6a22ee..5bc1daa34 100644 --- a/soroban-sdk-macros/src/lib.rs +++ b/soroban-sdk-macros/src/lib.rs @@ -29,13 +29,10 @@ use proc_macro2::{Literal, Span, TokenStream as TokenStream2}; use quote::{format_ident, quote}; use sha2::{Digest, Sha256}; use std::fs; -use syn::punctuated::Punctuated; -use syn::token::{Brace, Colon}; use syn::{ parse_macro_input, parse_str, spanned::Spanned, Data, DeriveInput, Error, Fields, ItemImpl, ItemStruct, LitStr, Path, Type, Visibility, }; -use syn::{Field, FieldMutability, FieldsNamed}; use syn_ext::HasFnsItem; use soroban_spec_rust::{generate_from_wasm, GenerateFromFileError}; @@ -128,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] @@ -143,6 +141,8 @@ pub fn contract(metadata: TokenStream, input: TokenStream) -> TokenStream { Err(e) => return e.write_errors().into(), }; + let input_orig: TokenStream2 = input.clone().into(); + let item = parse_macro_input!(input as ItemStruct); match &item.fields { Fields::Named(_) | Fields::Unnamed(_) => { @@ -158,35 +158,39 @@ pub fn contract(metadata: TokenStream, input: TokenStream) -> TokenStream { let crate_path = &args.crate_path; - let mut item_modified = item.clone(); - let mut fields = Punctuated::new(); - fields.push(Field { - attrs: Vec::new(), - vis: Visibility::Inherited, - mutability: FieldMutability::None, - ident: Some(format_ident!("env")), - colon_token: Some(Colon::default()), - ty: Type::Verbatim(quote! { #crate_path::Env }), - }); - item_modified.fields = Fields::Named(FieldsNamed { - brace_token: Brace::default(), - named: fields, - }); - 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 client = derive_client_type(&args.crate_path, &ty_str, &client_ident); - quote! { - #item_modified - impl #ty { - pub fn env(&self) -> &#crate_path::Env { - &self.env + 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! { + #code #client diff --git a/soroban-sdk/src/env.rs b/soroban-sdk/src/env.rs index 6312f11b8..511e8a900 100644 --- a/soroban-sdk/src/env.rs +++ b/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>, - contract: T, ) -> Address { struct InternalContractFunctionSet(pub(crate) T); impl internal::ContractFunctionSet for InternalContractFunctionSet { @@ -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::(a.address); } let auths = auths .iter() diff --git a/soroban-sdk/src/testutils.rs b/soroban-sdk/src/testutils.rs index 64541f16f..0b065ef7d 100644 --- a/soroban-sdk/src/testutils.rs +++ b/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;