make LazyCell implicit

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

Author: valentynfaychuk

contract_samples/rust/examples/counter_macros.rs

@@ -6,34 +6,31 @@ use alloc::vec::Vec;
 use alloc::string::ToString;
 use amadeus_sdk::*;
 
-#[derive(Contract, Default)]
+#[contract_state]
 struct SimpleCounter {
-    count: LazyCell<i128>,
-    owner: LazyCell<Vec<u8>>,
+    count: i128,
+    owner: Vec<u8>,
 }
 
 #[contract]
 impl SimpleCounter {
     pub fn init(&mut self) {
-        self.owner.set(account_current());
-        self.count.set(0);
+        *self.owner = account_current();
+        *self.count = 0;
         log("Counter initialized");
     }
 
     pub fn increment(&mut self, amount: Vec<u8>) {
-        let amount_val = i128::from_bytes(amount);
-        let current = self.count.get();
-        self.count.set(current + amount_val);
+        *self.count += i128::from_bytes(amount);
     }
 
     pub fn get(&self) -> Vec<u8> {
-        self.count.get().to_string().into_bytes()
+        (*self.count).to_string().into_bytes()
     }
 
     pub fn reset(&mut self) {
         let caller = account_caller();
-        let owner = self.owner.get();
-        amadeus_sdk::assert!(caller == owner, "unauthorized");
-        self.count.set(0);
+        amadeus_sdk::assert!(*self.owner == caller, "unauthorized");
+        *self.count = 0;
     }
 }

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

@@ -1,26 +1,57 @@
 use proc_macro::TokenStream;
 use quote::quote;
-use syn::{parse_macro_input, DeriveInput, Data, Fields, ItemImpl, ItemFn, ImplItem, FnArg, ReturnType, Type};
+use syn::{parse_macro_input, ItemImpl, ItemFn, ImplItem, FnArg, ReturnType, Type, ItemStruct, Fields};
 
-#[proc_macro_derive(Contract)]
-pub fn derive_contract(input: TokenStream) -> TokenStream {
-    let input = parse_macro_input!(input as DeriveInput);
+#[proc_macro_attribute]
+pub fn contract_state(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(item as ItemStruct);
     let name = &input.ident;
-    let Data::Struct(data) = &input.data else { return TokenStream::new() };
-    let Fields::Named(fields) = &data.fields else { return TokenStream::new() };
+    let vis = &input.vis;
+    let attrs = &input.attrs;
+
+    let Fields::Named(ref fields) = input.fields else {
+        return TokenStream::from(quote! { #input });
+    };
+
+    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 init_calls = fields.named.iter().map(|f| {
         let field = f.ident.as_ref().unwrap();
         let key = field.to_string();
-        quote! { self.#field = LazyCell::new(b!("__state__::", #key)); }
+        quote! { self.#field = LazyCell::new(#key.as_bytes().to_vec()); }
     });
 
     let flush_calls = fields.named.iter().map(|f| {
         let field = f.ident.as_ref().unwrap();
         quote! { self.#field.flush(); }
     });
 
+    let default_fields = fields.named.iter().map(|f| {
+        let field_name = &f.ident;
+        quote! { #field_name: LazyCell::default() }
+    });
+
     TokenStream::from(quote! {
+        #(#attrs)*
+        #vis struct #name {
+            #(#transformed_fields),*
+        }
+
+        impl Default for #name {
+            fn default() -> Self {
+                Self { #(#default_fields),* }
+            }
+        }
+
         impl #name {
             pub fn __init_lazy_fields(&mut self) { #(#init_calls)* }
             pub fn __flush_lazy_fields(&self) { #(#flush_calls)* }

contract_samples/rust/sdk/src/lib.rs

@@ -8,7 +8,7 @@ pub mod encoding;
 pub use context::*;
 pub use storage::*;
 pub use encoding::*;
-pub use amadeus_sdk_macros::{contract, Contract};
+pub use amadeus_sdk_macros::{contract, contract_state};
 
 use core::panic::PanicInfo;
 
@@ -108,6 +108,39 @@ pub struct LazyCell<T> {
     dirty: core::cell::Cell<bool>,
 }
 
+impl<T> core::ops::Deref for LazyCell<T>
+where
+    T: FromKvBytes + Default + Clone
+{
+    type Target = T;
+
+    fn deref(&self) -> &T {
+        if self.value.borrow().is_none() {
+            let loaded = kv_get::<T>(&self.key).unwrap_or_default();
+            *self.value.borrow_mut() = Some(loaded);
+        }
+        unsafe {
+            (*self.value.as_ptr()).as_ref().unwrap()
+        }
+    }
+}
+
+impl<T> core::ops::DerefMut for LazyCell<T>
+where
+    T: FromKvBytes + Default + Clone
+{
+    fn deref_mut(&mut self) -> &mut T {
+        if self.value.borrow().is_none() {
+            let loaded = kv_get::<T>(&self.key).unwrap_or_default();
+            *self.value.borrow_mut() = Some(loaded);
+        }
+        self.dirty.set(true);
+        unsafe {
+            (*self.value.as_ptr()).as_mut().unwrap()
+        }
+    }
+}
+
 impl<T> Default for LazyCell<T> {
     fn default() -> Self {
         Self::new(Vec::new())
@@ -143,4 +176,15 @@ impl<T> LazyCell<T> {
         *self.value.borrow_mut() = Some(val);
         self.dirty.set(true);
     }
+
+    pub fn update<F>(&self, f: F) where T: FromKvBytes + Default + Clone, F: FnOnce(T) -> T {
+        let current = self.get();
+        let new_value = f(current);
+        self.set(new_value);
+    }
+
+    pub fn add(&self, amount: T) where T: FromKvBytes + Default + Clone + core::ops::Add<Output = T> {
+        let current = self.get();
+        self.set(current + amount);
+    }
 }