add contract state nesting

Signed-off-by: Valentyn Faychuk <valy@faychuk.com>

Author: valentynfaychuk

contract_samples/rust/examples/counter_macros.rs

@@ -7,17 +7,24 @@ use alloc::string::ToString;
 use amadeus_sdk::*;
 
 #[contract_state]
-struct SimpleCounter {
-    count: i128,
+struct Metadata {
     owner: Vec<u8>,
+    created_at: i128,
+}
+
+#[contract_state]
+struct Counter {
+    count: i128,
+    #[nested]
+    metadata: Metadata,
 }
 
 #[contract]
-impl SimpleCounter {
+impl Counter {
     pub fn init(&mut self) {
-        *self.owner = account_current();
+        *self.metadata.owner = account_current();
+        *self.metadata.created_at = entry_slot() as i128;
         *self.count = 0;
-        log("Counter initialized");
     }
 
     pub fn increment(&mut self, amount: Vec<u8>) {
@@ -30,7 +37,7 @@ impl SimpleCounter {
 
     pub fn reset(&mut self) {
         let caller = account_caller();
-        amadeus_sdk::assert!(*self.owner == caller, "unauthorized");
+        amadeus_sdk::assert!(*self.metadata.owner == caller, "unauthorized");
         *self.count = 0;
     }
 }

contract_samples/rust/sdk-macros/src/lib.rs

@@ -13,31 +13,64 @@ pub fn contract_state(_attr: TokenStream, item: TokenStream) -> TokenStream {
         return TokenStream::from(quote! { #input });
     };
 
+    let is_nested = |f: &syn::Field| -> bool {
+        f.attrs.iter().any(|attr| attr.path().is_ident("nested"))
+    };
+
     let transformed_fields = fields.named.iter().map(|f| {
         let field_name = &f.ident;
         let field_vis = &f.vis;
-        let field_attrs = &f.attrs;
         let field_ty = &f.ty;
-        quote! {
-            #(#field_attrs)*
-            #field_vis #field_name: LazyCell<#field_ty>
+        let filtered_attrs: Vec<_> = f.attrs.iter()
+            .filter(|attr| !attr.path().is_ident("nested"))
+            .collect();
+
+        if is_nested(f) {
+            quote! {
+                #(#filtered_attrs)*
+                #field_vis #field_name: #field_ty
+            }
+        } else {
+            quote! {
+                #(#filtered_attrs)*
+                #field_vis #field_name: LazyCell<#field_ty>
+            }
         }
     });
 
     let init_calls = fields.named.iter().map(|f| {
         let field = f.ident.as_ref().unwrap();
         let key = field.to_string();
-        quote! { self.#field = LazyCell::new(#key.as_bytes().to_vec()); }
+
+        if is_nested(f) {
+            quote! {
+                let mut key = prefix.clone();
+                key.extend_from_slice(#key.as_bytes());
+                key.push(b':');
+                self.#field.__init_lazy_fields(key);
+            }
+        } else {
+            quote! {
+                let mut key = prefix.clone();
+                key.extend_from_slice(#key.as_bytes());
+                self.#field = LazyCell::new(key);
+            }
+        }
     });
 
     let flush_calls = fields.named.iter().map(|f| {
         let field = f.ident.as_ref().unwrap();
-        quote! { self.#field.flush(); }
+
+        if is_nested(f) {
+            quote! { self.#field.__flush_lazy_fields(); }
+        } else {
+            quote! { self.#field.flush(); }
+        }
     });
 
     let default_fields = fields.named.iter().map(|f| {
         let field_name = &f.ident;
-        quote! { #field_name: LazyCell::default() }
+        quote! { #field_name: Default::default() }
     });
 
     TokenStream::from(quote! {
@@ -52,9 +85,14 @@ pub fn contract_state(_attr: TokenStream, item: TokenStream) -> TokenStream {
             }
         }
 
-        impl #name {
-            pub fn __init_lazy_fields(&mut self) { #(#init_calls)* }
-            pub fn __flush_lazy_fields(&self) { #(#flush_calls)* }
+        impl ContractState for #name {
+            fn __init_lazy_fields(&mut self, prefix: alloc::vec::Vec<u8>) {
+                #(#init_calls)*
+            }
+
+            fn __flush_lazy_fields(&self) {
+                #(#flush_calls)*
+            }
         }
     })
 }
@@ -123,7 +161,7 @@ fn handle_impl_block(impl_block: ItemImpl) -> TokenStream {
             quote! {
                 #(#deserializations)*
                 let mut state = #self_ty::default();
-                state.__init_lazy_fields();
+                state.__init_lazy_fields(alloc::vec::Vec::new());
                 let result = state.#call;
                 state.__flush_lazy_fields();
                 ret(result);
@@ -132,7 +170,7 @@ fn handle_impl_block(impl_block: ItemImpl) -> TokenStream {
             quote! {
                 #(#deserializations)*
                 let mut state = #self_ty::default();
-                state.__init_lazy_fields();
+                state.__init_lazy_fields(alloc::vec::Vec::new());
                 state.#call;
                 state.__flush_lazy_fields();
             }

contract_samples/rust/sdk/src/lib.rs

@@ -10,6 +10,11 @@ pub use storage::*;
 pub use encoding::*;
 pub use amadeus_sdk_macros::{contract, contract_state};
 
+pub trait ContractState {
+    fn __init_lazy_fields(&mut self, prefix: Vec<u8>);
+    fn __flush_lazy_fields(&self);
+}
+
 use core::panic::PanicInfo;
 
 use alloc::{borrow::Cow, vec::Vec, string::String, string::ToString};
@@ -141,6 +146,7 @@ where
     }
 }
 
+
 impl<T> Default for LazyCell<T> {
     fn default() -> Self {
         Self::new(Vec::new())
@@ -156,14 +162,6 @@ impl<T> LazyCell<T> {
         }
     }
 
-    pub fn flush(&self) where T: Payload + Clone {
-        if self.dirty.get() {
-            if let Some(val) = self.value.borrow().as_ref() {
-                kv_put(&self.key, val.clone());
-            }
-        }
-    }
-
     pub fn get(&self) -> T where T: FromKvBytes + Default + Clone {
         if self.value.borrow().is_none() {
             let loaded = kv_get::<T>(&self.key).unwrap_or_default();
@@ -188,3 +186,47 @@ impl<T> LazyCell<T> {
         self.set(current + amount);
     }
 }
+
+impl<T: Payload + Clone> LazyCell<T> {
+    pub fn flush(&self) {
+        if self.dirty.get() {
+            if let Some(val) = self.value.borrow().as_ref() {
+                kv_put(&self.key, val.clone());
+            }
+        }
+    }
+}
+
+impl<T: ContractState + Default> LazyCell<T> {
+
+    pub fn with_mut<F, R>(&mut self, f: F) -> R
+    where
+        F: FnOnce(&mut T) -> R
+    {
+        if self.value.borrow().is_none() {
+            let mut loaded = T::default();
+            loaded.__init_lazy_fields(self.key.clone());
+            *self.value.borrow_mut() = Some(loaded);
+        }
+        self.dirty.set(true);
+        unsafe {
+            let ptr = self.value.as_ptr();
+            f((*ptr).as_mut().unwrap())
+        }
+    }
+
+    pub fn with<F, R>(&self, f: F) -> R
+    where
+        F: FnOnce(&T) -> R
+    {
+        if self.value.borrow().is_none() {
+            let mut loaded = T::default();
+            loaded.__init_lazy_fields(self.key.clone());
+            *self.value.borrow_mut() = Some(loaded);
+        }
+        unsafe {
+            let ptr = self.value.as_ptr();
+            f((*ptr).as_ref().unwrap())
+        }
+    }
+}