wip self

stellar · Jan 19, 2024 · 49c11c3 · 49c11c3
1 parent 049ed1d
commit 49c11c3
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 27 deletions.
diff --git 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
@@ -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()

diff --git 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<Val>;