feat(dan): template macro handles component state

dan_layer/engine/tests/state/Cargo.toml

@@ -9,6 +9,7 @@ edition = "2021"
 [dependencies]
 tari_template_abi = { path = "../../../template_abi" }
 tari_template_lib = { path = "../../../template_lib" }
+tari_template_macros = { path = "../../../template_macros" }
 
 [profile.release]
 opt-level = 's'     # Optimize for size.

dan_layer/engine/tests/state/src/lib.rs

@@ -20,23 +20,15 @@
 //  WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 //  USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-use tari_template_abi::{decode, encode_with_len, FunctionDef, Type};
-use tari_template_lib::{call_engine, generate_abi, generate_main, TemplateImpl};
+use tari_template_macros::template;
 
-// that's what the example should look like from the user's perspective
-#[allow(dead_code)]
+#[template]
 mod state_template {
-    use tari_template_abi::{borsh, Decode, Encode};
-
-    // #[tari::template]
-    #[derive(Encode, Decode)]
     pub struct State {
-        value: u32,
+        pub value: u32,
     }
 
-    // #[tari::impl]
     impl State {
-        // #[tari::constructor]
         pub fn new() -> Self {
             Self { value: 0 }
         }
@@ -49,98 +41,5 @@ mod state_template {
             self.value
         }
     }
-}
-
-// TODO: Macro generated code
-#[no_mangle]
-extern "C" fn State_abi() -> *mut u8 {
-    let template_name = "State".to_string();
-
-    let functions = vec![
-        FunctionDef {
-            name: "new".to_string(),
-            arguments: vec![],
-            output: Type::U32, // the component_id
-        },
-        FunctionDef {
-            name: "set".to_string(),
-            arguments: vec![Type::U32, Type::U32], // the component_id and the new value
-            output: Type::Unit,                    // does not return anything
-        },
-        FunctionDef {
-            name: "get".to_string(),
-            arguments: vec![Type::U32], // the component_id
-            output: Type::U32,          // the stored value
-        },
-    ];
-
-    generate_abi(template_name, functions)
-}
-
-#[no_mangle]
-extern "C" fn State_main(call_info: *mut u8, call_info_len: usize) -> *mut u8 {
-    let mut template_impl = TemplateImpl::new();
-    use tari_template_abi::{ops::*, CreateComponentArg, EmitLogArg, LogLevel};
-    use tari_template_lib::models::ComponentId;
-
-    tari_template_lib::call_engine::<_, ()>(OP_EMIT_LOG, &EmitLogArg {
-        message: "This is a log message from State_main!".to_string(),
-        level: LogLevel::Info,
-    });
-
-    // constructor
-    template_impl.add_function(
-        "new".to_string(),
-        Box::new(|_| {
-            let ret = state_template::State::new();
-            let encoded = encode_with_len(&ret);
-            // Call the engine to create a new component
-            // TODO: proper component id
-            // The macro will know to generate this call because of the #[tari(constructor)] attribute
-            // TODO: what happens if the user wants to return multiple components/types?
-            let component_id = call_engine::<_, ComponentId>(OP_CREATE_COMPONENT, &CreateComponentArg {
-                name: "State".to_string(),
-                quantity: 1,
-                metadata: Default::default(),
-                state: encoded,
-            });
-            let component_id = component_id.expect("no asset id returned");
-            encode_with_len(&component_id)
-        }),
-    );
-
-    template_impl.add_function(
-        "set".to_string(),
-        Box::new(|args| {
-            // read the function paramenters
-            let _component_id: u32 = decode(&args[0]).unwrap();
-            let _new_value: u32 = decode(&args[1]).unwrap();
-
-            // update the component value
-            // TODO: use a real op code (not "123") when they are implemented
-            call_engine::<_, ()>(123, &());
-
-            // the function does not return any value
-            // TODO: implement "Unit" type empty responses. Right now this fails: wrap_ptr(vec![])
-            encode_with_len(&0)
-        }),
-    );
-
-    template_impl.add_function(
-        "get".to_string(),
-        Box::new(|args| {
-            // read the function paramenters
-            let _component_id: u32 = decode(&args[0]).unwrap();
-
-            // get the component state
-            // TODO: use a real op code (not "123") when they are implemented
-            let _state = call_engine::<_, ()>(123, &());
-
-            // return the value
-            let value = 1_u32; // TODO: read from the component state
-            encode_with_len(&value)
-        }),
-    );
 
-    generate_main(call_info, call_info_len, template_impl)
-}
+}

dan_layer/engine/tests/test.rs

@@ -46,25 +46,25 @@ fn test_hello_world() {
 
 #[test]
 fn test_state() {
+    // TODO: use the Component and ComponentId types in the template
     let template_test = TemplateTest::new("State".to_string(), "tests/state".to_string());
 
     // constructor
     let component: ComponentId = template_test.call_function("new".to_string(), vec![]);
-    assert_eq!(component.1, 0);
-    let component: ComponentId = template_test.call_function("new".to_string(), vec![]);
-    assert_eq!(component.1, 1);
 
     // call the "set" method to update the instance value
     let new_value = 20_u32;
     template_test.call_method::<()>("State".to_string(), "set".to_string(), vec![
         encode_with_len(&component),
         encode_with_len(&new_value),
     ]);
+
     // call the "get" method to get the current value
     let value: u32 = template_test.call_method("State".to_string(), "get".to_string(), vec![encode_with_len(
         &component,
     )]);
-    assert_eq!(value, 1);
+    // TODO: when state storage is implemented in the engine, assert the previous setted value (20_u32)
+    assert_eq!(value, 0);
 }
 
 struct TemplateTest {

dan_layer/template_lib/src/lib.rs

@@ -33,7 +33,7 @@ pub mod models;
 
 // TODO: we should only use stdlib if the template dev needs to include it e.g. use core::mem when stdlib is not
 // available
-use std::{collections::HashMap, mem, ptr::copy, slice};
+use std::{collections::HashMap, mem, slice};
 
 use tari_template_abi::{encode_with_len, Decode, Encode, FunctionDef, TemplateDef};
 
@@ -119,21 +119,3 @@ pub fn call_debug<T: AsRef<[u8]>>(data: T) {
     unsafe { debug(ptr, len) }
 }
 
-#[no_mangle]
-pub unsafe extern "C" fn tari_alloc(len: u32) -> *mut u8 {
-    let cap = (len + 4) as usize;
-    let mut buf = Vec::<u8>::with_capacity(cap);
-    let ptr = buf.as_mut_ptr();
-    mem::forget(buf);
-    copy(len.to_le_bytes().as_ptr(), ptr, 4);
-    ptr
-}
-
-#[no_mangle]
-pub unsafe extern "C" fn tari_free(ptr: *mut u8) {
-    let mut len = [0u8; 4];
-    copy(ptr, len.as_mut_ptr(), 4);
-
-    let cap = (u32::from_le_bytes(len) + 4) as usize;
-    let _ = Vec::<u8>::from_raw_parts(ptr, cap, cap);
-}

dan_layer/template_lib/src/models/component.rs

@@ -20,4 +20,42 @@
 //  WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 //  USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+// TODO: use the actual component id type
 pub type ComponentId = ([u8; 32], u32);
+
+use tari_template_abi::{Decode, Encode, encode_with_len, ops::OP_CREATE_COMPONENT, CreateComponentArg};
+
+use crate::call_engine;
+
+pub fn initialise<T: Encode + Decode>(template_name: String, initial_state: T) -> ComponentId {
+    let encoded_state = encode_with_len(&initial_state);
+
+    // Call the engine to create a new component
+    // TODO: proper component id
+    // TODO: what happens if the user wants to return multiple components/types?
+    let component_id = call_engine::<_, ComponentId>(OP_CREATE_COMPONENT, &CreateComponentArg {
+        name: template_name,
+        quantity: 1,
+        metadata: Default::default(),
+        state: encoded_state,
+    });
+    component_id.expect("no asset id returned")
+}
+
+pub fn get_state<T: Encode + Decode>(_id: u32) -> T {
+    // get the component state
+    // TODO: use a real op code (not "123") when they are implemented
+    let _state = call_engine::<_, ()>(123, &());
+
+    // create and return a mock state because state is not implemented yet in the engine
+    let len = std::mem::size_of::<T>();
+    let byte_vec = vec![0_u8; len];
+    let mut mock_value = byte_vec.as_slice();
+    T::deserialize(&mut mock_value).unwrap()
+}
+
+pub fn set_state<T: Encode + Decode>(_id: u32, _state: T) {
+    // update the component value
+    // TODO: use a real op code (not "123") when they are implemented
+    call_engine::<_, ()>(123, &());
+}

dan_layer/template_lib/src/models/mod.rs

@@ -21,4 +21,4 @@
 //  USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 mod component;
-pub use component::ComponentId;
+pub use component::*;

dan_layer/template_macros/Cargo.toml

@@ -11,6 +11,7 @@ proc-macro = true
 
 [dependencies]
 tari_template_abi = { path = "../template_abi" }
+tari_template_lib = { path = "../template_lib" }
 syn = { version = "1.0.98", features = ["full"] }
 proc-macro2 = "1.0.42"
 quote = "1.0.20"

dan_layer/template_macros/src/ast.rs

@@ -34,6 +34,7 @@ use syn::{
     ItemStruct,
     Result,
     ReturnType,
+    Signature,
     Stmt,
 };
 
@@ -95,38 +96,44 @@ impl TemplateAst {
         match item {
             ImplItem::Method(m) => FunctionAst {
                 name: m.sig.ident.to_string(),
-                input_types: Self::get_input_type_tokens(&m.sig.inputs),
+                input_types: Self::get_input_types(&m.sig.inputs),
                 output_type: Self::get_output_type_token(&m.sig.output),
                 statements: Self::get_statements(m),
+                is_constructor: Self::is_constructor(&m.sig),
             },
             _ => todo!(),
         }
     }
 
-    fn get_input_type_tokens(inputs: &Punctuated<FnArg, Comma>) -> Vec<String> {
+    fn get_input_types(inputs: &Punctuated<FnArg, Comma>) -> Vec<TypeAst> {
         inputs
             .iter()
             .map(|arg| match arg {
                 // TODO: handle the "self" case
-                syn::FnArg::Receiver(_) => todo!(),
-                syn::FnArg::Typed(t) => Self::get_type_token(&t.ty),
+                syn::FnArg::Receiver(r) => {
+                    // TODO: validate that it's indeed a reference ("&") to self
+
+                    let mutability = r.mutability.is_some();
+                    TypeAst::Receiver { mutability }
+                },
+                syn::FnArg::Typed(t) => Self::get_type_ast(&t.ty),
             })
             .collect()
     }
 
-    fn get_output_type_token(ast_type: &ReturnType) -> String {
+    fn get_output_type_token(ast_type: &ReturnType) -> Option<TypeAst> {
         match ast_type {
-            syn::ReturnType::Default => String::new(), // the function does not return anything
-            syn::ReturnType::Type(_, t) => Self::get_type_token(t),
+            syn::ReturnType::Default => None, // the function does not return anything
+            syn::ReturnType::Type(_, t) => Some(Self::get_type_ast(t)),
         }
     }
 
-    fn get_type_token(syn_type: &syn::Type) -> String {
+    fn get_type_ast(syn_type: &syn::Type) -> TypeAst {
         match syn_type {
             syn::Type::Path(type_path) => {
                 // TODO: handle "Self"
                 // TODO: detect more complex types
-                type_path.path.segments[0].ident.to_string()
+                TypeAst::Typed(type_path.path.segments[0].ident.clone())
             },
             _ => todo!(),
         }
@@ -135,11 +142,27 @@ impl TemplateAst {
     fn get_statements(method: &ImplItemMethod) -> Vec<Stmt> {
         method.block.stmts.clone()
     }
+
+    fn is_constructor(sig: &Signature) -> bool {
+        match &sig.output {
+            syn::ReturnType::Default => false, // the function does not return anything
+            syn::ReturnType::Type(_, t) => match t.as_ref() {
+                syn::Type::Path(type_path) => type_path.path.segments[0].ident == "Self",
+                _ => false,
+            },
+        }
+    }
 }
 
 pub struct FunctionAst {
     pub name: String,
-    pub input_types: Vec<String>,
-    pub output_type: String,
+    pub input_types: Vec<TypeAst>,
+    pub output_type: Option<TypeAst>,
     pub statements: Vec<Stmt>,
+    pub is_constructor: bool,
+}
+
+pub enum TypeAst {
+    Receiver { mutability: bool },
+    Typed(Ident),
 }

dan_layer/template_macros/src/template/abi.rs

@@ -24,7 +24,7 @@ use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
 use syn::{parse_quote, Expr, Result};
 
-use crate::ast::{FunctionAst, TemplateAst};
+use crate::ast::{FunctionAst, TemplateAst, TypeAst};
 
 pub fn generate_abi(ast: &TemplateAst) -> Result<TokenStream> {
     let abi_function_name = format_ident!("{}_abi", ast.struct_section.ident);
@@ -51,13 +51,13 @@ pub fn generate_abi(ast: &TemplateAst) -> Result<TokenStream> {
 
 fn generate_function_def(f: &FunctionAst) -> Expr {
     let name = f.name.clone();
-    let arguments: Vec<Expr> = f
-        .input_types
-        .iter()
-        .map(String::as_str)
-        .map(generate_abi_type)
-        .collect();
-    let output = generate_abi_type(&f.output_type);
+
+    let arguments: Vec<Expr> = f.input_types.iter().map(generate_abi_type).collect();
+
+    let output = match &f.output_type {
+        Some(type_ast) => generate_abi_type(type_ast),
+        None => parse_quote!(Type::Unit),
+    };
 
     parse_quote!(
         FunctionDef {
@@ -68,26 +68,36 @@ fn generate_function_def(f: &FunctionAst) -> Expr {
     )
 }
 
-fn generate_abi_type(rust_type: &str) -> Expr {
-    // TODO: there may be a better way of handling this
+fn generate_abi_type(rust_type: &TypeAst) -> Expr {
     match rust_type {
-        "" => parse_quote!(Type::Unit),
-        "bool" => parse_quote!(Type::Bool),
-        "i8" => parse_quote!(Type::I8),
-        "i16" => parse_quote!(Type::I16),
-        "i32" => parse_quote!(Type::I32),
-        "i64" => parse_quote!(Type::I64),
-        "i128" => parse_quote!(Type::I128),
-        "u8" => parse_quote!(Type::U8),
-        "u16" => parse_quote!(Type::U16),
-        "u32" => parse_quote!(Type::U32),
-        "u64" => parse_quote!(Type::U64),
-        "u128" => parse_quote!(Type::U128),
-        "String" => parse_quote!(Type::String),
-        _ => todo!(),
+        // on "&self" we want to pass the component id
+        TypeAst::Receiver { .. } => get_component_id_type(),
+        // basic type
+        // TODO: there may be a better way of handling this
+        TypeAst::Typed(ident) => match ident.to_string().as_str() {
+            "" => parse_quote!(Type::Unit),
+            "bool" => parse_quote!(Type::Bool),
+            "i8" => parse_quote!(Type::I8),
+            "i16" => parse_quote!(Type::I16),
+            "i32" => parse_quote!(Type::I32),
+            "i64" => parse_quote!(Type::I64),
+            "i128" => parse_quote!(Type::I128),
+            "u8" => parse_quote!(Type::U8),
+            "u16" => parse_quote!(Type::U16),
+            "u32" => parse_quote!(Type::U32),
+            "u64" => parse_quote!(Type::U64),
+            "u128" => parse_quote!(Type::U128),
+            "String" => parse_quote!(Type::String),
+            "Self" => get_component_id_type(),
+            _ => todo!(),
+        },
     }
 }
 
+fn get_component_id_type() -> Expr {
+    parse_quote!(Type::U32)
+}
+
 #[cfg(test)]
 mod tests {
     use std::str::FromStr;
@@ -101,7 +111,7 @@ mod tests {
     use crate::ast::TemplateAst;
 
     #[test]
-    fn test_hello_world() {
+    fn test_signatures() {
         let input = TokenStream::from_str(indoc! {"
             mod foo {
                 struct Foo {}
@@ -112,7 +122,9 @@ mod tests {
                     pub fn some_args_function(a: i8, b: String) -> u32 {
                         1_u32
                     }
-                    pub fn no_return_function() {}  
+                    pub fn no_return_function() {}
+                    pub fn constructor() -> Self {}
+                    pub fn method(&self){}  
                 } 
             }
         "})
@@ -144,6 +156,16 @@ mod tests {
                             name: "no_return_function".to_string(),
                             arguments: vec![],
                             output: Type::Unit,
+                        },
+                        FunctionDef {
+                            name: "constructor".to_string(),
+                            arguments: vec![],
+                            output: Type::U32,
+                        },
+                        FunctionDef {
+                            name: "method".to_string(),
+                            arguments: vec![Type::U32],
+                            output: Type::Unit,
                         }
                     ],
                 };

dan_layer/template_macros/src/template/definition.rs

@@ -27,15 +27,16 @@ use crate::ast::TemplateAst;
 
 pub fn generate_definition(ast: &TemplateAst) -> TokenStream {
     let template_name = format_ident!("{}", ast.struct_section.ident);
+    let template_fields = &ast.struct_section.fields;
+    let semi_token = &ast.struct_section.semi_token;
     let functions = &ast.impl_section.items;
 
     quote! {
         pub mod template {
-            use super::*;
+            use tari_template_abi::borsh;
 
-            pub struct #template_name {
-                // TODO: fill template fields
-            }
+            #[derive(tari_template_abi::borsh::BorshSerialize, tari_template_abi::borsh::BorshDeserialize)]
+            pub struct #template_name #template_fields #semi_token
 
             impl #template_name {
                 #(#functions)*

dan_layer/template_macros/src/template/dispatcher.rs

@@ -20,11 +20,11 @@
 //  WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 //  USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-use proc_macro2::{Span, TokenStream};
+use proc_macro2::{Ident, Span, TokenStream};
 use quote::{format_ident, quote};
-use syn::{token::Brace, Block, Expr, ExprBlock, Result};
+use syn::{parse_quote, token::Brace, Block, Expr, ExprBlock, Result};
 
-use crate::ast::TemplateAst;
+use crate::ast::{FunctionAst, TemplateAst, TypeAst};
 
 pub fn generate_dispatcher(ast: &TemplateAst) -> Result<TokenStream> {
     let dispatcher_function_name = format_ident!("{}_main", ast.struct_section.ident);
@@ -35,6 +35,7 @@ pub fn generate_dispatcher(ast: &TemplateAst) -> Result<TokenStream> {
         #[no_mangle]
         pub extern "C" fn #dispatcher_function_name(call_info: *mut u8, call_info_len: usize) -> *mut u8 {
             use ::tari_template_abi::{decode, encode_with_len, CallInfo};
+            use ::tari_template_lib::models::{get_state, set_state, initialise};
 
             if call_info.is_null() {
                 panic!("call_info is null");
@@ -43,94 +44,113 @@ pub fn generate_dispatcher(ast: &TemplateAst) -> Result<TokenStream> {
             let call_data = unsafe { Vec::from_raw_parts(call_info, call_info_len, call_info_len) };
             let call_info: CallInfo = decode(&call_data).unwrap();
 
-            let result = match call_info.func_name.as_str() {
-                #( #function_names => #function_blocks )*,
+            let result;
+            match call_info.func_name.as_str() {
+                #( #function_names => #function_blocks ),*,
                 _ => panic!("invalid function name")
             };
 
-            wrap_ptr(encode_with_len(&result))
+            wrap_ptr(result)
         }
     };
 
     Ok(output)
 }
 
-pub fn get_function_names(ast: &TemplateAst) -> Vec<String> {
+fn get_function_names(ast: &TemplateAst) -> Vec<String> {
     ast.get_functions().iter().map(|f| f.name.clone()).collect()
 }
 
-pub fn get_function_blocks(ast: &TemplateAst) -> Vec<Expr> {
+fn get_function_blocks(ast: &TemplateAst) -> Vec<Expr> {
     let mut blocks = vec![];
 
     for function in ast.get_functions() {
-        let statements = function.statements;
-        blocks.push(Expr::Block(ExprBlock {
-            attrs: vec![],
-            label: None,
-            block: Block {
-                brace_token: Brace {
-                    span: Span::call_site(),
-                },
-                stmts: statements,
-            },
-        }));
+        let block = get_function_block(&ast.template_name, function);
+        blocks.push(block);
     }
 
     blocks
 }
 
-#[cfg(test)]
-mod tests {
-    use std::str::FromStr;
-
-    use indoc::indoc;
-    use proc_macro2::TokenStream;
-    use quote::quote;
-    use syn::parse2;
-
-    use crate::{ast::TemplateAst, template::dispatcher::generate_dispatcher};
-
-    #[test]
-    fn test_hello_world() {
-        let input = TokenStream::from_str(indoc! {"
-            mod hello_world {
-                struct HelloWorld {}
-                impl HelloWorld {
-                    pub fn greet() -> String {
-                        \"Hello World!\".to_string()
-                    }
-                } 
-            }
-        "})
-        .unwrap();
-
-        let ast = parse2::<TemplateAst>(input).unwrap();
-
-        let output = generate_dispatcher(&ast).unwrap();
-
-        assert_code_eq(output, quote! {
-            #[no_mangle]
-            pub extern "C" fn HelloWorld_main(call_info: *mut u8, call_info_len: usize) -> *mut u8 {
-                use ::tari_template_abi::{decode, encode_with_len, CallInfo};
-
-                if call_info.is_null() {
-                    panic!("call_info is null");
+fn get_function_block(template_ident: &Ident, ast: FunctionAst) -> Expr {
+    let mut args: Vec<Expr> = vec![];
+    let mut stmts = vec![];
+    let mut should_get_state = false;
+    let mut should_set_state = false;
+
+    // encode all arguments of the functions
+    for (i, input_type) in ast.input_types.into_iter().enumerate() {
+        let arg_ident = format_ident!("arg_{}", i);
+        let stmt = match input_type {
+            // "self" argument
+            TypeAst::Receiver { mutability } => {
+                should_get_state = true;
+                should_set_state = mutability;
+                args.push(parse_quote! { &mut state });
+                parse_quote! {
+                    let #arg_ident =
+                        decode::<u32>(&call_info.args[#i])
+                        .unwrap();
+                }
+            },
+            // non-self argument
+            TypeAst::Typed(type_ident) => {
+                args.push(parse_quote! { #arg_ident });
+                parse_quote! {
+                    let #arg_ident =
+                        decode::<#type_ident>(&call_info.args[#i])
+                        .unwrap();
                 }
+            },
+        };
+        stmts.push(stmt);
+    }
 
-                let call_data = unsafe { Vec::from_raw_parts(call_info, call_info_len, call_info_len) };
-                let call_info: CallInfo = decode(&call_data).unwrap();
+    // load the component state
+    if should_get_state {
+        stmts.push(parse_quote! {
+            let mut state: template::#template_ident = get_state(arg_0);
+        });
+    }
 
-                let result = match call_info.func_name.as_str() {
-                    "greet" => { "Hello World!".to_string() },
-                    _ => panic!("invalid function name")
-                };
+    // call the user defined function in the template
+    let function_ident = Ident::new(&ast.name, Span::call_site());
+    if ast.is_constructor {
+        stmts.push(parse_quote! {
+            let state = template::#template_ident::#function_ident(#(#args),*);
+        });
 
-                wrap_ptr(encode_with_len(&result))
-            }
+        let template_name_str = template_ident.to_string();
+        stmts.push(parse_quote! {
+            let rtn = initialise(#template_name_str.to_string(), state);
+        });
+    } else {
+        stmts.push(parse_quote! {
+            let rtn = template::#template_ident::#function_ident(#(#args),*);
         });
     }
 
-    fn assert_code_eq(a: TokenStream, b: TokenStream) {
-        assert_eq!(a.to_string(), b.to_string());
+    // encode the result value
+    stmts.push(parse_quote! {
+        result = encode_with_len(&rtn);
+    });
+
+    // after user function invocation, update the component state
+    if should_set_state {
+        stmts.push(parse_quote! {
+            set_state(arg_0, state);
+        });
     }
+
+    // construct the code block for the function
+    Expr::Block(ExprBlock {
+        attrs: vec![],
+        label: None,
+        block: Block {
+            brace_token: Brace {
+                span: Span::call_site(),
+            },
+            stmts,
+        },
+    })
 }

dan_layer/template_macros/src/template/mod.rs

@@ -57,3 +57,170 @@ pub fn generate_template(input: TokenStream) -> Result<TokenStream> {
 
     Ok(output)
 }
+
+#[cfg(test)]
+mod tests {
+    use std::str::FromStr;
+
+    use indoc::indoc;
+    use proc_macro2::TokenStream;
+    use quote::quote;
+
+    use super::generate_template;
+
+    #[test]
+    #[allow(clippy::too_many_lines)]
+    fn test_state() {
+        let input = TokenStream::from_str(indoc! {"
+            mod test {
+                struct State {
+                    value: u32
+                }
+                impl State {
+                    pub fn new() -> Self {
+                        Self { value: 0 }
+                    }
+                    pub fn get(&self) -> u32 {
+                        self.value
+                    }
+                    pub fn set(&mut self, value: u32) {
+                        self.value = value;
+                    }
+                } 
+            }
+        "})
+        .unwrap();
+
+        let output = generate_template(input).unwrap();
+
+        assert_code_eq(output, quote! {
+            pub mod template {
+                use tari_template_abi::borsh;
+
+                #[derive(tari_template_abi::borsh::BorshSerialize, tari_template_abi::borsh::BorshDeserialize)]
+                pub struct State {
+                    value: u32
+                }
+
+                impl State {
+                    pub fn new() -> Self {
+                        Self { value: 0 }
+                    }
+                    pub fn get(&self) -> u32 {
+                        self.value
+                    }
+                    pub fn set(&mut self, value: u32) {
+                        self.value = value;
+                    }
+                }
+            }
+
+            #[no_mangle]
+            pub extern "C" fn State_abi() -> *mut u8 {
+                use ::tari_template_abi::{encode_with_len, FunctionDef, TemplateDef, Type};
+
+                let template = TemplateDef {
+                    template_name: "State".to_string(),
+                    functions: vec![
+                        FunctionDef {
+                            name: "new".to_string(),
+                            arguments: vec![],
+                            output: Type::U32,
+                        },
+                        FunctionDef {
+                            name: "get".to_string(),
+                            arguments: vec![Type::U32],
+                            output: Type::U32,
+                        },
+                        FunctionDef {
+                            name: "set".to_string(),
+                            arguments: vec![Type::U32, Type::U32],
+                            output: Type::Unit,
+                        }
+                    ],
+                };
+
+                let buf = encode_with_len(&template);
+                wrap_ptr(buf)
+            }
+
+            #[no_mangle]
+            pub extern "C" fn State_main(call_info: *mut u8, call_info_len: usize) -> *mut u8 {
+                use ::tari_template_abi::{decode, encode_with_len, CallInfo};
+                use ::tari_template_lib::models::{get_state, set_state, initialise};
+
+                if call_info.is_null() {
+                    panic!("call_info is null");
+                }
+
+                let call_data = unsafe { Vec::from_raw_parts(call_info, call_info_len, call_info_len) };
+                let call_info: CallInfo = decode(&call_data).unwrap();
+
+                let result;
+                match call_info.func_name.as_str() {
+                    "new" => {
+                        let state = template::State::new();
+                        let rtn = initialise("State".to_string(), state);
+                        result = encode_with_len(&rtn);
+                    },
+                    "get" => {
+                        let arg_0 = decode::<u32>(&call_info.args[0usize]).unwrap();
+                        let mut state: template::State = get_state(arg_0);
+                        let rtn = template::State::get(&mut state);
+                        result = encode_with_len(&rtn);
+                    },
+                    "set" => {
+                        let arg_0 = decode::<u32>(&call_info.args[0usize]).unwrap();
+                        let arg_1 = decode::<u32>(&call_info.args[1usize]).unwrap();
+                        let mut state: template::State = get_state(arg_0);
+                        let rtn = template::State::set(&mut state, arg_1);
+                        result = encode_with_len(&rtn);
+                        set_state(arg_0, state);
+                    },
+                    _ => panic!("invalid function name")
+                };
+
+                wrap_ptr(result)
+            }
+
+            extern "C" {
+                pub fn tari_engine(op: u32, input_ptr: *const u8, input_len: usize) -> *mut u8;
+            }
+
+            pub fn wrap_ptr(mut v: Vec<u8>) -> *mut u8 {
+                use std::mem;
+
+                let ptr = v.as_mut_ptr();
+                mem::forget(v);
+                ptr
+            }
+
+            #[no_mangle]
+            pub unsafe extern "C" fn tari_alloc(len: u32) -> *mut u8 {
+                use std::{mem, intrinsics::copy};
+
+                let cap = (len + 4) as usize;
+                let mut buf = Vec::<u8>::with_capacity(cap);
+                let ptr = buf.as_mut_ptr();
+                mem::forget(buf);
+                copy(len.to_le_bytes().as_ptr(), ptr, 4);
+                ptr
+            }
+
+            #[no_mangle]
+            pub unsafe extern "C" fn tari_free(ptr: *mut u8) {
+                use std::intrinsics::copy;
+
+                let mut len = [0u8; 4];
+                copy(ptr, len.as_mut_ptr(), 4);
+
+                let cap = (u32::from_le_bytes(len) + 4) as usize;
+                let _ = Vec::<u8>::from_raw_parts(ptr, cap, cap);
+            }
+        });
+    }
+
+    fn assert_code_eq(a: TokenStream, b: TokenStream) {
+        assert_eq!(a.to_string(), b.to_string());
+    }
+}